From bd302d773c50552531bc7f11f782f8ed876e8fab Mon Sep 17 00:00:00 2001
From: Nghia Truong <7416935+ttnghia@users.noreply.github.com>
Date: Mon, 29 Jul 2024 17:07:33 -0700
Subject: [PATCH] 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 https://github.com/rapidsai/cudf/issues/16426.

---------

Signed-off-by: Nghia Truong <nghiat@nvidia.com>
---
 cpp/include/cudf/utilities/prefetch.hpp |  6 ++++++
 cpp/src/utilities/prefetch.cpp          | 15 +++++++++------
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/cpp/include/cudf/utilities/prefetch.hpp b/cpp/include/cudf/utilities/prefetch.hpp
index 49fca73a2c8..3384181fc37 100644
--- a/cpp/include/cudf/utilities/prefetch.hpp
+++ b/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,6 +48,8 @@ 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.
    */
@@ -54,6 +57,8 @@ class prefetch_config {
   /**
    * @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
 };
 
 /**
diff --git a/cpp/src/utilities/prefetch.cpp b/cpp/src/utilities/prefetch.cpp
index 16f2c3a1202..86d6cc00764 100644
--- a/cpp/src/utilities/prefetch.cpp
+++ b/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,