Deferred reference counting

Initial partial implementation of deferred reference counting.

Author: colesbury

Include/internal/pycore_object.h

@@ -360,6 +360,34 @@ _PyObject_SetMaybeWeakref(PyObject *op)
     }
 }
 
+/* Marks the object as support deferred reference counting.
+ *
+ * The object's type must be GC-enabled. This function is not thread-safe with
+ * respect to concurrent modifications; it must be called before the object
+ * becomes visible to other threads.
+ *
+ * Deferred refcounted objects are marked as "queued" to prevent merging
+ * reference count fields outside the garbage collector.
+ */
+static inline void
+_PyObject_SetDeferredRefcount(PyObject *op)
+{
+    assert(_Py_ThreadLocal(op) && "non thread-safe");
+    assert(!_PyObject_HasDeferredRefcount(op) && "already uses deferred refcounting");
+    assert(PyType_IS_GC(Py_TYPE(op)));
+    op->ob_ref_local += _Py_REF_DEFERRED_MASK + 1;
+    op->ob_ref_shared = (op->ob_ref_shared & ~_Py_REF_SHARED_FLAG_MASK) | _Py_REF_QUEUED;
+}
+
+#define _PyObject_SET_DEFERRED_REFCOUNT(op) _PyObject_SetDeferredRefcount(_PyObject_CAST(op))
+
+// Check is refcount is deferred or immortal
+static inline int
+_Py_REF_NON_IMMEDIATE(uint32_t local)
+{
+    return _Py_STATIC_CAST(int32_t, local) <= Py_REF_IMMORTAL;
+}
+
 #ifdef Py_REF_DEBUG
 extern void _PyDebug_PrintTotalRefs(void);
 #endif

Include/object.h

@@ -135,7 +135,7 @@ PyAPI_FUNC(int) Py_Is(PyObject *x, PyObject *y);
 #define Py_Is(x, y) ((x) == (y))
 
 static inline void
-_PyRef_UnpackLocal(uint32_t bits, Py_ssize_t *refcount, int *immortal);
+_PyRef_UnpackLocal(uint32_t bits, Py_ssize_t *refcount, int *immortal, int *deferred);
 
 static inline void
 _PyRef_UnpackShared(uint32_t bits, Py_ssize_t *refcount, int *queued, int *merged);
@@ -148,7 +148,7 @@ static inline Py_ssize_t Py_REFCNT(PyObject *ob) {
     int immortal;
 
     uint32_t local = _Py_atomic_load_uint32_relaxed(&ob->ob_ref_local);
-    _PyRef_UnpackLocal(local, &local_refcount, &immortal);
+    _PyRef_UnpackLocal(local, &local_refcount, &immortal, NULL);
 
     if (immortal) {
         return 999;
@@ -617,7 +617,7 @@ _Py_ThreadLocal(PyObject *op)
 #define _Py_REF_LOCAL_SHIFT     0
 #define _Py_REF_LOCAL_INIT      1
 #define Py_REF_IMMORTAL         -1
-#define _Py_REF_DEFERRED_MASK   0x40000000
+#define _Py_REF_DEFERRED_MASK 0xC0000000
 
 // 30         2
 // [refcount] [bits]
@@ -884,10 +884,17 @@ static inline PyObject* _Py_XNewRef(PyObject *obj)
 #endif
 
 static inline void
-_PyRef_UnpackLocal(uint32_t bits, Py_ssize_t *refcount, int *immortal)
+_PyRef_UnpackLocal(uint32_t bits, Py_ssize_t *refcount, int *immortal, int *deferred)
 {
     *refcount = bits >> _Py_REF_LOCAL_SHIFT;
     *immortal = _Py_REF_IS_IMMORTAL(bits);
+    int is_deferred = _Py_STATIC_CAST(int32_t, bits) < Py_REF_IMMORTAL;
+    if (is_deferred) {
+        *refcount -= _Py_REF_DEFERRED_MASK;
+    }
+    if (deferred) {
+        *deferred = is_deferred;
+    }
 }
 
 static inline void
@@ -925,7 +932,7 @@ _PyObject_IsReferened(PyObject *op)
     int immortal;
 
     uint32_t local = _Py_atomic_load_uint32_relaxed(&op->ob_ref_local);
-    _PyRef_UnpackLocal(local, &local_refcount, &immortal);
+    _PyRef_UnpackLocal(local, &local_refcount, &immortal, NULL);
 
     if (immortal) {
         return 1;
@@ -940,6 +947,13 @@ _PyObject_IsReferened(PyObject *op)
     return (local_refcount + shared_refcount) > 0;
 }
 
+static inline int
+_PyObject_HasDeferredRefcount(PyObject *op)
+{
+    uint32_t local = _Py_atomic_load_uint32_relaxed(&op->ob_ref_local);
+    return _Py_STATIC_CAST(int32_t, local) < Py_REF_IMMORTAL;
+}
+
 static inline void
 _Py_RESURRECT(PyObject *op, Py_ssize_t refcnt)
 {

Modules/gcmodule.c

@@ -318,13 +318,13 @@ static Py_ssize_t
 _Py_GC_REFCNT(PyObject *op)
 {
     Py_ssize_t local, shared;
-    int immortal;
+    int immortal, deferred;
 
-    _PyRef_UnpackLocal(op->ob_ref_local, &local, &immortal);
+    _PyRef_UnpackLocal(op->ob_ref_local, &local, &immortal, &deferred);
     _PyRef_UnpackShared(op->ob_ref_shared, &shared, NULL, NULL);
     assert(!immortal);
 
-    return local + shared;
+    return local + shared - deferred;
 }
 
 typedef int (gc_visit_fn)(PyGC_Head* gc, void *arg);
@@ -682,9 +682,35 @@ find_dead_shared_keys(_PyObjectQueue **queue, int *num_unmarked)
     }
 }
 
+static void
+merge_refcount(PyObject *op, Py_ssize_t extra)
+{
+    Py_ssize_t local_refcount, shared_refcount;
+    int immortal, deferred;
+
+    assert(_PyRuntime.stop_the_world);
+
+    _PyRef_UnpackLocal(op->ob_ref_local, &local_refcount, &immortal, &deferred);
+    _PyRef_UnpackShared(op->ob_ref_shared, &shared_refcount, NULL, NULL);
+    assert(!immortal && "immortal objects should not be in garbage");
+
+    Py_ssize_t refcount = local_refcount + shared_refcount;
+    refcount += extra;
+    refcount -= deferred;
+
+#ifdef Py_REF_DEBUG
+    _Py_IncRefTotalN(extra);
+#endif
+
+    op->ob_tid = 0;
+    op->ob_ref_local = 0;
+    op->ob_ref_shared = _Py_REF_PACK_SHARED(refcount, _Py_REF_MERGED);
+}
+
 struct update_refs_args {
     PyGC_Head *list;
     int split_keys_marked;
+    _PyGC_Reason gc_reason;
 };
 
 // Compute the number of external references to objects in the heap
@@ -730,6 +756,15 @@ update_refs(const mi_heap_t* heap, const mi_heap_area_t* area, void* block, size
         }
     }
 
+    if (arg->gc_reason == GC_REASON_SHUTDOWN) {
+        if (_PyObject_HasDeferredRefcount(op)) {
+            // Disable deferred reference counting when we're shutting down.
+            // This is useful for interp->sysdict because the last reference
+            // to it is cleared after the last GC cycle.
+            merge_refcount(op, 0);
+        }
+    }
+
     // Add the actual refcount to gc_refs.
     Py_ssize_t refcount = _Py_GC_REFCNT(op);
     _PyObject_ASSERT(op, refcount >= 0);
@@ -1008,23 +1043,7 @@ move_legacy_finalizer_reachable(PyGC_Head *finalizers)
 static void
 incref_merge(PyObject *op)
 {
-    Py_ssize_t local_refcount, shared_refcount;
-    int immortal;
-
-    assert(_PyRuntime.stop_the_world);
-
-#ifdef Py_REF_DEBUG
-    _Py_IncRefTotal();
-#endif
-
-    _PyRef_UnpackLocal(op->ob_ref_local, &local_refcount, &immortal);
-    _PyRef_UnpackShared(op->ob_ref_shared, &shared_refcount, NULL, NULL);
-    assert(!immortal && "immortal objects should not be in garbage");
-
-    Py_ssize_t refcount = local_refcount + shared_refcount + 1;
-    op->ob_tid = 0;
-    op->ob_ref_local = 0;
-    op->ob_ref_shared = _Py_REF_PACK_SHARED(refcount, _Py_REF_MERGED);
+    merge_refcount(op, 1);
 }
 
 /* Subtracts one from the refcount field. */
@@ -1594,7 +1613,11 @@ gc_collect_main(PyThreadState *tstate, int generation, _PyGC_Reason reason)
     validate_tracked_heap(_PyGC_PREV_MASK|_PyGC_PREV_MASK_UNREACHABLE, 0);
 
     gc_list_init(&young);
-    struct update_refs_args args = { .list = &young, .split_keys_marked = 0 };
+    struct update_refs_args args = {
+        .list = &young,
+        .split_keys_marked = 0,
+        .gc_reason = reason,
+    };
     visit_heaps2(mi_heap_tag_gc, update_refs, &args);
 
     _PyObjectQueue *dead_keys = NULL;

Objects/descrobject.c

@@ -903,6 +903,7 @@ descr_new(PyTypeObject *descrtype, PyTypeObject *type, const char *name)
 
     descr = (PyDescrObject *)PyType_GenericAlloc(descrtype, 0);
     if (descr != NULL) {
+        _PyObject_SET_DEFERRED_REFCOUNT(descr);
         descr->d_type = (PyTypeObject*)Py_XNewRef(type);
         descr->d_name = PyUnicode_InternFromString(name);
         if (descr->d_name == NULL) {

Objects/dictobject.c

@@ -1371,6 +1371,10 @@ _PyDict_MaybeUntrack(PyObject *op)
     if (!PyDict_CheckExact(op) || !_PyObject_GC_IS_TRACKED(op))
         return;
 
+    if (_PyObject_HasDeferredRefcount(op)) {
+        return;
+    }
+
     mp = (PyDictObject *) op;
     numentries = mp->ma_keys->dk_nentries;
     if (_PyDict_HasSplitTable(mp)) {

Objects/funcobject.c

@@ -77,7 +77,9 @@ PyFunction_ClearWatcher(int watcher_id)
 PyFunctionObject *
 _PyFunction_FromConstructor(PyFrameConstructor *constr)
 {
-
+    // NOTE(sgross): functions created via FrameConstructor do *not* use
+    // deferred reference counting because they are typically not part of cycles
+    // nor accessed by multiple threads.
     PyFunctionObject *op = PyObject_GC_New(PyFunctionObject, &PyFunction_Type);
     if (op == NULL) {
         return NULL;
@@ -172,6 +174,9 @@ PyFunction_NewWithQualName(PyObject *code, PyObject *globals, PyObject *qualname
     op->func_annotations = NULL;
     op->vectorcall = _PyFunction_Vectorcall;
     op->func_version = 0;
+    if ((code_obj->co_flags & CO_NESTED) == 0) {
+        _PyObject_SET_DEFERRED_REFCOUNT(op);
+    }
     _PyObject_GC_TRACK(op);
     handle_func_event(PyFunction_EVENT_CREATE, op, NULL);
     return (PyObject *)op;

Objects/moduleobject.c

@@ -107,6 +107,9 @@ new_module_notrack(PyTypeObject *mt)
     m->md_name = NULL;
     m->md_dict = PyDict_New();
     if (m->md_dict != NULL) {
+        PyObject_GC_Track(m->md_dict);
+        _PyObject_SET_DEFERRED_REFCOUNT(m->md_dict);
+        _PyObject_SET_DEFERRED_REFCOUNT(m);
         return m;
     }
     Py_DECREF(m);

Objects/typeobject.c

@@ -2887,6 +2887,7 @@ type_new_alloc(type_new_ctx *ctx)
     if (type == NULL) {
         return NULL;
     }
+    _PyObject_SET_DEFERRED_REFCOUNT(type);
     PyHeapTypeObject *et = (PyHeapTypeObject *)type;
 
     // Initialize tp_flags.
@@ -3775,6 +3776,7 @@ PyType_FromMetaclass(PyTypeObject *metaclass, PyObject *module,
     if (res == NULL) {
         goto finally;
     }
+    _PyObject_SET_DEFERRED_REFCOUNT(res);
     res_start = (char*)res;
 
     type = &res->ht_type;

Python/sysmodule.c

@@ -1876,16 +1876,17 @@ sys_getfullrefcount(PyObject *module, PyObject *object)
     uintptr_t tid = _PyObject_ThreadId(object);
 
     Py_ssize_t local, shared;
-    int immortal, queued, merged;
+    int immortal, deferred, queued, merged;
 
-    _PyRef_UnpackLocal(object->ob_ref_local, &local, &immortal);
+    _PyRef_UnpackLocal(object->ob_ref_local, &local, &immortal, &deferred);
     _PyRef_UnpackShared(object->ob_ref_shared, &shared, &queued, &merged);
 
     PyDict_SetItemString(res, "local", PyLong_FromSsize_t(local));
     PyDict_SetItemString(res, "shared", PyLong_FromSsize_t(shared));
     PyDict_SetItemString(res, "merged", PyBool_FromLong(merged));
     PyDict_SetItemString(res, "queued", PyBool_FromLong(queued));
     PyDict_SetItemString(res, "immortal", PyBool_FromLong(immortal));
+    PyDict_SetItemString(res, "deferred", PyBool_FromLong(deferred));
     if (queued) {
         Py_INCREF(Py_None);
         PyDict_SetItemString(res, "tid", Py_None);