Refactor the token deleter from the lambda to a custom type

rapidsai · rapids-bot · Feb 15, 2023 · Feb 2, 2023 · Feb 3, 2023 · Feb 3, 2023
commit 4873727c66c560e693dbbc74720d562f0c70141d
@@ -178,9 +178,44 @@ class interruptible {
   auto operator=(interruptible&&) -> interruptible& = delete;
 
  private:
+  /** Global registry of thread-local cancellation stores. */
   using registry_t =
     std::tuple<std::mutex, std::unordered_map<std::thread::id, std::weak_ptr<interruptible>>>;
-  /** Global registry of thread-local cancellation stores. */
+
+  /**
+   * The registry "garbage collector": a custom deleter for the interruptible tokens that removes
+   * the token from the registry, if the registry still exists.
+   */
+  struct registry_gc_t {
+    std::weak_ptr<registry_t> weak_registry;
+    std::thread::id thread_id;
+
+    inline void operator()(interruptible* thread_store) const noexcept
+    {
+      // the deleter kicks in at thread/program exit; in some cases, the registry_ (static variable)
+      // may have been destructed by this point of time.
+      // Hence, we use a weak pointer to check if the registry still exists.
+      auto registry = weak_registry.lock();
+      if (registry) {
+        std::lock_guard<std::mutex> guard_erase(std::get<0>(*registry));
+        auto& map  = std::get<1>(*registry);
+        auto found = map.find(thread_id);
+        if (found != map.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() == thread_store) { map.erase(found); }
+        }
+      }
+      delete thread_store;
+    }
+  };
+
+  /**
+   * The registry itself is stored in the static memory, in a shared pointer.
+   * This is to safely access it from the destructors of the thread-local tokens.
+   */
   static inline std::shared_ptr<registry_t> registry_{new registry_t{}};
 
   /**
@@ -213,26 +248,10 @@ class interruptible {
     auto& weak_store  = std::get<1>(*shared_registry)[thread_id];
     auto thread_store = weak_store.lock();
     if (!thread_store || (Claim && thread_store->claimed_)) {
-      std::weak_ptr<registry_t> weak_registry = shared_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(), [weak_registry, thread_id](auto ts) {
-        auto registry = weak_registry.lock();
-        if (registry) {
-          std::lock_guard<std::mutex> guard_erase(std::get<0>(*registry));
-          auto& map  = std::get<1>(*registry);
-          auto found = map.find(thread_id);
-          if (found != map.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) { map.erase(found); }
-          }
-        }
-        delete ts;
-      });
+      thread_store.reset(new interruptible(), registry_gc_t{shared_registry, thread_id});
       std::weak_ptr<interruptible>(thread_store).swap(weak_store);
     }
     // The thread_store is "claimed" by the thread