Support thread-safe for prefetch_config::get and `prefetch_config::…

…set` (#16425)

This adds muti-thread support for `prefetch_config` getter and setter
functions. This avoid the issue that the config map is corrupted in
multi-thread environments.

Closes #16426.

---------

Signed-off-by: Nghia Truong <[email protected]>

cpp/include/cudf/utilities/prefetch.hpp

@@ -21,6 +21,7 @@
 #include <rmm/device_uvector.hpp>
 
 #include <map>
+#include <shared_mutex>
 #include <string>
 #include <string_view>
 
@@ -47,13 +48,17 @@ class prefetch_config {
   /**
    * @brief Get the value of a configuration key.
    *
+   * If the key does not exist, a `false` value will be returned.
+   *
    * @param key The configuration key.
    * @return The value of the configuration key.
    */
   bool get(std::string_view key);
   /**
    * @brief Set the value of a configuration key.
    *
+   * This is a thread-safe operation.
+   *
    * @param key The configuration key.
    * @param value The value to set.
    */
@@ -68,6 +73,7 @@ class prefetch_config {
  private:
   prefetch_config() = default;                //< Private constructor to enforce singleton pattern
   std::map<std::string, bool> config_values;  //< Map of configuration keys to values
+  std::shared_mutex config_mtx;               //< Mutex for thread-safe config access
 };
 
 /**

cpp/src/utilities/prefetch.cpp

@@ -34,13 +34,16 @@ prefetch_config& prefetch_config::instance()
 
 bool prefetch_config::get(std::string_view key)
 {
-  // Default to not prefetching
-  if (config_values.find(key.data()) == config_values.end()) {
-    return (config_values[key.data()] = false);
-  }
-  return config_values[key.data()];
+  std::shared_lock<std::shared_mutex> lock(config_mtx);
+  auto const it = config_values.find(key.data());
+  return it == config_values.end() ? false : it->second;  // default to not prefetching
+}
+
+void prefetch_config::set(std::string_view key, bool value)
+{
+  std::lock_guard<std::shared_mutex> lock(config_mtx);
+  config_values[key.data()] = value;
 }
-void prefetch_config::set(std::string_view key, bool value) { config_values[key.data()] = value; }
 
 cudaError_t prefetch_noexcept(std::string_view key,
                               void const* ptr,