Hack to handle OMP_NUM_THREADS

1 files changed, 15 insertions, 3 deletions

diff --git a/threadpool.c b/threadpool.c
index 57d046b..fb6f0ee 100644
--- a/threadpool.c
+++ b/threadpool.c
@@ -23,7 +23,7 @@
 #include "threadpool.h"
 
 struct TPContext {
-    char dummy;
+    unsigned int nb_threads;
 };
 
 void tp_free(TPContext **pctx)
@@ -40,6 +40,7 @@ void tp_free(TPContext **pctx)
 int tp_init(TPContext **pctx, unsigned int nb_threads)
 {
     TPContext *ctx = NULL;
+    char *env_threads;
     int ret;
 
     if (!nb_threads) {
@@ -53,7 +54,18 @@ int tp_init(TPContext **pctx, unsigned int nb_threads)
         goto fail;
     }
 
-    omp_set_num_threads(nb_threads);
+    env_threads = getenv("OMP_NUM_THREADS");
+    if (env_threads) {
+        nb_threads = strtol(env_threads, NULL, 0);
+        if (!nb_threads) {
+            ret = -EINVAL;
+            goto fail;
+        }
+    } else {
+        omp_set_num_threads(nb_threads);
+    }
+
+    ctx->nb_threads = nb_threads;
 
     *pctx = ctx;
     return 0;
@@ -76,5 +88,5 @@ int tp_execute(TPContext *ctx, unsigned int nb_jobs,
 
 unsigned int tp_get_nb_threads(TPContext *ctx)
 {
-    return omp_get_num_threads();
+    return ctx->nb_threads;
 }