Fix the destruction of interruptible token registry

cpp/include/raft/core/interruptible.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -178,11 +178,13 @@ class interruptible {
   auto operator=(interruptible&&) -> interruptible& = delete;
 
  private:
-  /** Global registry of thread-local cancellation stores. */
-  static inline std::unordered_map<std::thread::id, std::weak_ptr<interruptible>> registry_;
   /** Protect the access to the registry. */
   static inline std::mutex mutex_;
 
+  using registry_t = std::unordered_map<std::thread::id, std::weak_ptr<interruptible>>;
+  /** Global registry of thread-local cancellation stores. */
+  static inline std::shared_ptr<registry_t> registry_{new registry_t{}};
+
   /**
    * Create a new interruptible token or get an existing from the global registry_.
    *
@@ -203,21 +205,25 @@ class interruptible {
   {
     std::lock_guard<std::mutex> guard_get(mutex_);
     // the following constructs an empty shared_ptr if the key does not exist.
-    auto& weak_store  = registry_[thread_id];
+    auto& weak_store  = (*registry_)[thread_id];
     auto thread_store = weak_store.lock();
     if (!thread_store || (Claim && thread_store->claimed_)) {
+      std::weak_ptr<registry_t> weak_registry = registry_;
       // Create a new thread_store in two cases:
       //  1. It does not exist in the map yet
       //  2. The previous store in the map has not yet been deleted
-      thread_store.reset(new interruptible(), [thread_id](auto ts) {
-        std::lock_guard<std::mutex> guard_erase(mutex_);
-        auto found = registry_.find(thread_id);
-        if (found != registry_.end()) {
-          auto stored = found->second.lock();
-          // thread_store is not moveable, thus retains its original location.
-          // Not equal pointers below imply the new store has been already placed
-          // in the registry_ by the same std::thread::id
-          if (!stored || stored.get() == ts) { registry_.erase(found); }
+      thread_store.reset(new interruptible(), [weak_registry, thread_id](auto ts) {
+        auto registry = weak_registry.lock();
+        if (registry) {
+          std::lock_guard<std::mutex> guard_erase(mutex_);
+          auto found = registry->find(thread_id);
+          if (found != registry->end()) {
+            auto stored = found->second.lock();
+            // thread_store is not moveable, thus retains its original location.
+            // Not equal pointers below imply the new store has been already placed
+            // in the registry by the same std::thread::id
+            if (!stored || stored.get() == ts) { registry->erase(found); }
+          }
         }
         delete ts;
       });
@@ -268,4 +274,4 @@ class interruptible {
 
 }  // namespace raft
 
-#endif
+#endif

cpp/include/raft/core/resource/cuda_stream.hpp

@@ -91,10 +91,7 @@ inline void set_cuda_stream(resources const& res, rmm::cuda_stream_view stream_v
  */
 inline void sync_stream(const resources& res, rmm::cuda_stream_view stream)
 {
-  // TODO: Fix interruptible segfault:
-  // https://github.com/rapidsai/raft/issues/1225
-  // interruptible::synchronize(stream);
-  RAFT_CUDA_TRY(cudaStreamSynchronize(stream));
+  interruptible::synchronize(stream);
 }
 
 /**
@@ -106,4 +103,4 @@ inline void sync_stream(const resources& res) { sync_stream(res, get_cuda_stream
  * @}
  */
 
-}  // namespace raft::resource
+}  // namespace raft::resource