Skip to content

Commit

Permalink
Fix data race when using shared variables (free threading)
Browse files Browse the repository at this point in the history
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 pybind#5489
  • Loading branch information
colesbury committed Jan 15, 2025
1 parent 945e251 commit e08cfa9
Show file tree
Hide file tree
Showing 4 changed files with 93 additions and 1 deletion.
22 changes: 22 additions & 0 deletions include/pybind11/detail/class.h
Original file line number Diff line number Diff line change
Expand Up @@ -312,7 +312,29 @@ inline void traverse_offset_bases(void *valueptr,
}
}

static inline void enable_try_inc_ref(PyObject *op) {
#ifdef Py_GIL_DISABLED
// TODO: Replace with PyUnstable_Object_EnableTryIncRef when available.
// See https://github.com/python/cpython/issues/128844
if (_Py_IsImmortal(op)) {
return;
}
for (;;) {
Py_ssize_t shared = _Py_atomic_load_ssize_relaxed(&op->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(
&op->ob_ref_shared, &shared, shared | _Py_REF_MAYBE_WEAKREF)) {
return;
}
}
#endif
}

inline bool register_instance_impl(void *ptr, instance *self) {
enable_try_inc_ref((PyObject *) self);
with_instance_map(ptr, [&](instance_map &instances) { instances.emplace(ptr, self); });
return true; // unused, but gives the same signature as the deregister func
}
Expand Down
48 changes: 47 additions & 1 deletion include/pybind11/detail/type_caster_base.h
Original file line number Diff line number Diff line change
Expand Up @@ -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);
#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) {
Expand All @@ -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();
PyObject *wrapper = (PyObject *) it_i->second;
if (try_incref(wrapper)) {
return handle(wrapper);
}
}
}
}
Expand Down
6 changes: 6 additions & 0 deletions tests/test_thread.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,9 @@ struct IntStruct {
int value;
};

struct EmptyStruct {};
static EmptyStruct SharedInstance;

} // namespace

TEST_SUBMODULE(thread, m) {
Expand Down Expand Up @@ -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.
}
18 changes: 18 additions & 0 deletions tests/test_thread.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,3 +47,21 @@ 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):
x = m.EmptyStruct.SharedInstance
del x

threads = [threading.Thread(target=access_shared_instance) for _ in range(nb_threads)]
for thread in threads:
thread.start()
for thread in threads:
thread.join()

0 comments on commit e08cfa9

Please sign in to comment.