[Environment Variable][1/N] Use thread-safe env variable API in c10 (#119449)

This PR is the beginning of attempts to wrap thread-unsafe getenv and set_env functions inside a RW mutex.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/119449
Approved by: https://github.com/malfet, https://github.com/albanD, https://github.com/eqy

c10/test/util/DeadlockDetection_test.cpp

@@ -1,9 +1,8 @@
 #include <c10/util/DeadlockDetection.h>
+#include <c10/util/env.h>
 
 #include <gtest/gtest.h>
 
-#include <cstdlib>
-
 using namespace ::testing;
 using namespace c10::impl;
 
@@ -23,7 +22,7 @@ TEST(DeadlockDetection, basic) {
 
 #ifndef _WIN32
 TEST(DeadlockDetection, disable) {
-  setenv("TORCH_DISABLE_DEADLOCK_DETECTION", "1", 1);
+  c10::utils::set_env("TORCH_DISABLE_DEADLOCK_DETECTION", "1");
   DummyPythonGILHooks hooks;
   SetPythonGILHooks(&hooks);
   SetPythonGILHooks(&hooks);

c10/util/DeadlockDetection.cpp

@@ -1,6 +1,5 @@
 #include <c10/util/DeadlockDetection.h>
-
-#include <cstdlib>
+#include <c10/util/env.h>
 
 namespace c10::impl {
 
@@ -8,7 +7,7 @@ namespace {
 PythonGILHooks* python_gil_hooks = nullptr;
 
 bool disable_detection() {
-  return std::getenv("TORCH_DISABLE_DEADLOCK_DETECTION") != nullptr;
+  return c10::utils::has_env("TORCH_DISABLE_DEADLOCK_DETECTION");
 }
 } // namespace
 

c10/util/Logging.cpp

@@ -3,6 +3,7 @@
 #include <c10/util/Flags.h>
 #include <c10/util/Lazy.h>
 #include <c10/util/Logging.h>
+#include <c10/util/env.h>
 #ifdef FBCODE_CAFFE2
 #include <folly/synchronization/SanitizeThread.h>
 #endif
@@ -12,7 +13,6 @@
 #endif
 
 #include <algorithm>
-#include <cstdlib>
 #include <iostream>
 
 // Common code that we use regardless of whether we use glog or not.
@@ -122,8 +122,8 @@ using DDPUsageLoggerType = std::function<void(const DDPLoggingData&)>;
 
 namespace {
 bool IsAPIUsageDebugMode() {
-  const char* val = getenv("PYTORCH_API_USAGE_STDERR");
-  return val && *val; // any non-empty value
+  auto val = c10::utils::get_env("PYTORCH_API_USAGE_STDERR");
+  return val.has_value() && !val.value().empty(); // any non-empty value
 }
 
 void APIUsageDebug(const string& event) {
@@ -504,10 +504,10 @@ namespace c10::detail {
 namespace {
 
 void setLogLevelFlagFromEnv() {
-  const char* level_str = std::getenv("TORCH_CPP_LOG_LEVEL");
+  auto level_env = c10::utils::get_env("TORCH_CPP_LOG_LEVEL");
 
   // Not set, fallback to the default level (i.e. WARNING).
-  std::string level{level_str != nullptr ? level_str : ""};
+  std::string level{level_env.has_value() ? level_env.value() : ""};
   if (level.empty()) {
     return;
   }

c10/util/env.cpp

@@ -0,0 +1,94 @@
+#include <c10/util/Exception.h>
+#include <c10/util/env.h>
+#include <fmt/format.h>
+#include <cstdlib>
+#include <shared_mutex>
+
+namespace c10::utils {
+
+static std::shared_mutex env_mutex;
+
+// Set an environment variable.
+void set_env(const char* name, const char* value, bool overwrite) {
+  std::lock_guard lk(env_mutex);
+#ifdef _MSC_VER
+#pragma warning(push)
+#pragma warning(disable : 4996)
+  if (!overwrite) {
+    // NOLINTNEXTLINE(concurrency-mt-unsafe)
+    if (std::getenv(name) != nullptr) {
+      return;
+    }
+  }
+  auto full_env_variable = fmt::format("{}={}", name, value);
+  // NOLINTNEXTLINE(concurrency-mt-unsafe)
+  auto err = putenv(full_env_variable.c_str());
+  TORCH_INTERNAL_ASSERT(
+      err == 0,
+      "putenv failed for environment \"",
+      name,
+      "\", the error is: ",
+      err);
+#pragma warning(pop)
+#else
+  // NOLINTNEXTLINE(concurrency-mt-unsafe)
+  auto err = setenv(name, value, static_cast<int>(overwrite));
+  TORCH_INTERNAL_ASSERT(
+      err == 0,
+      "setenv failed for environment \"",
+      name,
+      "\", the error is: ",
+      err);
+#endif
+  return;
+}
+
+// Reads an environment variable and returns the content if it is set
+std::optional<std::string> get_env(const char* name) noexcept {
+  std::shared_lock lk(env_mutex);
+#ifdef _MSC_VER
+#pragma warning(push)
+#pragma warning(disable : 4996)
+#endif
+  // NOLINTNEXTLINE(concurrency-mt-unsafe)
+  auto envar = std::getenv(name);
+#ifdef _MSC_VER
+#pragma warning(pop)
+#endif
+  if (envar != nullptr) {
+    return std::string(envar);
+  }
+  return std::nullopt;
+}
+
+// Checks an environment variable is set.
+bool has_env(const char* name) noexcept {
+  return get_env(name).has_value();
+}
+
+// Reads an environment variable and returns
+// - optional<true>,              if set equal to "1"
+// - optional<false>,             if set equal to "0"
+// - nullopt,   otherwise
+//
+// NB:
+// Issues a warning if the value of the environment variable is not 0 or 1.
+std::optional<bool> check_env(const char* name) {
+  auto env_opt = get_env(name);
+  if (env_opt.has_value()) {
+    if (*env_opt == "0") {
+      return false;
+    }
+    if (*env_opt == "1") {
+      return true;
+    }
+    TORCH_WARN(
+        "Ignoring invalid value for boolean flag ",
+        name,
+        ": ",
+        *env_opt,
+        "valid values are 0 or 1.");
+  }
+  return std::nullopt;
+}
+} // namespace c10::utils

c10/util/env.h

@@ -1,11 +1,20 @@
 #pragma once
 
-#include <c10/util/Exception.h>
-#include <cstdlib>
-#include <cstring>
+#include <c10/macros/Export.h>
 #include <optional>
+#include <string>
 
 namespace c10::utils {
+
+// Set an environment variable.
+C10_API void set_env(
+    const char* name,
+    const char* value,
+    bool overwrite = true);
+
+// Checks an environment variable is set.
+C10_API bool has_env(const char* name) noexcept;
+
 // Reads an environment variable and returns
 // - std::optional<true>,              if set equal to "1"
 // - std::optional<false>,             if set equal to "0"
@@ -13,29 +22,10 @@ namespace c10::utils {
 //
 // NB:
 // Issues a warning if the value of the environment variable is not 0 or 1.
-inline std::optional<bool> check_env(const char* name) {
-#ifdef _MSC_VER
-#pragma warning(push)
-#pragma warning(disable : 4996)
-#endif
-  auto envar = std::getenv(name);
-#ifdef _MSC_VER
-#pragma warning(pop)
-#endif
-  if (envar) {
-    if (strcmp(envar, "0") == 0) {
-      return false;
-    }
-    if (strcmp(envar, "1") == 0) {
-      return true;
-    }
-    TORCH_WARN(
-        "Ignoring invalid value for boolean flag ",
-        name,
-        ": ",
-        envar,
-        "valid values are 0 or 1.");
-  }
-  return std::nullopt;
-}
+C10_API std::optional<bool> check_env(const char* name);
+
+// Reads the value of an environment variable if it is set.
+// However, check_env should be used if the value is assumed to be a flag.
+C10_API std::optional<std::string> get_env(const char* name) noexcept;
+
 } // namespace c10::utils

c10/util/tempfile.cpp

@@ -1,4 +1,5 @@
 #include <c10/util/Exception.h>
+#include <c10/util/env.h>
 #include <c10/util/tempfile.h>
 #include <fmt/format.h>
 
@@ -22,10 +23,11 @@ static std::string make_filename(std::string_view name_prefix) {
   // We see if any of these environment variables is set and use their value, or
   // else default the temporary directory to `/tmp`.
 
-  const char* tmp_directory = "/tmp";
+  std::string tmp_directory = "/tmp";
   for (const char* variable : {"TMPDIR", "TMP", "TEMP", "TEMPDIR"}) {
-    if (const char* path = getenv(variable)) {
-      tmp_directory = path;
+    auto path_opt = c10::utils::get_env(variable);
+    if (path_opt.has_value()) {
+      tmp_directory = path_opt.value();
       break;
     }
   }