Add per-interpreter storage for gil_safe_call_once_and_store

Author: XuehaiPan

include/pybind11/detail/internals.h

@@ -39,7 +39,7 @@
 /// further ABI-incompatible changes may be made before the ABI is officially
 /// changed to the new version.
 #ifndef PYBIND11_INTERNALS_VERSION
-#    define PYBIND11_INTERNALS_VERSION 11
+#    define PYBIND11_INTERNALS_VERSION 12
 #endif
 
 #if PYBIND11_INTERNALS_VERSION < 11
@@ -234,6 +234,34 @@ inline uint64_t round_up_to_next_pow2(uint64_t x) {
 
 class loader_life_support;
 
+struct call_once_storage_base {
+    call_once_storage_base() = default;
+    virtual ~call_once_storage_base() = default;
+    call_once_storage_base(const call_once_storage_base &) = delete;
+    call_once_storage_base(call_once_storage_base &&) = delete;
+    call_once_storage_base &operator=(const call_once_storage_base &) = delete;
+    call_once_storage_base &operator=(call_once_storage_base &&) = delete;
+};
+
+template <typename T>
+struct call_once_storage : call_once_storage_base {
+    void (*finalize)(T &) = nullptr;
+    alignas(T) char storage[sizeof(T)] = {0};
+
+    call_once_storage() = default;
+    ~call_once_storage() override {
+        if (finalize != nullptr) {
+            finalize(*reinterpret_cast<T *>(storage));
+        }
+        memset(storage, 0, sizeof(T));
+        finalize = nullptr;
+    };
+    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;
+};
+
 /// Internal data structure used to track registered instances and types.
 /// Whenever binary incompatible changes are made to this structure,
 /// `PYBIND11_INTERNALS_VERSION` must be incremented.
@@ -283,6 +311,8 @@ struct internals {
 
     type_map<PyObject *> native_enum_type_map;
 
+    std::unordered_map<const void *, call_once_storage_base *> call_once_storage_map;
+
     internals()
         : static_property_type(make_static_property_type()),
           default_metaclass(make_default_metaclass()) {
@@ -308,7 +338,12 @@ struct internals {
     internals(internals &&other) = delete;
     internals &operator=(const internals &other) = delete;
     internals &operator=(internals &&other) = delete;
-    ~internals() = default;
+    ~internals() {
+        for (auto &[_, storage_ptr] : call_once_storage_map) {
+            delete storage_ptr;
+        }
+        call_once_storage_map.clear();
+    }
 };
 
 // the internals struct (above) is shared between all the modules. local_internals are only

include/pybind11/gil_safe_call_once.h

@@ -3,6 +3,7 @@
 #pragma once
 
 #include "detail/common.h"
+#include "detail/internals.h"
 #include "gil.h"
 
 #include <cassert>
@@ -52,13 +53,15 @@ 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
 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_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.
@@ -80,10 +83,10 @@ 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
     }
@@ -96,6 +99,7 @@ class gil_safe_call_once_and_store {
     }
 
 private:
+    // Global static storage (per process) when subinterpreter support is disabled.
     alignas(T) char storage_[sizeof(T)] = {};
     std::once_flag once_flag_;
     void (*finalize_fn_)(T &) = nullptr;
@@ -105,5 +109,88 @@ class gil_safe_call_once_and_store {
     // therefore `std::optional` is not a viable alternative here.
     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 object may not shareable
+// across interpreters (e.g., imported modules and their members).
+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_initialized_by_atleast_one_interpreter_
+            || detail::get_num_interpreters_seen() > 1) {
+            detail::with_internals([&](detail::internals &internals) {
+                const void *key = reinterpret_cast<const void *>(this);
+                auto &storage_map = internals.call_once_storage_map;
+                auto it = storage_map.find(key);
+                if (it == storage_map.end()) {
+                    gil_scoped_release gil_rel; // Needed to establish lock ordering.
+                    {
+                        // Only one thread will ever enter here.
+                        gil_scoped_acquire gil_acq;
+                        auto s = new detail::call_once_storage<T>{};
+                        ::new (s->storage) T(fn()); // fn may release, but will reacquire, the GIL.
+                        s->finalize = finalize_fn;
+                        last_storage_ = reinterpret_cast<T *>(s->storage);
+                        storage_map.emplace(key, s);
+                    };
+                }
+                is_initialized_by_atleast_one_interpreter_ = true;
+            });
+            // All threads will observe `is_initialized_by_atleast_one_interp_` 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_;
+        if (!is_initialized_by_atleast_one_interpreter_
+            || detail::get_num_interpreters_seen() > 1) {
+            detail::with_internals([&](detail::internals &internals) {
+                const void *key = reinterpret_cast<const void *>(this);
+                auto &storage_map = internals.call_once_storage_map;
+                auto it = storage_map.find(key);
+                assert(it != storage_map.end());
+                auto *s = static_cast<detail::call_once_storage<T> *>(it->second);
+                result = last_storage_ = reinterpret_cast<T *>(s->storage);
+            });
+        }
+        assert(result != nullptr);
+        return *result;
+    }
+
+    constexpr gil_safe_call_once_and_store() = default;
+    PYBIND11_DTOR_CONSTEXPR ~gil_safe_call_once_and_store() {
+        if (is_initialized_by_atleast_one_interpreter_) {
+            detail::with_internals([&](detail::internals &internals) {
+                const void *key = reinterpret_cast<const void *>(this);
+                auto &storage_map = internals.call_once_storage_map;
+                auto it = storage_map.find(key);
+                if (it != storage_map.end()) {
+                    delete it->second;
+                    storage_map.erase(it);
+                }
+            });
+        }
+    }
+
+private:
+    // No storage needed when subinterpreter support is enabled.
+    // The actual storage is stored in the per-interpreter state dict in
+    // `internals.call_once_storage_map`.
+
+    // 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.
+    T *last_storage_ = nullptr;
+    // This flag is true if the value has been initialized by any interpreter (may not be the
+    // current one).
+    atomic_bool is_initialized_by_atleast_one_interpreter_{false};
+};
+#endif
 PYBIND11_NAMESPACE_END(PYBIND11_NAMESPACE)