Skip to content

Commit b090137

Browse files
committed
Add PyUnstable_TryIncref and PyUnstable_EnableTryIncRef.
1 parent e510a7b commit b090137

File tree

2 files changed

+118
-0
lines changed

2 files changed

+118
-0
lines changed

pythoncapi_compat.h

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2231,6 +2231,76 @@ static inline int PyUnstable_Object_IsUniquelyReferenced(PyObject *obj)
22312231
}
22322232
#endif
22332233

2234+
// gh-128926 added PyUnstable_TryIncRef() and PyUnstable_EnableTryIncRef() to
2235+
// Python 3.14.0a5. Adapted from _Py_TryIncref and _PyObject_SetMaybeWeakref.
2236+
#if PY_VERSION_HEX < 0x030E00A5
2237+
static inline int PyUnstable_TryIncRef(PyObject *op)
2238+
{
2239+
#ifndef Py_GIL_DISABLED
2240+
if (Py_REFCNT(op) > 0) {
2241+
Py_INCREF(op);
2242+
return 1;
2243+
}
2244+
return 0;
2245+
#else
2246+
uint32_t local = _Py_atomic_load_uint32_relaxed(&op->ob_ref_local);
2247+
local += 1;
2248+
if (local == 0) {
2249+
return 1;
2250+
}
2251+
if (_Py_IsOwnedByCurrentThread(op)) {
2252+
_Py_INCREF_STAT_INC();
2253+
_Py_atomic_store_uint32_relaxed(&op->ob_ref_local, local);
2254+
#ifdef Py_REF_DEBUG
2255+
_Py_INCREF_IncRefTotal();
2256+
#endif
2257+
return 1;
2258+
}
2259+
Py_ssize_t shared = _Py_atomic_load_ssize_relaxed(&op->ob_ref_shared);
2260+
for (;;) {
2261+
// If the shared refcount is zero and the object is either merged
2262+
// or may not have weak references, then we cannot incref it.
2263+
if (shared == 0 || shared == _Py_REF_MERGED) {
2264+
return 0;
2265+
}
2266+
2267+
if (_Py_atomic_compare_exchange_ssize(
2268+
&op->ob_ref_shared,
2269+
&shared,
2270+
shared + (1 << _Py_REF_SHARED_SHIFT))) {
2271+
#ifdef Py_REF_DEBUG
2272+
_Py_INCREF_IncRefTotal();
2273+
#endif
2274+
_Py_INCREF_STAT_INC();
2275+
return 1;
2276+
}
2277+
}
2278+
#endif
2279+
}
2280+
2281+
static inline void PyUnstable_EnableTryIncRef(PyObject *op)
2282+
{
2283+
#ifdef Py_GIL_DISABLED
2284+
if (_Py_IsImmortal(op)) {
2285+
return;
2286+
}
2287+
for (;;) {
2288+
Py_ssize_t shared = _Py_atomic_load_ssize_relaxed(&op->ob_ref_shared);
2289+
if ((shared & _Py_REF_SHARED_FLAG_MASK) != 0) {
2290+
// Nothing to do if it's in WEAKREFS, QUEUED, or MERGED states.
2291+
return;
2292+
}
2293+
if (_Py_atomic_compare_exchange_ssize(
2294+
&op->ob_ref_shared, &shared, shared | _Py_REF_MAYBE_WEAKREF)) {
2295+
return;
2296+
}
2297+
}
2298+
#else
2299+
(void)op; // unused argument
2300+
#endif
2301+
}
2302+
#endif
2303+
22342304

22352305
#if PY_VERSION_HEX < 0x030F0000
22362306
static inline PyObject*

tests/test_pythoncapi_compat_cext.c

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2448,6 +2448,46 @@ test_tuple(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
24482448
return test_tuple_fromarray();
24492449
}
24502450

2451+
// Test adapted from CPython's _testcapi/object.c
2452+
static int MyObject_dealloc_called = 0;
2453+
2454+
static void
2455+
MyObject_dealloc(PyObject *op)
2456+
{
2457+
// PyUnstable_TryIncRef should return 0 if object is being deallocated
2458+
assert(Py_REFCNT(op) == 0);
2459+
assert(!PyUnstable_TryIncRef(op));
2460+
assert(Py_REFCNT(op) == 0);
2461+
2462+
MyObject_dealloc_called++;
2463+
Py_TYPE(op)->tp_free(op);
2464+
}
2465+
2466+
static PyTypeObject MyType;
2467+
2468+
static PyObject*
2469+
test_try_incref(PyObject *Py_UNUSED(module), PyObject *Py_UNUSED(args))
2470+
{
2471+
MyObject_dealloc_called = 0;
2472+
2473+
PyObject *obj = PyObject_New(PyObject, &MyType);
2474+
if (obj == _Py_NULL) {
2475+
return _Py_NULL;
2476+
}
2477+
2478+
PyUnstable_EnableTryIncRef(obj);
2479+
2480+
Py_ssize_t refcount = Py_REFCNT(obj);
2481+
assert(PyUnstable_TryIncRef(obj));
2482+
assert(Py_REFCNT(obj) == refcount + 1);
2483+
2484+
Py_DECREF(obj);
2485+
Py_DECREF(obj);
2486+
2487+
assert(MyObject_dealloc_called == 1);
2488+
Py_RETURN_NONE;
2489+
}
2490+
24512491

24522492
static struct PyMethodDef methods[] = {
24532493
{"test_object", test_object, METH_NOARGS, _Py_NULL},
@@ -2504,6 +2544,7 @@ static struct PyMethodDef methods[] = {
25042544
{"test_uniquely_referenced", test_uniquely_referenced, METH_NOARGS, _Py_NULL},
25052545
{"test_byteswriter", test_byteswriter, METH_NOARGS, _Py_NULL},
25062546
{"test_tuple", test_tuple, METH_NOARGS, _Py_NULL},
2547+
{"test_try_incref", test_try_incref, METH_NOARGS, _Py_NULL},
25072548
{_Py_NULL, _Py_NULL, 0, _Py_NULL}
25082549
};
25092550

@@ -2532,6 +2573,13 @@ module_exec(PyObject *module)
25322573
return -1;
25332574
}
25342575
#endif
2576+
MyType.tp_name = "MyType";
2577+
MyType.tp_basicsize = sizeof(PyObject);
2578+
MyType.tp_dealloc = MyObject_dealloc;
2579+
MyType.tp_free = PyObject_Del;
2580+
if (PyType_Ready(&MyType) < 0) {
2581+
return -1;
2582+
}
25352583
return 0;
25362584
}
25372585

0 commit comments

Comments
 (0)