Fix data race when using shared variables (free threading) (#5494)

* Fix data race when using shared variables (free threading)

In the free threading build, there's a race between wrapper re-use and
wrapper deallocation. This can happen with a static variable accessed by
multiple threads.

Fixing this requires using some private CPython APIs: _Py_TryIncref and
_PyObject_SetMaybeWeakref. The implementations of these functions are
included until they're made available as public ("unstable") APIs.

Fixes #5489

* style: pre-commit fixes

* Avoid unused parameter

* Add missing return for default build

* Changes from review

* Assign result to local variable

* s/clang-tidy/ruff

* clang-tidy: static is redundant

* Use 'noqa: B018'

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

include/pybind11/detail/class.h

@@ -312,7 +312,31 @@ inline void traverse_offset_bases(void *valueptr,
     }
 }
 
+#ifdef Py_GIL_DISABLED
+inline void enable_try_inc_ref(PyObject *obj) {
+    // TODO: Replace with PyUnstable_Object_EnableTryIncRef when available.
+    // See https://github.com/python/cpython/issues/128844
+    if (_Py_IsImmortal(obj)) {
+        return;
+    }
+    for (;;) {
+        Py_ssize_t shared = _Py_atomic_load_ssize_relaxed(&obj->ob_ref_shared);
+        if ((shared & _Py_REF_SHARED_FLAG_MASK) != 0) {
+            // Nothing to do if it's in WEAKREFS, QUEUED, or MERGED states.
+            return;
+        }
+        if (_Py_atomic_compare_exchange_ssize(
+                &obj->ob_ref_shared, &shared, shared | _Py_REF_MAYBE_WEAKREF)) {
+            return;
+        }
+    }
+}
+#endif
+
 inline bool register_instance_impl(void *ptr, instance *self) {
+#ifdef Py_GIL_DISABLED
+    enable_try_inc_ref(reinterpret_cast<PyObject *>(self));
+#endif
     with_instance_map(ptr, [&](instance_map &instances) { instances.emplace(ptr, self); });
     return true; // unused, but gives the same signature as the deregister func
 }

include/pybind11/detail/type_caster_base.h

@@ -241,6 +241,49 @@ PYBIND11_NOINLINE handle get_type_handle(const std::type_info &tp, bool throw_if
     return handle(type_info ? ((PyObject *) type_info->type) : nullptr);
 }
 
+inline bool try_incref(PyObject *obj) {
+    // Tries to increment the reference count of an object if it's not zero.
+    // TODO: Use PyUnstable_TryIncref when available.
+    // See https://github.com/python/cpython/issues/128844
+#ifdef Py_GIL_DISABLED
+    // See
+    // https://github.com/python/cpython/blob/d05140f9f77d7dfc753dd1e5ac3a5962aaa03eff/Include/internal/pycore_object.h#L761
+    uint32_t local = _Py_atomic_load_uint32_relaxed(&obj->ob_ref_local);
+    local += 1;
+    if (local == 0) {
+        // immortal
+        return true;
+    }
+    if (_Py_IsOwnedByCurrentThread(obj)) {
+        _Py_atomic_store_uint32_relaxed(&obj->ob_ref_local, local);
+#    ifdef Py_REF_DEBUG
+        _Py_INCREF_IncRefTotal();
+#    endif
+        return true;
+    }
+    Py_ssize_t shared = _Py_atomic_load_ssize_relaxed(&obj->ob_ref_shared);
+    for (;;) {
+        // If the shared refcount is zero and the object is either merged
+        // or may not have weak references, then we cannot incref it.
+        if (shared == 0 || shared == _Py_REF_MERGED) {
+            return false;
+        }
+
+        if (_Py_atomic_compare_exchange_ssize(
+                &obj->ob_ref_shared, &shared, shared + (1 << _Py_REF_SHARED_SHIFT))) {
+#    ifdef Py_REF_DEBUG
+            _Py_INCREF_IncRefTotal();
+#    endif
+            return true;
+        }
+    }
+#else
+    assert(Py_REFCNT(obj) > 0);
+    Py_INCREF(obj);
+    return true;
+#endif
+}
+
 // Searches the inheritance graph for a registered Python instance, using all_type_info().
 PYBIND11_NOINLINE handle find_registered_python_instance(void *src,
                                                          const detail::type_info *tinfo) {
@@ -249,7 +292,10 @@ PYBIND11_NOINLINE handle find_registered_python_instance(void *src,
         for (auto it_i = it_instances.first; it_i != it_instances.second; ++it_i) {
             for (auto *instance_type : detail::all_type_info(Py_TYPE(it_i->second))) {
                 if (instance_type && same_type(*instance_type->cpptype, *tinfo->cpptype)) {
-                    return handle((PyObject *) it_i->second).inc_ref();
+                    auto *wrapper = reinterpret_cast<PyObject *>(it_i->second);
+                    if (try_incref(wrapper)) {
+                        return handle(wrapper);
+                    }
                 }
             }
         }

tests/test_thread.cpp

@@ -28,6 +28,9 @@ struct IntStruct {
     int value;
 };
 
+struct EmptyStruct {};
+EmptyStruct SharedInstance;
+
 } // namespace
 
 TEST_SUBMODULE(thread, m) {
@@ -61,6 +64,9 @@ TEST_SUBMODULE(thread, m) {
         },
         py::call_guard<py::gil_scoped_release>());
 
+    py::class_<EmptyStruct>(m, "EmptyStruct")
+        .def_readonly_static("SharedInstance", &SharedInstance);
+
     // NOTE: std::string_view also uses loader_life_support to ensure that
     // the string contents remain alive, but that's a C++ 17 feature.
 }

tests/test_thread.py

@@ -47,3 +47,22 @@ def test_implicit_conversion_no_gil():
         x.start()
     for x in [c, b, a]:
         x.join()
+
+
+@pytest.mark.skipif(sys.platform.startswith("emscripten"), reason="Requires threads")
+def test_bind_shared_instance():
+    nb_threads = 4
+    b = threading.Barrier(nb_threads)
+
+    def access_shared_instance():
+        b.wait()
+        for _ in range(1000):
+            m.EmptyStruct.SharedInstance  # noqa: B018
+
+    threads = [
+        threading.Thread(target=access_shared_instance) for _ in range(nb_threads)
+    ]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join()