From 70fa7b37ca7eb85bc36b34880d09801e562b635f Mon Sep 17 00:00:00 2001
From: senhuang42 <senhuang96@fb.com>
Date: Tue, 25 Aug 2020 18:49:52 -0400
Subject: [PATCH] Added env var function for threads

---
 programs/zstdcli.c | 27 ++++++++++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/programs/zstdcli.c b/programs/zstdcli.c
index 1aea63dca..27cc752db 100644
--- a/programs/zstdcli.c
+++ b/programs/zstdcli.c
@@ -16,6 +16,10 @@
 #  define ZSTDCLI_CLEVEL_DEFAULT 3
 #endif
 
+#ifndef ZSTD_THREADS_DEFAULT
+#  define ZSTD_THREADS_DEFAULT 3
+#endif
+
 #ifndef ZSTDCLI_CLEVEL_MAX
 #  define ZSTDCLI_CLEVEL_MAX 19   /* without using --ultra */
 #endif
@@ -586,6 +590,7 @@ static void printVersion(void)
 
 /* Environment variables for parameter setting */
 #define ENV_CLEVEL "ZSTD_CLEVEL"
+#define ENV_THREADS "ZSTD_THREADS"
 
 /* pick up environment variable */
 static int init_cLevel(void) {
@@ -615,6 +620,26 @@ static int init_cLevel(void) {
     return ZSTDCLI_CLEVEL_DEFAULT;
 }
 
+/* pick up environment variable */
+static unsigned init_numThreads(void) {
+    const char* const env = getenv(ENV_THREADS);
+    if (env != NULL) {
+        const char* ptr = env;
+        if ((*ptr>='0') && (*ptr<='9')) {
+            unsigned numThreads;
+            if (readU32FromCharChecked(&ptr, &numThreads)) {
+                DISPLAYLEVEL(2, "Ignore environment variable setting %s=%s: numeric value too large \n", ENV_THREADS, env);
+                return ZSTD_THREADS_DEFAULT;
+            } else if (*ptr == 0) {
+                return numThreads;
+            }
+        }
+        DISPLAYLEVEL(2, "Ignore environment variable setting %s=%s: not a valid integer value \n", ENV_THREADS, env);
+    }
+
+    return ZSTD_THREADS_DEFAULT;
+}
+
 #define NEXT_FIELD(ptr) {         \
     if (*argument == '=') {       \
         ptr = ++argument;         \
@@ -721,7 +746,7 @@ int main(int const argCount, const char* argv[])
     if ((filenames==NULL) || (file_of_names==NULL)) { DISPLAY("zstd: allocation error \n"); exit(1); }
     programName = lastNameFromPath(programName);
 #ifdef ZSTD_MULTITHREAD
-    nbWorkers = 1;
+    nbWorkers = init_numThreads();
 #endif
 
     /* preset behaviors */