Add PyUnstable_TryIncref() and PyUnstable_EnableTryIncRef() (#159)

Co-authored-by: Victor Stinner <[email protected]>

Author: colesbury

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*

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;
 }