Tracing region

Include/cpython/weakrefobject.h

@@ -2,6 +2,18 @@
 #  error "this header file must not be included directly"
 #endif
 
+/* A region reference is a weak reference that survives its target's region
+ * being closed. Instead of keeping the region open it checks on every
+ * dereference whether this interpreter may reach the target, and opens the
+ * region tree if it may. The metadata carrying that information lives in
+ * `pycore_regionref.h`; it is opaque here.
+ */
+PyAPI_DATA(PyTypeObject) _PyRegionref_RefType;
+
+#define _PyRegionRef_CheckExact(op) Py_IS_TYPE((op), &_PyRegionref_RefType)
+
+struct _PyRegionRefMetadata;
+
 /* PyWeakReference is the base struct for the Python ReferenceType, ProxyType,
  * and CallableProxyType.
  */
@@ -43,12 +55,23 @@ struct _PyWeakReference {
      */
     PyMutex *weakrefs_lock;
 #endif
+
+    /* The ownership domain of `wr_object`, or NULL if this object doesn't have an ownership
+     * domain. This can happen if this is a normal weakref or if the object is immutable.
+     */
+    struct _PyRegionRefMetadata *region_ref;
 };
 
 PyAPI_FUNC(void) _PyWeakref_ClearRef(PyWeakReference *self);
 
+/* Region references reuse this struct but are deliberately not a subtype of
+ * `_PyWeakref_RefType`, so that `PyWeakref_Check()` stays false for them and
+ * the region close trace does not follow them. */
+#define _PyWeakrefOrRegionRef_Check(op) \
+    (PyWeakref_Check(op) || _PyRegionRef_CheckExact(op))
+
 #define _PyWeakref_CAST(op) \
-    (assert(PyWeakref_Check(op)), _Py_CAST(PyWeakReference*, (op)))
+    (assert(_PyWeakrefOrRegionRef_Check(op)), _Py_CAST(PyWeakReference*, (op)))
 
 // Test if a weak reference is dead.
 PyAPI_FUNC(int) PyWeakref_IsDead(PyObject *ref);

Include/internal/pycore_cown.h

@@ -0,0 +1,39 @@
+#ifndef Py_INTERNAL_COWN_H
+#define Py_INTERNAL_COWN_H
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#ifndef Py_BUILD_CORE
+#  error "Py_BUILD_CORE must be defined to include this header"
+#endif
+
+#include "object.h"
+#include "exports.h"
+
+typedef struct _PyCownObject _PyCownObject;
+#define _PyCownObject_CAST(op) _Py_CAST(_PyCownObject*, op)
+
+PyAPI_DATA(PyTypeObject) _PyCown_Type;
+
+typedef uint64_t _PyCown_ipid_t;
+typedef uint64_t _PyCown_thread_id_t;
+
+PyAPI_FUNC(_PyCown_ipid_t) _PyCown_ThisInterpreterId(void);
+PyAPI_FUNC(_PyCown_thread_id_t) _PyCown_ThisThreadId(void);
+
+/* The interpreter currently owning the cown, or `_PyCown_ReleasedIpid()` when
+ * no interpreter does. Safe to call from any interpreter. */
+PyAPI_FUNC(_PyCown_ipid_t) _PyCown_Owner(PyObject *cown);
+PyAPI_FUNC(_PyCown_ipid_t) _PyCown_ReleasedIpid(void);
+
+/* The thread that acquired the cown, or `_PyCown_UnsetThreadId()` when it was
+ * acquired without the GIL. Not enforced, only reported. */
+PyAPI_FUNC(_PyCown_thread_id_t) _PyCown_LockingThread(PyObject *cown);
+PyAPI_FUNC(_PyCown_thread_id_t) _PyCown_UnsetThreadId(void);
+
+
+#ifdef __cplusplus
+}
+#endif
+#endif /* !Py_INTERNAL_COWN_H */

Include/internal/pycore_gc.h

@@ -352,6 +352,25 @@ extern PyObject *_PyGC_GetObjects(PyInterpreterState *interp, int generation);
 extern PyObject *_PyGC_GetReferrers(PyInterpreterState *interp, PyObject *objs);
 
 // Functions to clear types free lists
+/* Disposal of a list of objects that are known to be unreachable. Used by the
+ * collector itself and by anything else that owns a set of objects it has
+ * established to be garbage, such as a closed tracing region.
+ *
+ * `_PyGC_FinalizeGarbage()` runs the finalizer of every object in `collectable`,
+ * before anything is cleared, so that a `__del__` still sees its object intact.
+ *
+ * `_PyGC_DeleteGarbage()` then breaks the references between them, deallocating
+ * every object whose reference count reaches zero. Objects that a finalizer kept
+ * alive are moved to `old` instead.
+ *
+ * Neither may be called with an exception set. Only available in the default
+ * build; the free-threaded collector has its own implementation.
+ */
+#ifndef Py_GIL_DISABLED
+extern void _PyGC_FinalizeGarbage(PyGC_Head *collectable);
+extern void _PyGC_DeleteGarbage(PyGC_Head *collectable, PyGC_Head *old);
+#endif
+
 extern void _PyGC_ClearAllFreeLists(PyInterpreterState *interp);
 extern void _Py_ScheduleGC(PyThreadState *tstate);
 extern void _Py_RunGC(PyThreadState *tstate);

Include/internal/pycore_immutability.h

@@ -8,6 +8,33 @@ extern "C" {
 #  error "Py_BUILD_CORE must be defined to include this header"
 #endif
 
+struct _PyRegionRefMetadata;
+
+PyAPI_DATA(PyTypeObject) _PyTracingRegion_Type;
+PyAPI_FUNC(int) _PyTracingRegion_Close(PyObject* region);
+PyAPI_FUNC(int) _PyTracingRegion_IsClosed(PyObject* region);
+PyAPI_FUNC(void) _PyTracingRegion_Open(PyObject* region);
+
+/* Returns the region's metadata node, allocating it if this is the first
+ * region reference the current close has found. Borrowed, and only valid while
+ * the region stays closed. The caller must hold `_PyWeakref_Lock`. */
+PyAPI_FUNC(struct _PyRegionRefMetadata*)
+    _PyTracingRegion_MetaLockHeld(PyObject* region);
+
+/* Hands a closed region's node to `cown`, used when a cown takes ownership of
+ * the region. Does nothing for an open region. */
+PyAPI_FUNC(void) _PyTracingRegion_SetMetaCown(PyObject* region, PyObject* cown);
+
+/* Records who owns a closed region, used when it leaves the cown that owned it.
+ * The owner is NOT necessarily the calling interpreter: a cown is immutable and
+ * may be deallocated by anyone holding a reference, including an interpreter
+ * that never owned it. Pass `_PyCown_ReleasedIpid()` when nobody owns it.
+ * Does nothing for an open region. */
+// FIXME: The deallocation will be fixed in a follow-up, then we can remove the
+// owner argument and assert that it's always local.
+PyAPI_FUNC(void) _PyTracingRegion_SetMetaOwner(
+    PyObject* region, uint64_t owner);
+
 struct _Py_immutability_state {
     int late_init_done;
     struct _Py_hashtable_t *shallow_immutable_types;

Include/internal/pycore_regionref.h

@@ -0,0 +1,86 @@
+#ifndef Py_INTERNAL_REGIONREF_H
+#define Py_INTERNAL_REGIONREF_H
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#ifndef Py_BUILD_CORE
+#  error "Py_BUILD_CORE must be defined to include this header"
+#endif
+
+#include "pycore_cown.h"          // _PyCown_ipid_t
+#include "pycore_hashtable.h"     // _Py_hashtable_t
+
+/* Every `RegionRef` points at a `_PyRegionRefMetadata` node, and those nodes
+ * form a tree that mirrors the region hierarchy: a nested region's node
+ * delegates to the node of its parent, and the outermost node names the owner,
+ * either an interpreter or the cown holding the region. A cown node looks the
+ * owner up on the cown, which is what makes acquiring and releasing free;
+ * moving a closed region between owners restamps a single node.
+ *
+ * A reference to an object that is in no region at all owns its own node.
+ */
+
+typedef enum {
+    /* A close is in progress. The referenced data is unavailable but should be
+     * soon; dereferencing fails and the caller may retry. Only ever reachable
+     * through a child's `parent`, never directly from a reference. */
+    _Py_REGION_REF_WIP,
+    /* Delegates to `value.parent`, the node of the enclosing region. */
+    _Py_REGION_REF_META,
+    /* Terminal. The region is held by `value.cown`, on which the owner is
+     * looked up dynamically. */
+    _Py_REGION_REF_COWN,
+    /* Terminal, owned by one interpreter. */
+    _Py_REGION_REF_IPID,
+} _PyRegionRefKind;
+
+typedef struct _PyRegionRefMetadata {
+    uint32_t rc;
+    /* A `_PyRegionRefKind`. */
+    uint8_t kind;
+    /* Borrowed. Set only on region nodes, and only while that region is
+     * closed. Names the region a dereference has to open on its way down.
+     *
+     * Borrowing is safe because the pointer is only followed after the terminal
+     * check established that this interpreter owns the region, and a region can
+     * only be deallocated by its owner. */
+    PyObject *region;
+    union {
+        struct _PyRegionRefMetadata *parent;  /* META */
+        PyObject *cown;                       /* COWN, borrowed */
+        _PyCown_ipid_t ipid;                  /* IPID */
+    } value;
+} _PyRegionRefMetadata;
+
+/* Creates the node of a region that is being closed. Returns a new reference. */
+extern _PyRegionRefMetadata *_PyRegionRef_NewRegionMetaLockHeld(PyObject *region);
+
+extern void _PyRegionRef_MetaDecref(_PyRegionRefMetadata *meta);
+
+// Ownership transitions.
+extern void _PyRegionRef_MetaSetParentLockHeld(_PyRegionRefMetadata *meta,
+                                               _PyRegionRefMetadata *parent);
+/* Hands the node to `cown`, which is borrowed. The owner is from then on
+ * whoever holds the cown. */
+extern void _PyRegionRef_MetaSetCown(_PyRegionRefMetadata *meta, PyObject *cown);
+/* Stamps an explicit owner, which need not be the current interpreter and may
+ * be `_PyCown_ReleasedIpid()` to mean nobody owns the region. */
+extern void _PyRegionRef_MetaSetIpid(_PyRegionRefMetadata *meta,
+                                     _PyCown_ipid_t ipid);
+extern void _PyRegionRef_MetaRegionOpened(_PyRegionRefMetadata *meta);
+extern void _PyRegionRef_MetaResolveWip(_PyRegionRefMetadata *meta);
+
+/* Re-homes every `RegionRef` pointing at `obj` onto `region`'s node, allocating
+ * that node if this is the first reference the close has found. Every other
+ * weak reference to `obj` is cleared unless it is listed in `keep`.
+ *
+ * This is the region close hook; it replaces `_PyWeakref_ClearWeakRefsExcept()`
+ * for objects that are being closed into a region. */
+extern void _PyRegionRef_CloseWeakRefs(PyObject *obj, _Py_hashtable_t *keep,
+                                       PyObject *region);
+
+#ifdef __cplusplus
+}
+#endif
+#endif /* !Py_INTERNAL_REGIONREF_H */

Include/internal/pycore_weakref.h

@@ -13,6 +13,15 @@ extern "C" {
 #include "pycore_object.h"           // _Py_REF_IS_MERGED()
 #include "pycore_pyatomic_ft_wrappers.h"
 
+/* Guards weakrefs to immutable objects, and all region reference metadata.
+ * Declared for both builds because `_PyRegionRefMetadata` uses it either way,
+ * while the weakref lists themselves are striped in free-threaded builds. */
+extern PyMutex _PyWeakref_Lock;
+
+#define LOCK_REGION_REF_META() \
+    PyMutex_LockFlags(&_PyWeakref_Lock, _Py_LOCK_DONT_DETACH)
+#define UNLOCK_REGION_REF_META() PyMutex_Unlock(&_PyWeakref_Lock)
+
 #ifdef Py_GIL_DISABLED
 
 #define WEAKREF_LIST_LOCK(obj) \
@@ -37,9 +46,6 @@ extern "C" {
 
 #else
 
-// Lock used for weakrefs to immutable objects
-extern PyMutex _PyWeakref_Lock;
-
 #define LOCK_WEAKREFS(obj) PyMutex_LockFlags(&_PyWeakref_Lock, _Py_LOCK_DONT_DETACH)
 #define UNLOCK_WEAKREFS(obj) PyMutex_Unlock(&_PyWeakref_Lock)
 
@@ -108,7 +114,7 @@ static inline PyObject* get_ref_lock_held(PyWeakReference *ref, PyObject *obj)
 
 static inline PyObject* _PyWeakref_GET_REF(PyObject *ref_obj)
 {
-    assert(PyWeakref_Check(ref_obj));
+    assert(_PyWeakrefOrRegionRef_Check(ref_obj));
     PyWeakReference *ref = _Py_CAST(PyWeakReference*, ref_obj);
 
     PyObject *obj = _Py_atomic_load_ptr(&ref->wr_object);
@@ -125,7 +131,7 @@ static inline PyObject* _PyWeakref_GET_REF(PyObject *ref_obj)
 
 static inline int _PyWeakref_IS_DEAD(PyObject *ref_obj)
 {
-    assert(PyWeakref_Check(ref_obj));
+    assert(_PyWeakrefOrRegionRef_Check(ref_obj));
     int ret = 0;
     PyWeakReference *ref = _Py_CAST(PyWeakReference*, ref_obj);
     PyObject *obj = FT_ATOMIC_LOAD_PTR(ref->wr_object);
@@ -149,6 +155,10 @@ extern Py_ssize_t _PyWeakref_GetWeakrefCount(PyObject *obj);
 // intact.
 extern void _PyWeakref_ClearWeakRefsNoCallbacks(PyObject *obj);
 
+// Same, but every weak reference listed in `keep` survives. The table is keyed
+// by the weak reference objects themselves, not by their referents.
+extern void _PyWeakref_ClearWeakRefsExcept(PyObject *obj, _Py_hashtable_t *keep);
+
 PyAPI_FUNC(void) _PyWeakref_OnObjectFreeze(PyObject *object);
 PyAPI_FUNC(void) _PyImmutability_ClearWeakRefsWithCallback(PyObject *object, PyWeakReference **callbacks);
 PyAPI_FUNC(int) _PyWeakref_IsDead(PyObject *weakref);

Lib/immutable.py

@@ -21,6 +21,9 @@
 FREEZABLE_PROXY = _c.FREEZABLE_PROXY
 InterpreterLocal = _c.InterpreterLocal
 SharedField = _c.SharedField
+TracingRegion = _c.TracingRegion
+Cown = _c.Cown
+RegionRef = _c.RegionRef
 
 # FIXME(immutable): For the longest time we used the name `isfrozen`
 # without the underscore. This keeps the function name for now, but
@@ -139,6 +142,9 @@ def __enter__(self):
     "FREEZABLE_PROXY",
     "InterpreterLocal",
     "SharedField",
+    "TracingRegion",
+    "Cown",
+    "RegionRef",
     "freezable",
     "unfreezable",
     "explicitlyFreezable",

Lib/test/test_freeze/test_implicit.py

@@ -1,3 +1,4 @@
+import sys
 import unittest
 from immutable import freeze, is_frozen
 
@@ -139,6 +140,26 @@ def test_deeply_nested_no_stack_overflow(self):
             obj = (obj,)
         self.assertTrue(is_frozen(obj))
 
+    def test_abandoned_walk_keeps_references(self):
+        """An aborted walk must not drop references it never took.
+
+        The walk pushes objects onto a worklist without increfing them, so
+        anything still on the worklist when a mutable object aborts the walk
+        used to be decrefed when the worklist was released. That freed the
+        object while its real owners were still pointing at it, which showed
+        up much later as a negative refcount.
+        """
+        # Built at runtime so it is neither interned nor immortal, which makes
+        # its reference count fully accounted for by this test.
+        item = "".join(["abandoned", "-", "worklist", "-", "entry"])
+        # Tuples are traversed back to front, so `item` reaches the worklist
+        # before the dict aborts the walk.
+        obj = ({"mutable": 1}, item)
+
+        before = sys.getrefcount(item)
+        self.assertFalse(is_frozen(obj))
+        self.assertEqual(sys.getrefcount(item), before)
+
 
 if __name__ == '__main__':
     unittest.main()