From 8d529e859e2abba843cfcb1b7f2b6906db9b362d Mon Sep 17 00:00:00 2001
From: Melody Ren <melodyr@nvidia.com>
Date: Wed, 8 Jan 2025 18:10:07 +0000
Subject: [PATCH] add python callback to ensure correct shutdown order

Signed-off-by: Melody Ren <melodyr@nvidia.com>
---
 libs/qec/include/cudaq/qec/plugin_loader.h |  8 ++---
 libs/qec/lib/plugin_loader.cpp             | 39 +++++++++++-----------
 libs/qec/python/bindings/py_decoder.cpp    |  8 +++++
 3 files changed, 31 insertions(+), 24 deletions(-)

diff --git a/libs/qec/include/cudaq/qec/plugin_loader.h b/libs/qec/include/cudaq/qec/plugin_loader.h
index 810bed1..280b1f1 100644
--- a/libs/qec/include/cudaq/qec/plugin_loader.h
+++ b/libs/qec/include/cudaq/qec/plugin_loader.h
@@ -11,6 +11,7 @@
 
 #include <dlfcn.h>
 #include <map>
+#include <memory>
 #include <string>
 
 /// @brief Enum to define different types of plugins
@@ -22,10 +23,9 @@ enum class PluginType {
 
 /// @brief A struct to store plugin handle with its type
 struct PluginHandle {
-  void *handle; // Pointer to the shared library handle. This is the result of
-                // dlopen() function.
-  PluginType type; // Type of the plugin (e.g., decoder, code, etc)
-  bool is_closed;  // Flag indicating if the handle is closed
+  std::shared_ptr<void> handle; // Pointer to the shared library handle. This is
+                                // the result of dlopen() function.
+  PluginType type;              // Type of the plugin (e.g., decoder, code, etc)
 };
 
 /// @brief Function to load plugins from a directory based on type
diff --git a/libs/qec/lib/plugin_loader.cpp b/libs/qec/lib/plugin_loader.cpp
index f450cad..df870cd 100644
--- a/libs/qec/lib/plugin_loader.cpp
+++ b/libs/qec/lib/plugin_loader.cpp
@@ -19,37 +19,36 @@ static std::map<std::string, PluginHandle> &get_plugin_handles() {
 
 // Function to load plugins from a directory based on their type
 void load_plugins(const std::string &plugin_dir, PluginType type) {
-  if (!fs::exists(plugin_dir)) {
-    std::cerr << "WARNING: Plugin directory does not exist: " << plugin_dir
-              << std::endl;
-    return;
-  }
   for (const auto &entry : fs::directory_iterator(plugin_dir)) {
     if (entry.path().extension() == ".so") {
+      void *raw_handle = dlopen(entry.path().c_str(), RTLD_NOW);
+      if (raw_handle) {
+        // Custom deleter ensures dlclose is called
+        auto deleter = [](void *h) {
+          if (h)
+            dlclose(h);
+        };
 
-      void *handle = dlopen(entry.path().c_str(), RTLD_NOW);
-
-      if (!handle) {
+        get_plugin_handles().emplace(
+            entry.path().filename().string(),
+            PluginHandle{std::shared_ptr<void>(raw_handle, deleter), type});
+      } else {
         std::cerr << "ERROR: Failed to load plugin: " << entry.path()
                   << " Error: " << dlerror() << std::endl;
-      } else {
-        get_plugin_handles().emplace(entry.path().filename().string(),
-                                     PluginHandle{handle, type, false});
       }
     }
   }
 }
 
+// Function to clean up the plugin handles
 void cleanup_plugins(PluginType type) {
-  for (auto &[key, plugin] : get_plugin_handles()) {
-    if (plugin.type == type) {
-      if (plugin.handle && !plugin.is_closed) {
-        dlclose(plugin.handle);
-        plugin.is_closed = true;
-      } else {
-        std::cerr << "WARNING: Invalid or null handle for plugin: " << key
-                  << "\n";
-      }
+  auto &handles = get_plugin_handles();
+  auto it = handles.begin();
+  while (it != handles.end()) {
+    if (it->second.type == type) {
+      it = handles.erase(it); // dlclose is handled by the custom deleter
+    } else {
+      ++it;
     }
   }
 }
diff --git a/libs/qec/python/bindings/py_decoder.cpp b/libs/qec/python/bindings/py_decoder.cpp
index 8b6d591..11b6789 100644
--- a/libs/qec/python/bindings/py_decoder.cpp
+++ b/libs/qec/python/bindings/py_decoder.cpp
@@ -14,6 +14,7 @@
 #include "common/Logger.h"
 
 #include "cudaq/qec/decoder.h"
+#include "cudaq/qec/plugin_loader.h"
 
 #include "type_casters.h"
 #include "utils.h"
@@ -70,6 +71,13 @@ std::unordered_map<std::string, std::function<py::object(
     PyDecoderRegistry::registry;
 
 void bindDecoder(py::module &mod) {
+  // Required by all plugin classes
+  auto cleanup_callback = []() {
+    // Change the type to the correct plugin type
+    cleanup_plugins(PluginType::DECODER);
+  };
+  // This ensures the correct shutdown sequence
+  mod.add_object("_cleanup", py::capsule(cleanup_callback));
 
   auto qecmod = py::hasattr(mod, "qecrt")
                     ? mod.attr("qecrt").cast<py::module_>()