Add per-interpreter storage for gil_safe_call_once_and_store

include/pybind11/detail/internals.h

@@ -420,8 +420,8 @@ inline PyThreadState *get_thread_state_unchecked() {
 
 /// We use this counter to figure out if there are or have been multiple subinterpreters active at
 /// any point. This must never decrease while any interpreter may be running in any thread!
-inline std::atomic<int> &get_num_interpreters_seen() {
-    static std::atomic<int> counter(0);
+inline std::atomic<int64_t> &get_num_interpreters_seen() {
+    static std::atomic<int64_t> counter(0);
     return counter;
 }
 
@@ -564,7 +564,7 @@ class internals_pp_manager {
     /// acquire the GIL. Will never return nullptr.
     std::unique_ptr<InternalsType> *get_pp() {
 #ifdef PYBIND11_HAS_SUBINTERPRETER_SUPPORT
-        if (get_num_interpreters_seen() > 1) {
+        if (get_num_interpreters_seen() > 1 || last_istate_tls() == nullptr) {
             // Whenever the interpreter changes on the current thread we need to invalidate the
             // internals_pp so that it can be pulled from the interpreter's state dict.  That is
             // slow, so we use the current PyThreadState to check if it is necessary.
@@ -590,11 +590,8 @@ class internals_pp_manager {
     /// Drop all the references we're currently holding.
     void unref() {
 #ifdef PYBIND11_HAS_SUBINTERPRETER_SUPPORT
-        if (get_num_interpreters_seen() > 1) {
-            last_istate_tls() = nullptr;
-            internals_p_tls() = nullptr;
-            return;
-        }
+        last_istate_tls() = nullptr;
+        internals_p_tls() = nullptr;
 #endif
         internals_singleton_pp_ = nullptr;
     }
@@ -606,7 +603,6 @@ class internals_pp_manager {
             // this could be called without an active interpreter, just use what was cached
             if (!tstate || tstate->interp == last_istate_tls()) {
                 auto tpp = internals_p_tls();
-
                 delete tpp;
             }
             unref();
@@ -660,6 +656,8 @@ class internals_pp_manager {
 
     char const *holder_id_ = nullptr;
     on_fetch_function *on_fetch_ = nullptr;
+    // Pointer-to-pointer to the singleton internals for the first seen interpreter (may not be the
+    // main interpreter)
     std::unique_ptr<InternalsType> *internals_singleton_pp_;
 };
 

include/pybind11/gil_safe_call_once.h

@@ -3,17 +3,31 @@
 #pragma once
 
 #include "detail/common.h"
+#include "detail/internals.h"
 #include "gil.h"
 
 #include <cassert>
 #include <mutex>
 
-#ifdef Py_GIL_DISABLED
+#if defined(Py_GIL_DISABLED) || defined(PYBIND11_HAS_SUBINTERPRETER_SUPPORT)
 #    include <atomic>
 #endif
+#ifdef PYBIND11_HAS_SUBINTERPRETER_SUPPORT
+#    include <cstdint>
+#    include <memory>
+#    include <string>
+#endif
 
 PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
 
+PYBIND11_NAMESPACE_BEGIN(detail)
+#if defined(Py_GIL_DISABLED) || defined(PYBIND11_HAS_SUBINTERPRETER_SUPPORT)
+using atomic_bool = std::atomic_bool;
+#else
+using atomic_bool = bool;
+#endif
+PYBIND11_NAMESPACE_END(detail)
+
 // Use the `gil_safe_call_once_and_store` class below instead of the naive
 //
 //   static auto imported_obj = py::module_::import("module_name"); // BAD, DO NOT USE!
@@ -48,12 +62,23 @@ PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
 // functions, which is usually the case.
 //
 // For in-depth background, see docs/advanced/deadlock.md
+#ifndef PYBIND11_HAS_SUBINTERPRETER_SUPPORT
+// Subinterpreter support is disabled.
+// In this case, we can store the result globally, because there is only a single interpreter.
+//
+// The life span of the stored result is the entire process lifetime. It is leaked on process
+// termination to avoid destructor calls after the Python interpreter was finalized.
 template <typename T>
 class gil_safe_call_once_and_store {
 public:
     // PRECONDITION: The GIL must be held when `call_once_and_store_result()` is called.
+    //
+    // NOTE: The second parameter (finalize callback) is intentionally unused when subinterpreter
+    // support is disabled. In that case, storage is process-global and intentionally leaked to
+    // avoid calling destructors after the Python interpreter has been finalized.
     template <typename Callable>
-    gil_safe_call_once_and_store &call_once_and_store_result(Callable &&fn) {
+    gil_safe_call_once_and_store &call_once_and_store_result(Callable &&fn,
+                                                             void (*)(T &) /*unused*/ = nullptr) {
         if (!is_initialized_) { // This read is guarded by the GIL.
             // Multiple threads may enter here, because the GIL is released in the next line and
             // CPython API calls in the `fn()` call below may release and reacquire the GIL.
@@ -74,29 +99,232 @@ class gil_safe_call_once_and_store {
     T &get_stored() {
         assert(is_initialized_);
         PYBIND11_WARNING_PUSH
-#if !defined(__clang__) && defined(__GNUC__) && __GNUC__ < 5
+#    if !defined(__clang__) && defined(__GNUC__) && __GNUC__ < 5
         // Needed for gcc 4.8.5
         PYBIND11_WARNING_DISABLE_GCC("-Wstrict-aliasing")
-#endif
+#    endif
         return *reinterpret_cast<T *>(storage_);
         PYBIND11_WARNING_POP
     }
 
     constexpr gil_safe_call_once_and_store() = default;
+    // The instance is a global static, so its destructor runs when the process
+    // is terminating. Therefore, do nothing here because the Python interpreter
+    // may have been finalized already.
     PYBIND11_DTOR_CONSTEXPR ~gil_safe_call_once_and_store() = default;
 
+    // Disable copy and move operations.
+    gil_safe_call_once_and_store(const gil_safe_call_once_and_store &) = delete;
+    gil_safe_call_once_and_store(gil_safe_call_once_and_store &&) = delete;
+    gil_safe_call_once_and_store &operator=(const gil_safe_call_once_and_store &) = delete;
+    gil_safe_call_once_and_store &operator=(gil_safe_call_once_and_store &&) = delete;
+
 private:
+    // The global static storage (per-process) when subinterpreter support is disabled.
     alignas(T) char storage_[sizeof(T)] = {};
     std::once_flag once_flag_;
-#ifdef Py_GIL_DISABLED
-    std::atomic_bool
-#else
-    bool
-#endif
-        is_initialized_{false};
+
     // The `is_initialized_`-`storage_` pair is very similar to `std::optional`,
     // but the latter does not have the triviality properties of former,
     // therefore `std::optional` is not a viable alternative here.
+    detail::atomic_bool is_initialized_{false};
 };
+#else
+// Subinterpreter support is enabled.
+// In this case, we should store the result per-interpreter instead of globally, because each
+// subinterpreter has its own separate state. The cached result may not shareable across
+// interpreters (e.g., imported modules and their members).
+
+PYBIND11_NAMESPACE_BEGIN(detail)
+
+template <typename T>
+struct call_once_storage {
+    alignas(T) char storage[sizeof(T)] = {};
+    std::once_flag once_flag;
+    void (*finalize)(T &) = nullptr;
+    std::atomic_bool is_initialized{false};
+
+    call_once_storage() = default;
+    ~call_once_storage() {
+        if (is_initialized) {
+            if (finalize != nullptr) {
+                finalize(*reinterpret_cast<T *>(storage));
+            } else {
+                reinterpret_cast<T *>(storage)->~T();
+            }
+        }
+    }
+    call_once_storage(const call_once_storage &) = delete;
+    call_once_storage(call_once_storage &&) = delete;
+    call_once_storage &operator=(const call_once_storage &) = delete;
+    call_once_storage &operator=(call_once_storage &&) = delete;
+};
+
+PYBIND11_NAMESPACE_END(detail)
+
+// Prefix for storage keys in the interpreter state dict.
+#    define PYBIND11_CALL_ONCE_STORAGE_KEY_PREFIX PYBIND11_INTERNALS_ID "_call_once_storage__"
+
+// The life span of the stored result is the entire interpreter lifetime. An additional
+// `finalize_fn` can be provided to clean up the stored result when the interpreter is destroyed.
+template <typename T>
+class gil_safe_call_once_and_store {
+public:
+    // PRECONDITION: The GIL must be held when `call_once_and_store_result()` is called.
+    template <typename Callable>
+    gil_safe_call_once_and_store &call_once_and_store_result(Callable &&fn,
+                                                             void (*finalize_fn)(T &) = nullptr) {
+        if (!is_last_storage_valid()) {
+            // Multiple threads may enter here, because the GIL is released in the next line and
+            // CPython API calls in the `fn()` call below may release and reacquire the GIL.
+            gil_scoped_release gil_rel; // Needed to establish lock ordering.
+            // There can be multiple threads going through here.
+            storage_type *value = nullptr;
+            {
+                gil_scoped_acquire gil_acq; // Restore lock ordering.
+                // This function is thread-safe under free-threading.
+                value = get_or_create_storage_in_state_dict();
+            }
+            assert(value != nullptr);
+            std::call_once(value->once_flag, [&] {
+                // Only one thread will ever enter here.
+                gil_scoped_acquire gil_acq;
+                // fn may release, but will reacquire, the GIL.
+                ::new (value->storage) T(fn());
+                value->finalize = finalize_fn;
+                value->is_initialized = true;
+                last_storage_ptr_ = reinterpret_cast<T *>(value->storage);
+                is_initialized_by_at_least_one_interpreter_ = true;
+            });
+            // All threads will observe `is_initialized_by_at_least_one_interpreter_` as true here.
+        }
+        // Intentionally not returning `T &` to ensure the calling code is self-documenting.
+        return *this;
+    }
+
+    // This must only be called after `call_once_and_store_result()` was called.
+    T &get_stored() {
+        T *result = last_storage_ptr_;
+        if (!is_last_storage_valid()) {
+            gil_scoped_acquire gil_acq;
+            auto *value = get_or_create_storage_in_state_dict();
+            result = last_storage_ptr_ = reinterpret_cast<T *>(value->storage);
+        }
+        assert(result != nullptr);
+        return *result;
+    }
+
+    gil_safe_call_once_and_store() = default;
+    // The instance is a global static, so its destructor runs when the process
+    // is terminating. Therefore, do nothing here because the Python interpreter
+    // may have been finalized already.
+    PYBIND11_DTOR_CONSTEXPR ~gil_safe_call_once_and_store() = default;
+
+    // Disable copy and move operations because the memory address is used as key.
+    gil_safe_call_once_and_store(const gil_safe_call_once_and_store &) = delete;
+    gil_safe_call_once_and_store(gil_safe_call_once_and_store &&) = delete;
+    gil_safe_call_once_and_store &operator=(const gil_safe_call_once_and_store &) = delete;
+    gil_safe_call_once_and_store &operator=(gil_safe_call_once_and_store &&) = delete;
+
+private:
+    using storage_type = detail::call_once_storage<T>;
+
+    // Indicator of fast path for single-interpreter case.
+    bool is_last_storage_valid() const {
+        return is_initialized_by_at_least_one_interpreter_
+               && detail::get_num_interpreters_seen() == 1;
+    }
+
+    // Get the unique key for this storage instance in the interpreter's state dict.
+    // The return type should not be `py::str` because PyObject is interpreter-dependent.
+    std::string get_storage_key() const {
+        // The instance is expected to be global static, so using its address as unique identifier.
+        // The typical usage is like:
+        //
+        //   PYBIND11_CONSTINIT static gil_safe_call_once_and_store<T> storage;
+        //
+        return PYBIND11_CALL_ONCE_STORAGE_KEY_PREFIX
+               + std::to_string(reinterpret_cast<std::uintptr_t>(this));
+    }
+
+    // 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);
+    }
+
+    // No storage needed when subinterpreter support is enabled.
+    // The actual storage is stored in the per-interpreter state dict via
+    // `get_or_create_storage_in_state_dict()`.
+
+    // Fast local cache to avoid repeated lookups when there are no multiple interpreters.
+    // This is only valid if there is a single interpreter. Otherwise, it is not used.
+    // WARNING: We cannot use thread local cache similar to `internals_pp_manager::internals_p_tls`
+    //          because the thread local storage cannot be explicitly invalidated when interpreters
+    //          are destroyed (unlike `internals_pp_manager` which has explicit hooks for that).
+    T *last_storage_ptr_ = nullptr;
+    // This flag is true if the value has been initialized by any interpreter (may not be the
+    // current one).
+    detail::atomic_bool is_initialized_by_at_least_one_interpreter_{false};
+};
+#endif
 
 PYBIND11_NAMESPACE_END(PYBIND11_NAMESPACE)