Add an overload getenv_or that supports env var alias. Add new env var KVIKIO_NUM_THREADS. Fix UB.

cpp/include/kvikio/defaults.hpp

@@ -18,11 +18,13 @@
 
 #include <cstddef>
 #include <cstdlib>
+#include <initializer_list>
 #include <sstream>
 #include <stdexcept>
 #include <string>
 
 #include <kvikio/compat_mode.hpp>
+#include <kvikio/error.hpp>
 #include <kvikio/http_status_codes.hpp>
 #include <kvikio/shim/cufile.hpp>
 #include <kvikio/threadpool_wrapper.hpp>
@@ -56,6 +58,52 @@ CompatMode getenv_or(std::string_view env_var_name, CompatMode default_val);
 template <>
 std::vector<int> getenv_or(std::string_view env_var_name, std::vector<int> default_val);
 
+/**
+ * @brief Get the environment variable value from a candidate list
+ *
+ * @tparam T Type of the environment variable value
+ * @param env_var_names Candidate list containing the names of environment variable
+ * @param default_val Default value of the environment variable, if none of the candidates has been
+ * found
+ * @return A tuple of (`env_var_name`, `result`, `has_found`), where:
+ *   - If the environment variable is not set by any of the candidates, `has_found` will be false,
+ * `result` be `default_val`, and `env_var_name` be empty.
+ *   - If the environment variable is set by `env_var_name`, then `has_found` will be true, and
+ * `result` be the set value.
+ *
+ * @throws std::invalid_argument if:
+ *   - `env_var_names` is empty
+ *   - Multiple candidates have been set at the same time
+ *   - An invalid value is given, e.g. value that cannot be converted to type T
+ */
+template <typename T>
+std::tuple<std::string_view, T, bool> getenv_or(
+  std::initializer_list<std::string_view> env_var_names, T default_val)
+{
+  KVIKIO_EXPECT(env_var_names.size() > 0,
+                "`env_var_names` must contain at least one environment variable name.",
+                std::invalid_argument);
+  std::string_view env_name_target;
+  std::string_view env_val_target;
+  bool has_already_been_set{false};
+  for (auto const& env_var_name : env_var_names) {
+    auto const* env_val = std::getenv(env_var_name.data());
+    if (env_val == nullptr) { continue; }
+    KVIKIO_EXPECT(
+      !has_already_been_set,
+      "Environment variable " + std::string{env_var_name} + " has already been set by its alias.",
+      std::invalid_argument);
+    has_already_been_set = true;
+    env_name_target      = env_var_name;
+    env_val_target       = env_val;
+  }
+
+  if (env_name_target.empty()) { return {env_name_target, default_val, false}; }
+
+  auto res = getenv_or<T>(env_name_target, default_val);
+  return {env_name_target, res, true};
+}
+
 /**
  * @brief Singleton class of default values used throughout KvikIO.
  *
@@ -183,6 +231,20 @@ class defaults {
    */
   static void set_thread_pool_nthreads(unsigned int nthreads);
 
+  /**
+   * @brief Alias of `thread_pool_nthreads`
+   *
+   * @return The number of threads
+   */
+  [[nodiscard]] static unsigned int num_threads();
+
+  /**
+   * @brief Alias of `set_thread_pool_nthreads`
+   *
+   * @param nthreads The number of threads to use
+   */
+  static void set_num_threads(unsigned int nthreads);
+
   /**
    * @brief Get the default task size used for parallel IO operations.
    *

cpp/src/defaults.cpp

@@ -28,6 +28,7 @@
 #include <kvikio/error.hpp>
 #include <kvikio/http_status_codes.hpp>
 #include <kvikio/shim/cufile.hpp>
+#include <string_view>
 
 namespace kvikio {
 template <>
@@ -88,9 +89,13 @@ std::vector<int> getenv_or(std::string_view env_var_name, std::vector<int> defau
 unsigned int defaults::get_num_threads_from_env()
 {
   KVIKIO_NVTX_FUNC_RANGE();
-  int const ret = getenv_or("KVIKIO_NTHREADS", 1);
-  KVIKIO_EXPECT(ret > 0, "KVIKIO_NTHREADS has to be a positive integer", std::invalid_argument);
-  return ret;
+
+  auto const [env_var_name, num_threads, _] =
+    getenv_or({"KVIKIO_NTHREADS", "KVIKIO_NUM_THREADS"}, 1);
+  KVIKIO_EXPECT(num_threads > 0,
+                std::string{env_var_name} + " has to be a positive integer",
+                std::invalid_argument);
+  return num_threads;
 }
 
 defaults::defaults()
@@ -183,6 +188,10 @@ void defaults::set_thread_pool_nthreads(unsigned int nthreads)
   thread_pool().reset(nthreads);
 }
 
+unsigned int defaults::num_threads() { return thread_pool_nthreads(); }
+
+void defaults::set_num_threads(unsigned int nthreads) { set_thread_pool_nthreads(nthreads); }
+
 std::size_t defaults::task_size() { return instance()->_task_size; }
 
 void defaults::set_task_size(std::size_t nbytes)

cpp/tests/CMakeLists.txt

@@ -72,7 +72,7 @@ endfunction()
 
 kvikio_add_test(NAME BASIC_IO_TEST SOURCES test_basic_io.cpp)
 
-kvikio_add_test(NAME DEFAULTS_TEST SOURCES test_defaults.cpp)
+kvikio_add_test(NAME DEFAULTS_TEST SOURCES test_defaults.cpp utils/env.cpp)
 
 kvikio_add_test(NAME ERROR_TEST SOURCES test_error.cpp)
 

cpp/tests/test_defaults.cpp

@@ -15,10 +15,18 @@
  */
 
 #include <stdexcept>
+#include <vector>
 
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include <kvikio/defaults.hpp>
 
+#include "kvikio/compat_mode.hpp"
+#include "utils/env.hpp"
+
+using ::testing::HasSubstr;
+using ::testing::ThrowsMessage;
+
 TEST(DefaultsTest, parse_compat_mode_str)
 {
   {
@@ -72,3 +80,120 @@ TEST(DefaultsTest, parse_http_status_codes)
     }
   }
 }
+
+TEST(DefaultsTest, alias_for_getenv_or)
+{
+  // Passed initializer list is empty
+  {
+    EXPECT_THAT([=] { kvikio::getenv_or({}, 123); },
+                ThrowsMessage<std::invalid_argument>(HasSubstr(
+                  "`env_var_names` must contain at least one environment variable name")));
+  }
+
+  // Env var has an empty value
+  {
+    kvikio::test::EnvVarContext env_var_ctx{{{"KVIKIO_TEST_ALIAS", ""}}};
+    EXPECT_THAT(
+      [=] { kvikio::getenv_or({"KVIKIO_TEST_ALIAS"}, 123); },
+      ThrowsMessage<std::invalid_argument>(HasSubstr("unknown config value KVIKIO_TEST_ALIAS=")));
+  }
+
+  // Env var has already been set by its alias
+  {
+    kvikio::test::EnvVarContext env_var_ctx{
+      {{"KVIKIO_TEST_ALIAS_1", "10"}, {"KVIKIO_TEST_ALIAS_2", "20"}}};
+    EXPECT_THAT([=] { kvikio::getenv_or({"KVIKIO_TEST_ALIAS_1", "KVIKIO_TEST_ALIAS_2"}, 123); },
+                ThrowsMessage<std::invalid_argument>(HasSubstr(
+                  "Environment variable KVIKIO_TEST_ALIAS_2 has already been set by its alias")));
+  }
+
+  // Env var has invalid value
+  {
+    kvikio::test::EnvVarContext env_var_ctx{{{"KVIKIO_TEST_ALIAS", "abc"}}};
+    EXPECT_THAT([=] { kvikio::getenv_or({"KVIKIO_TEST_ALIAS"}, 123); },
+                ThrowsMessage<std::invalid_argument>(
+                  HasSubstr("unknown config value KVIKIO_TEST_ALIAS=abc")));
+  }
+
+  // 1st alias has a set value
+  {
+    kvikio::test::EnvVarContext env_var_ctx{{{"KVIKIO_TEST_ALIAS_1", "654.321"}}};
+    auto const [env_var_name, result, has_found] =
+      kvikio::getenv_or({"KVIKIO_TEST_ALIAS_1", "KVIKIO_TEST_ALIAS_2"}, 123.456);
+    EXPECT_EQ(env_var_name, std::string_view{"KVIKIO_TEST_ALIAS_1"});
+    EXPECT_EQ(result, 654.321);
+    EXPECT_TRUE(has_found);
+  }
+
+  // 2nd alias has a set value
+  {
+    kvikio::test::EnvVarContext env_var_ctx{{{"KVIKIO_TEST_ALIAS_2", "654.321"}}};
+    auto const [env_var_name, result, has_found] =
+      kvikio::getenv_or({"KVIKIO_TEST_ALIAS_1", "KVIKIO_TEST_ALIAS_2"}, 123.456);
+    EXPECT_EQ(env_var_name, std::string_view{"KVIKIO_TEST_ALIAS_2"});
+    EXPECT_EQ(result, 654.321);
+    EXPECT_TRUE(has_found);
+  }
+
+  // Neither alias has a set value
+  {
+    auto const [env_var_name, result, has_found] =
+      kvikio::getenv_or({"KVIKIO_TEST_ALIAS_1", "KVIKIO_TEST_ALIAS_2"}, 123.456);
+    EXPECT_TRUE(env_var_name.empty());
+    EXPECT_EQ(result, 123.456);
+    EXPECT_FALSE(has_found);
+  }
+
+  // Special type: bool
+  {
+    kvikio::test::EnvVarContext env_var_ctx{{{"KVIKIO_TEST_ALIAS", "yes"}}};
+    auto const [env_var_name, result, has_found] = kvikio::getenv_or({"KVIKIO_TEST_ALIAS"}, false);
+    EXPECT_EQ(env_var_name, std::string_view{"KVIKIO_TEST_ALIAS"});
+    EXPECT_TRUE(result);
+    EXPECT_TRUE(has_found);
+  }
+  {
+    kvikio::test::EnvVarContext env_var_ctx{{{"KVIKIO_TEST_ALIAS", "OFF"}}};
+    auto const [env_var_name, result, has_found] = kvikio::getenv_or({"KVIKIO_TEST_ALIAS"}, false);
+    EXPECT_EQ(env_var_name, std::string_view{"KVIKIO_TEST_ALIAS"});
+    EXPECT_FALSE(result);
+    EXPECT_TRUE(has_found);
+  }
+
+  // Special type: CompatMode
+  {
+    kvikio::test::EnvVarContext env_var_ctx{{{"KVIKIO_TEST_ALIAS", "yes"}}};
+    auto const [env_var_name, result, has_found] =
+      kvikio::getenv_or({"KVIKIO_TEST_ALIAS"}, kvikio::CompatMode::AUTO);
+    EXPECT_EQ(env_var_name, std::string_view{"KVIKIO_TEST_ALIAS"});
+    EXPECT_EQ(result, kvikio::CompatMode::ON);
+    EXPECT_TRUE(has_found);
+  }
+  {
+    kvikio::test::EnvVarContext env_var_ctx{{{"KVIKIO_TEST_ALIAS", "FALSE"}}};
+    auto const [env_var_name, result, has_found] =
+      kvikio::getenv_or({"KVIKIO_TEST_ALIAS"}, kvikio::CompatMode::AUTO);
+    EXPECT_EQ(env_var_name, std::string_view{"KVIKIO_TEST_ALIAS"});
+    EXPECT_EQ(result, kvikio::CompatMode::OFF);
+    EXPECT_TRUE(has_found);
+  }
+  {
+    kvikio::test::EnvVarContext env_var_ctx{{{"KVIKIO_TEST_ALIAS", "aUtO"}}};
+    auto const [env_var_name, result, has_found] =
+      kvikio::getenv_or({"KVIKIO_TEST_ALIAS"}, kvikio::CompatMode::ON);
+    EXPECT_EQ(env_var_name, std::string_view{"KVIKIO_TEST_ALIAS"});
+    EXPECT_EQ(result, kvikio::CompatMode::AUTO);
+    EXPECT_TRUE(has_found);
+  }
+
+  // Special type: std::vector<int>
+  {
+    kvikio::test::EnvVarContext env_var_ctx{{{"KVIKIO_TEST_ALIAS", "109, 108, 107"}}};
+    auto const [env_var_name, result, has_found] =
+      kvikio::getenv_or({"KVIKIO_TEST_ALIAS"}, std::vector<int>{111, 112, 113});
+    EXPECT_EQ(env_var_name, std::string_view{"KVIKIO_TEST_ALIAS"});
+    std::vector<int> expected{109, 108, 107};
+    EXPECT_EQ(result, expected);
+    EXPECT_TRUE(has_found);
+  }
+}

cpp/tests/utils/env.cpp

@@ -23,21 +23,22 @@
 
 namespace kvikio::test {
 EnvVarContext::EnvVarContext(
-  std::initializer_list<std::pair<std::string, std::string>> env_var_entries)
+  std::initializer_list<std::pair<std::string_view, std::string_view>> env_var_entries)
 {
   for (auto const& [key, current_value] : env_var_entries) {
     EnvVarState env_var_state;
-    if (auto const res = std::getenv(key.c_str()); res != nullptr) {
+    if (auto const res = std::getenv(key.data()); res != nullptr) {
       env_var_state.existed_before = true;
       env_var_state.previous_value = res;
     }
-    SYSCALL_CHECK(setenv(key.c_str(), current_value.c_str(), 1 /* allow overwrite */));
-    if (_env_var_map.find(key) != _env_var_map.end()) {
+    SYSCALL_CHECK(setenv(key.data(), current_value.data(), 1 /* allow overwrite */));
+    std::string key_str{key};
+    if (_env_var_map.find(key_str) != _env_var_map.end()) {
       std::stringstream ss;
       ss << "Environment variable " << key << " has already been set in this context.";
       KVIKIO_FAIL(ss.str());
     }
-    _env_var_map.insert({key, std::move(env_var_state)});
+    _env_var_map.insert({std::move(key_str), std::move(env_var_state)});
   }
 }
 

cpp/tests/utils/env.hpp

@@ -45,7 +45,8 @@ class EnvVarContext {
    * @param env_var_entries User-specified environment variables. Each entry includes the variable
    * name and value.
    */
-  EnvVarContext(std::initializer_list<std::pair<std::string, std::string>> env_var_entries);
+  EnvVarContext(
+    std::initializer_list<std::pair<std::string_view, std::string_view>> env_var_entries);
 
   /**
    * @brief Restore the environment variables to previous values