add locking for remaining internal fields

.github/workflows/ci.yml

@@ -73,6 +73,7 @@ jobs:
     - uses: actions/checkout@v4
       with:
         submodules: true
+
     - uses: deadsnakes/[email protected]
       with:
         python-version: 3.13-dev

src/nb_enum.cpp

@@ -17,12 +17,21 @@ using enum_map = tsl::robin_map<int64_t, int64_t, int64_hash>;
 
 PyObject *enum_create(enum_init_data *ed) noexcept {
     // Update hash table that maps from std::type_info to Python type
-    auto [it, success] = internals->type_c2p_slow.try_emplace(ed->type, nullptr);
-    if (!success) {
-        PyErr_WarnFormat(PyExc_RuntimeWarning, 1, "nanobind: type '%s' was already registered!\n", ed->name);
-        PyObject *tp = (PyObject *) it->second->type_py;
-        Py_INCREF(tp);
-        return tp;
+    nb_internals *internals_ = internals;
+    bool success;
+    nb_type_map_slow::iterator it;
+
+    {
+        lock_internals guard(internals_);
+        std::tie(it, success) = internals->type_c2p_slow.try_emplace(ed->type, nullptr);
+        if (!success) {
+            PyErr_WarnFormat(PyExc_RuntimeWarning, 1,
+                             "nanobind: type '%s' was already registered!\n",
+                             ed->name);
+            PyObject *tp = (PyObject *) it->second->type_py;
+            Py_INCREF(tp);
+            return tp;
+        }
     }
 
     handle scope(ed->scope);
@@ -71,8 +80,14 @@ PyObject *enum_create(enum_init_data *ed) noexcept {
 
     it.value() = t;
 
-    internals->type_c2p_fast[ed->type] = t;
-    internals->type_c2p_slow[ed->type] = t;
+    {
+        lock_internals guard(internals_);
+        internals_->type_c2p_slow[ed->type] = t;
+
+        #if !defined(NB_FREE_THREADED)
+            internals_->type_c2p_fast[ed->type] = t;
+        #endif
+    }
 
     result.attr("__nb_enum__") = capsule(t, [](void *p) noexcept {
         type_init_data *t = (type_init_data *) p;

src/nb_func.cpp

@@ -82,11 +82,15 @@ void nb_func_dealloc(PyObject *self) {
         func_data *f = nb_func_data(self);
 
         // Delete from registered function list
-        size_t n_deleted = internals->funcs.erase(self);
-        check(n_deleted == 1,
-              "nanobind::detail::nb_func_dealloc(\"%s\"): function not found!",
-              ((f->flags & (uint32_t) func_flags::has_name) ? f->name
+        nb_internals *internals_ = internals;
+        {
+            lock_internals guard(internals_);
+            size_t n_deleted = internals_->funcs.erase(self);
+            check(n_deleted == 1,
+                  "nanobind::detail::nb_func_dealloc(\"%s\"): function not found!",
+                  ((f->flags & (uint32_t) func_flags::has_name) ? f->name
                                                             : "<anonymous>"));
+        }
         for (size_t i = 0; i < size; ++i) {
             if (f->flags & (uint32_t) func_flags::has_free)
                 f->free_capture(f->capture);
@@ -216,14 +220,15 @@ PyObject *nb_func_new(const void *in_) noexcept {
     }
 
     // Check for previous overloads
+    nb_internals *internals_ = internals;
     if (has_scope && has_name) {
         name = PyUnicode_InternFromString(name_cstr);
         check(name, "nb::detail::nb_func_new(\"%s\"): invalid name.", name_cstr);
 
         func_prev = PyObject_GetAttr(f->scope, name);
         if (func_prev) {
-            if (Py_TYPE(func_prev) == internals->nb_func ||
-                Py_TYPE(func_prev) == internals->nb_method) {
+            if (Py_TYPE(func_prev) == internals_->nb_func ||
+                Py_TYPE(func_prev) == internals_->nb_method) {
                 func_data *fp = nb_func_data(func_prev);
 
                 check((fp->flags & (uint32_t) func_flags::is_method) ==
@@ -274,7 +279,7 @@ PyObject *nb_func_new(const void *in_) noexcept {
     // Create a new function and destroy the old one
     Py_ssize_t to_copy = func_prev ? Py_SIZE(func_prev) : 0;
     nb_func *func = (nb_func *) PyType_GenericAlloc(
-        is_method ? internals->nb_method : internals->nb_func, to_copy + 1);
+        is_method ? internals_->nb_method : internals_->nb_func, to_copy + 1);
     check(func, "nb::detail::nb_func_new(\"%s\"): alloc. failed (1).",
           name_cstr);
 
@@ -294,7 +299,8 @@ PyObject *nb_func_new(const void *in_) noexcept {
 
         ((PyVarObject *) func_prev)->ob_size = 0;
 
-        size_t n_deleted = internals->funcs.erase(func_prev);
+        lock_internals guard(internals_);
+        size_t n_deleted = internals_->funcs.erase(func_prev);
         check(n_deleted == 1,
               "nanobind::detail::nb_func_new(): internal update failed (1)!");
 
@@ -307,9 +313,14 @@ PyObject *nb_func_new(const void *in_) noexcept {
                                           : nb_func_vectorcall_simple;
 
     // Register the function
-    auto [it, success] = internals->funcs.try_emplace(func, nullptr);
-    check(success,
-          "nanobind::detail::nb_func_new(): internal update failed (2)!");
+    nb_ptr_map::iterator it;
+    {
+        lock_internals guard(internals_);
+        bool success;
+        std::tie(it, success) = internals_->funcs.try_emplace(func, nullptr);
+        check(success,
+              "nanobind::detail::nb_func_new(): internal update failed (2)!");
+    }
 
     func_data *fc = nb_func_data(func) + to_copy;
     memcpy(fc, f, sizeof(func_data_prelim<0>));
@@ -1092,14 +1103,21 @@ static uint32_t nb_func_render_signature(const func_data *f,
                       "nb::detail::nb_func_render_signature(): missing type!");
 
                 if (!(is_method && arg_index == 0)) {
-                    auto it = internals->type_c2p_slow.find(*descr_type);
-
-                    if (it != internals->type_c2p_slow.end()) {
-                        handle th((PyObject *) it->second->type_py);
-                        buf.put_dstr((borrow<str>(th.attr("__module__"))).c_str());
-                        buf.put('.');
-                        buf.put_dstr((borrow<str>(th.attr("__qualname__"))).c_str());
-                    } else {
+                    bool found = false;
+                    {
+                        nb_internals *internals_ = internals;
+                        lock_internals guard(internals_);
+                        auto it = internals_->type_c2p_slow.find(*descr_type);
+
+                        if (it != internals_->type_c2p_slow.end()) {
+                            handle th((PyObject *) it->second->type_py);
+                            buf.put_dstr((borrow<str>(th.attr("__module__"))).c_str());
+                            buf.put('.');
+                            buf.put_dstr((borrow<str>(th.attr("__qualname__"))).c_str());
+                            found = true;
+                        }
+                    }
+                    if (!found) {
                         if (nb_signature_mode)
                             buf.put('"');
                         char *name = type_name(*descr_type);

src/nb_internals.cpp

@@ -18,7 +18,7 @@
 
 /// Tracks the ABI of nanobind
 #ifndef NB_INTERNALS_VERSION
-#  define NB_INTERNALS_VERSION 14
+#  define NB_INTERNALS_VERSION 15
 #endif
 
 /// On MSVC, debug and release builds are not ABI-compatible!
@@ -234,7 +234,8 @@ static bool *is_alive_ptr = &is_alive_value;
 bool is_alive() noexcept { return *is_alive_ptr; }
 
 static void internals_cleanup() {
-    if (!internals)
+    nb_internals *p = internals;
+    if (!p)
         return;
 
     *is_alive_ptr = false;
@@ -243,12 +244,12 @@ static void internals_cleanup() {
     /* The memory leak checker is unsupported on PyPy, see
        see https://foss.heptapod.net/pypy/pypy/-/issues/3855 */
 
-    bool print_leak_warnings = internals->print_leak_warnings;
+    bool print_leak_warnings = p->print_leak_warnings;
     size_t inst_leaks = 0, keep_alive_leaks = 0;
 
     // Shard locking no longer needed, Py_AtExit is single-threaded
-    for (size_t i = 0; i < internals->shard_count; ++i) {
-        nb_shard &s = internals->shards[i];
+    for (size_t i = 0; i < p->shard_count; ++i) {
+        nb_shard &s = p->shards[i];
         inst_leaks += s.inst_c2p.size();
         keep_alive_leaks += s.keep_alive.size();
     }
@@ -260,12 +261,12 @@ static void internals_cleanup() {
 
 #if !defined(Py_LIMITED_API)
         auto print_leak = [](void* k, PyObject* v) {
-            PyTypeObject *tp = Py_TYPE(v);
-            fprintf(stderr, " - leaked instance %p of type \"%s\"\n", k, tp->tp_name);
+            type_data *tp = nb_type_data(Py_TYPE(v));
+            fprintf(stderr, " - leaked instance %p of type \"%s\"\n", k, tp->name);
         };
 
-        for (size_t i = 0; i < internals->shard_count; ++i) {
-            for (auto [k, v]: internals->shards[i].inst_c2p) {
+        for (size_t i = 0; i < p->shard_count; ++i) {
+            for (auto [k, v]: p->shards[i].inst_c2p) {
                 if (NB_UNLIKELY(nb_is_seq(v))) {
                     nb_inst_seq* seq = nb_get_seq(v);
                     for(; seq != nullptr; seq = seq->next)
@@ -288,13 +289,12 @@ static void internals_cleanup() {
         print_leak_warnings = false;
 #endif
 
-    if (!internals->type_c2p_slow.empty() ||
-        !internals->type_c2p_fast.empty()) {
+    if (!p->type_c2p_slow.empty()) {
         if (print_leak_warnings) {
             fprintf(stderr, "nanobind: leaked %zu types!\n",
-                    internals->type_c2p_slow.size());
+                    p->type_c2p_slow.size());
             int ctr = 0;
-            for (const auto &kv : internals->type_c2p_slow) {
+            for (const auto &kv : p->type_c2p_slow) {
                 fprintf(stderr, " - leaked type \"%s\"\n", kv.second->name);
                 if (ctr++ == 10) {
                     fprintf(stderr, " - ... skipped remainder\n");
@@ -305,12 +305,12 @@ static void internals_cleanup() {
         leak = true;
     }
 
-    if (!internals->funcs.empty()) {
+    if (!p->funcs.empty()) {
         if (print_leak_warnings) {
             fprintf(stderr, "nanobind: leaked %zu functions!\n",
-                    internals->funcs.size());
+                    p->funcs.size());
             int ctr = 0;
-            for (auto [f, p] : internals->funcs) {
+            for (auto [f, p] : p->funcs) {
                 fprintf(stderr, " - leaked function \"%s\"\n",
                         nb_func_data(f)->name);
                 if (ctr++ == 10) {
@@ -323,13 +323,19 @@ static void internals_cleanup() {
     }
 
     if (!leak) {
-        nb_translator_seq* t = internals->translators.next;
+        nb_translator_seq* t = p->translators.next;
         while (t) {
             nb_translator_seq *next = t->next;
             delete t;
             t = next;
         }
-        delete internals;
+
+#if defined(NB_FREE_THREADED)
+        PyThread_tss_delete(p->nb_static_property_disabled);
+        PyThread_tss_free(p->nb_static_property_disabled);
+#endif
+
+        delete p;
         internals = nullptr;
         nb_meta_cache = nullptr;
     } else {
@@ -396,6 +402,11 @@ NB_NOINLINE void init(const char *name) {
     p->nb_method = (PyTypeObject *) PyType_FromSpec(&nb_method_spec);
     p->nb_bound_method = (PyTypeObject *) PyType_FromSpec(&nb_bound_method_spec);
 
+#if defined(NB_FREE_THREADED)
+    p->nb_static_property_disabled = PyThread_tss_alloc();
+    PyThread_tss_create(p->nb_static_property_disabled);
+#endif
+
     for (size_t i = 0; i < shard_count; ++i) {
         p->shards[i].keep_alive.min_load_factor(.1f);
         p->shards[i].inst_c2p.min_load_factor(.1f);