diff --git a/pythoncapi_compat.h b/pythoncapi_compat.h index b16075f..04378bd 100644 --- a/pythoncapi_compat.h +++ b/pythoncapi_compat.h @@ -2231,6 +2231,81 @@ static inline int PyUnstable_Object_IsUniquelyReferenced(PyObject *obj) } #endif +// gh-128926 added PyUnstable_TryIncRef() and PyUnstable_EnableTryIncRef() to +// Python 3.14.0a5. Adapted from _Py_TryIncref() and _PyObject_SetMaybeWeakref(). +#if PY_VERSION_HEX < 0x030E00A5 +static inline int PyUnstable_TryIncRef(PyObject *op) +{ +#ifndef Py_GIL_DISABLED + if (Py_REFCNT(op) > 0) { + Py_INCREF(op); + return 1; + } + return 0; +#else + // _Py_TryIncrefFast() + uint32_t local = _Py_atomic_load_uint32_relaxed(&op->ob_ref_local); + local += 1; + if (local == 0) { + // immortal + return 1; + } + if (_Py_IsOwnedByCurrentThread(op)) { + _Py_INCREF_STAT_INC(); + _Py_atomic_store_uint32_relaxed(&op->ob_ref_local, local); +#ifdef Py_REF_DEBUG + _Py_INCREF_IncRefTotal(); +#endif + return 1; + } + + // _Py_TryIncRefShared() + Py_ssize_t shared = _Py_atomic_load_ssize_relaxed(&op->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 0; + } + + if (_Py_atomic_compare_exchange_ssize( + &op->ob_ref_shared, + &shared, + shared + (1 << _Py_REF_SHARED_SHIFT))) { +#ifdef Py_REF_DEBUG + _Py_INCREF_IncRefTotal(); +#endif + _Py_INCREF_STAT_INC(); + return 1; + } + } +#endif +} + +static inline void PyUnstable_EnableTryIncRef(PyObject *op) +{ +#ifdef Py_GIL_DISABLED + // _PyObject_SetMaybeWeakref() + 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; + } + } +#else + (void)op; // unused argument +#endif +} +#endif + #if PY_VERSION_HEX < 0x030F0000 static inline PyObject* diff --git a/tests/test_pythoncapi_compat_cext.c b/tests/test_pythoncapi_compat_cext.c index fed0822..7d9fea4 100644 --- a/tests/test_pythoncapi_compat_cext.c +++ b/tests/test_pythoncapi_compat_cext.c @@ -2448,6 +2448,46 @@ test_tuple(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args)) return test_tuple_fromarray(); } +// Test adapted from CPython's _testcapi/object.c +static int TryIncref_dealloc_called = 0; + +static void +TryIncref_dealloc(PyObject *op) +{ + // PyUnstable_TryIncRef should return 0 if object is being deallocated + assert(Py_REFCNT(op) == 0); + assert(!PyUnstable_TryIncRef(op)); + assert(Py_REFCNT(op) == 0); + + TryIncref_dealloc_called++; + Py_TYPE(op)->tp_free(op); +} + +static PyTypeObject TryIncrefType; + +static PyObject* +test_try_incref(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args)) +{ + TryIncref_dealloc_called = 0; + + PyObject *obj = PyObject_New(PyObject, &TryIncrefType); + if (obj == _Py_NULL) { + return _Py_NULL; + } + + PyUnstable_EnableTryIncRef(obj); + + Py_ssize_t refcount = Py_REFCNT(obj); + assert(PyUnstable_TryIncRef(obj)); + assert(Py_REFCNT(obj) == refcount + 1); + + Py_DECREF(obj); + Py_DECREF(obj); + + assert(TryIncref_dealloc_called == 1); + Py_RETURN_NONE; +} + static struct PyMethodDef methods[] = { {"test_object", test_object, METH_NOARGS, _Py_NULL}, @@ -2504,6 +2544,7 @@ static struct PyMethodDef methods[] = { {"test_uniquely_referenced", test_uniquely_referenced, METH_NOARGS, _Py_NULL}, {"test_byteswriter", test_byteswriter, METH_NOARGS, _Py_NULL}, {"test_tuple", test_tuple, METH_NOARGS, _Py_NULL}, + {"test_try_incref", test_try_incref, METH_NOARGS, _Py_NULL}, {_Py_NULL, _Py_NULL, 0, _Py_NULL} }; @@ -2532,6 +2573,13 @@ module_exec(PyObject *module) return -1; } #endif + TryIncrefType.tp_name = "TryIncrefType"; + TryIncrefType.tp_basicsize = sizeof(PyObject); + TryIncrefType.tp_dealloc = TryIncref_dealloc; + TryIncrefType.tp_free = PyObject_Del; + if (PyType_Ready(&TryIncrefType) < 0) { + return -1; + } return 0; }