gh-155515: Use GC tracking for HAMT iterators (#155517)

nascheme · web-flow · commit f0562124927e · 2026-08-13T17:45:02.000+05:30
diff --git a/Lib/test/test_context.py b/Lib/test/test_context.py
@@ -1282,6 +1282,31 @@ def test_hamt_gc_2(self):
 
         self.assertIsNone(ref())
 
+    def test_hamt_gc_3(self):
+        # gh-154535: the iterators must be tracked by the GC, otherwise a
+        # cycle running through one is never collected and the HAMT it
+        # holds -- and everything in it -- leaks.
+        A = HashKey(100, 'A')
+
+        container = []
+        h = hamt()
+        h = h.set(A, container)
+
+        hi = h.items()
+        self.assertTrue(gc.is_tracked(hi))
+
+        # Close the cycle: hi -> h -> container -> hi.
+        container.append(hi)
+        ref = weakref.ref(h)
+
+        del h, hi, container
+
+        gc.collect()
+        gc.collect()
+        gc.collect()
+
+        self.assertIsNone(ref())
+
     def test_hamt_in_1(self):
         A = HashKey(100, 'A')
         AA = HashKey(100, 'A')
diff --git a/Misc/NEWS.d/next/Core_and_Builtins/2026-08-10-11-20-00.gh-issue-155515.Lp3qWc.rst b/Misc/NEWS.d/next/Core_and_Builtins/2026-08-10-11-20-00.gh-issue-155515.Lp3qWc.rst
@@ -0,0 +1,4 @@
+Track the internal HAMT iterators, which back iteration over a
+:class:`contextvars.Context`, with the garbage collector.  A reference cycle
+running through such an iterator was never collected, leaking the whole
+context it iterated over.
diff --git a/Python/hamt.c b/Python/hamt.c
@@ -2451,6 +2451,10 @@ hamt_baseiter_tp_clear(PyObject *op)
 {
     PyHamtIterator *it = (PyHamtIterator*)op;
     Py_CLEAR(it->hi_obj);
+    /* i_nodes holds borrowed pointers into the tree that hi_obj was keeping
+       alive, so the cursor must not be used again.  A negative i_level makes
+       hamt_iterator_next() report I_END without touching i_nodes. */
+    it->hi_iter.i_level = -1;
     return 0;
 }
 
@@ -2497,6 +2501,10 @@ static Py_ssize_t
 hamt_baseiter_tp_len(PyObject *op)
 {
     PyHamtIterator *it = (PyHamtIterator*)op;
+    if (it->hi_obj == NULL) {
+        /* tp_clear() ran on this iterator. */
+        return 0;
+    }
     return it->hi_obj->h_count;
 }
 
@@ -2517,6 +2525,7 @@ hamt_baseiter_new(PyTypeObject *type, binaryfunc yield, PyHamtObject *o)
 
     hamt_iterator_init(&it->hi_iter, o->h_root);
 
+    PyObject_GC_Track(it);
     return (PyObject*)it;
 }
 

Original file line number	Diff line number	Diff line change
`@@ -2451,6 +2451,10 @@ hamt_baseiter_tp_clear(PyObject *op)`
`2451`	`2451`	`{`
`2452`	`2452`	`PyHamtIterator it = (PyHamtIterator)op;`
`2453`	`2453`	`Py_CLEAR(it->hi_obj);`
	`2454`	`+ /* i_nodes holds borrowed pointers into the tree that hi_obj was keeping`
	`2455`	`+ alive, so the cursor must not be used again. A negative i_level makes`
	`2456`	`+ hamt_iterator_next() report I_END without touching i_nodes. */`
	`2457`	`+ it->hi_iter.i_level = -1;`
`2454`	`2458`	`return 0;`
`2455`	`2459`	`}`
`2456`	`2460`
`@@ -2497,6 +2501,10 @@ static Py_ssize_t`
`2497`	`2501`	`hamt_baseiter_tp_len(PyObject *op)`
`2498`	`2502`	`{`
`2499`	`2503`	`PyHamtIterator it = (PyHamtIterator)op;`
	`2504`	`+ if (it->hi_obj == NULL) {`
	`2505`	`+ /* tp_clear() ran on this iterator. */`
	`2506`	`+ return 0;`
	`2507`	`+ }`
`2500`	`2508`	`return it->hi_obj->h_count;`
`2501`	`2509`	`}`
`2502`	`2510`
`@@ -2517,6 +2525,7 @@ hamt_baseiter_new(PyTypeObject type, binaryfunc yield, PyHamtObject o)`
`2517`	`2525`
`2518`	`2526`	`hamt_iterator_init(&it->hi_iter, o->h_root);`
`2519`	`2527`
	`2528`	`+ PyObject_GC_Track(it);`
`2520`	`2529`	`return (PyObject*)it;`
`2521`	`2530`	`}`
`2522`	`2531`