Extract reusable functions

Author: XuehaiPan

include/pybind11/detail/internals.h

@@ -550,6 +550,83 @@ inline object get_python_state_dict() {
     return state_dict;
 }
 
+// Get or create per-storage capsule in the current interpreter's state dict.
+// Use test-and-set pattern with `PyDict_SetDefault` for thread-safe concurrent access.
+// WARNING: There can be multiple threads creating the storage at the same time, while only one
+//          will succeed in inserting its capsule into the dict. Therefore, the deleter will be
+//          used to clean up the storage of the unused capsules.
+template <typename Payload, bool LeakOnInterpreterShutdown = false>
+Payload *atomic_get_or_create_in_state_dict(const char *key) {
+    error_scope err_scope; // preserve any existing Python error states
+
+    auto state_dict = reinterpret_borrow<dict>(get_python_state_dict());
+    PyObject *capsule_obj = nullptr;
+
+    // First, try to get existing storage (fast path).
+    {
+        capsule_obj = dict_getitemstring(state_dict.ptr(), key);
+        if (capsule_obj == nullptr && PyErr_Occurred()) {
+            throw error_already_set();
+        }
+        // Fallthrough if capsule_obj is nullptr (not found).
+        // Otherwise, we have found the existing storage (most common case) and return it below.
+    }
+
+    if (capsule_obj == nullptr) {
+        // Storage doesn't exist yet, create a new one.
+        // Use unique_ptr for exception safety: if capsule creation throws, the storage is
+        // automatically deleted.
+        auto storage_ptr = std::unique_ptr<Payload>(new Payload{});
+        // Create capsule with destructor to clean up when the interpreter shuts down.
+        auto new_capsule = capsule(
+            storage_ptr.get(),
+            // The destructor will be called when the capsule is GC'ed.
+            //   - If our capsule is inserted into the dict below, it will be kept alive until
+            //     interpreter shutdown, so the destructor will be called at that time.
+            //   - If our capsule is NOT inserted (another thread inserted first), it will be
+            //     destructed when going out of scope here, so the destructor will be called
+            //     immediately, which will also free the storage.
+            /*destructor=*/[](void *ptr) -> void { delete static_cast<Payload *>(ptr); });
+        // At this point, the capsule object is created successfully.
+        // Release the unique_ptr and let the capsule object own the storage to avoid
+        // double-free.
+        (void) storage_ptr.release();
+
+        // Use `PyDict_SetDefault` for atomic test-and-set:
+        //   - If key doesn't exist, inserts our capsule and returns it.
+        //   - If key exists (another thread inserted first), returns the existing value.
+        // This is thread-safe because `PyDict_SetDefault` will hold a lock on the dict.
+        //
+        // NOTE: Here we use `PyDict_SetDefault` instead of `PyDict_SetDefaultRef` because the
+        //       capsule is kept alive until interpreter shutdown, so we do not need to handle
+        //       incref and decref here.
+        capsule_obj = dict_setdefaultstring(state_dict.ptr(), key, new_capsule.ptr());
+        if (capsule_obj == nullptr) {
+            throw error_already_set();
+        }
+        if (LeakOnInterpreterShutdown && capsule_obj == new_capsule.ptr()) {
+            // Our capsule was inserted.
+            // Remove the destructor to leak the storage on interpreter shutdown.
+            if (PyCapsule_SetDestructor(capsule_obj, nullptr) < 0) {
+                throw error_already_set();
+            }
+        }
+    }
+    // - If key already existed, our `new_capsule` is not inserted, it will be destructed
+    //   when going out of scope here, which will also free the storage.
+    // - Otherwise, our `new_capsule` is now in the dict, and it owns the storage and the
+    //   state dict will incref it.
+}
+
+// Get the storage pointer from the capsule.
+void *raw_ptr = PyCapsule_GetPointer(capsule_obj, /*name=*/nullptr);
+if (!raw_ptr) {
+    raise_from(PyExc_SystemError, "pybind11::detail::atomic_get_or_create_in_state_dict() FAILED");
+    throw error_already_set();
+}
+return static_cast<Payload *>(raw_ptr);
+}
+
 template <typename InternalsType>
 class internals_pp_manager {
 public:
@@ -618,26 +695,11 @@ class internals_pp_manager {
         : holder_id_(id), on_fetch_(on_fetch) {}
 
     std::unique_ptr<InternalsType> *get_or_create_pp_in_state_dict() {
-        error_scope err_scope;
-        dict state_dict = get_python_state_dict();
-        auto internals_obj
-            = reinterpret_steal<object>(dict_getitemstringref(state_dict.ptr(), holder_id_));
-        std::unique_ptr<InternalsType> *pp = nullptr;
-        if (internals_obj) {
-            void *raw_ptr = PyCapsule_GetPointer(internals_obj.ptr(), /*name=*/nullptr);
-            if (!raw_ptr) {
-                raise_from(PyExc_SystemError,
-                           "pybind11::detail::internals_pp_manager::get_pp_from_dict() FAILED");
-                throw error_already_set();
-            }
-            pp = reinterpret_cast<std::unique_ptr<InternalsType> *>(raw_ptr);
-            if (on_fetch_ && pp) {
-                on_fetch_(pp->get());
-            }
-        } else {
-            pp = new std::unique_ptr<InternalsType>;
-            // NOLINTNEXTLINE(bugprone-casting-through-void)
-            state_dict[holder_id_] = capsule(reinterpret_cast<void *>(pp));
+        auto *pp
+            = atomic_get_or_create_in_state_dict<std::unique_ptr<InternalsType>,
+                                                 /*LeakOnInterpreterShutdown=*/true>(holder_id_);
+        if (on_fetch_ && pp) {
+            on_fetch_(pp->get());
         }
         return pp;
     }

include/pybind11/gil_safe_call_once.h

@@ -248,67 +248,8 @@ class gil_safe_call_once_and_store {
     }
 
     // Get or create per-storage capsule in the current interpreter's state dict.
-    // Use test-and-set pattern with `PyDict_SetDefault` for thread-safe concurrent access.
     storage_type *get_or_create_storage_in_state_dict() {
-        error_scope err_scope; // preserve any existing Python error states
-
-        auto state_dict = reinterpret_borrow<dict>(detail::get_python_state_dict());
-        const std::string key = get_storage_key();
-        PyObject *capsule_obj = nullptr;
-
-        // First, try to get existing storage (fast path).
-        {
-            capsule_obj = detail::dict_getitemstring(state_dict.ptr(), key.c_str());
-            if (capsule_obj == nullptr && PyErr_Occurred()) {
-                throw error_already_set();
-            }
-            // Fallthrough if capsule_obj is nullptr (not found).
-            // Otherwise, we have found the existing storage (most common case) and return it
-            // below.
-        }
-
-        if (capsule_obj == nullptr) {
-            // Storage doesn't exist yet, create a new one.
-            // Use unique_ptr for exception safety: if capsule creation throws, the storage is
-            // automatically deleted.
-            auto storage_ptr = std::unique_ptr<storage_type>(new storage_type{});
-            // Create capsule with destructor to clean up when the interpreter shuts down.
-            auto new_capsule = capsule(storage_ptr.get(), [](void *ptr) -> void {
-                delete static_cast<storage_type *>(ptr);
-            });
-            // At this point, the capsule object is created successfully.
-            // Release the unique_ptr and let the capsule object own the storage to avoid
-            // double-free.
-            (void) storage_ptr.release();
-
-            // Use `PyDict_SetDefault` for atomic test-and-set:
-            //   - If key doesn't exist, inserts our capsule and returns it.
-            //   - If key exists (another thread inserted first), returns the existing value.
-            // This is thread-safe because `PyDict_SetDefault` will hold a lock on the dict.
-            //
-            // NOTE: Here we use `PyDict_SetDefault` instead of `PyDict_SetDefaultRef` because the
-            // capsule is kept alive until interpreter shutdown, so we do not need to handle incref
-            // and decref here.
-            capsule_obj
-                = detail::dict_setdefaultstring(state_dict.ptr(), key.c_str(), new_capsule.ptr());
-            if (capsule_obj == nullptr) {
-                throw error_already_set();
-            }
-            // - If key already existed, our `new_capsule` is not inserted, it will be destructed
-            //   when going out of scope here, which will also free the storage.
-            // - Otherwise, our `new_capsule` is now in the dict, and it owns the storage and the
-            //   state dict will incref it.
-        }
-
-        // Get the storage pointer from the capsule.
-        void *raw_ptr = PyCapsule_GetPointer(capsule_obj, /*name=*/nullptr);
-        if (!raw_ptr) {
-            raise_from(PyExc_SystemError,
-                       "pybind11::gil_safe_call_once_and_store::"
-                       "get_or_create_storage_in_state_dict() FAILED");
-            throw error_already_set();
-        }
-        return static_cast<storage_type *>(raw_ptr);
+        return detail::atomic_get_or_create_in_state_dict<storage_type>(get_storage_key().c_str());
     }
 
     // No storage needed when subinterpreter support is enabled.