From 010d729f0e7ce88152142606ce056c9e64e04d03 Mon Sep 17 00:00:00 2001
From: Shreyash Bhatkar <12shreyashh@gmail.com>
Date: Mon, 12 Oct 2026 04:33:32 +0000
Subject: [PATCH] Drop dead iterator refs from the stats registry

---
 grain/_src/python/dataset/stats.py      | 18 ++++++++-
 grain/_src/python/dataset/stats_test.py | 53 +++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/grain/_src/python/dataset/stats.py b/grain/_src/python/dataset/stats.py
index de6c389de..ec69082f0 100644
--- a/grain/_src/python/dataset/stats.py
+++ b/grain/_src/python/dataset/stats.py
@@ -441,7 +441,16 @@ class HashableWeakRef:
 
   def __init__(self, obj):
     self._obj_id = id(obj)
-    self._weakref = weakref.ref(obj)
+    # Drop the registry entry when the referent is collected so a later
+    # iterator that reuses this address is not treated as already present.
+    self._weakref = weakref.ref(obj, self._remove_from_registry)
+
+  def _remove_from_registry(self, _ref):
+    # discard is a no-op if a child node already removed this ref. At
+    # interpreter shutdown the module global may already have been cleared.
+    registry = _iter_weakref_registry
+    if registry is not None:
+      registry.discard(self)
 
   def __call__(self):
     return self._weakref()
@@ -452,7 +461,12 @@ def __hash__(self):
   def __eq__(self, other):
     if not isinstance(other, HashableWeakRef):
       return False
-    return self._obj_id == other._obj_id
+    if self is other:
+      return True
+    if self._obj_id != other._obj_id:
+      return False
+    # Dead and live refs that share an address are different objects.
+    return self() is not None and other() is not None
 
 
 class _Table:
diff --git a/grain/_src/python/dataset/stats_test.py b/grain/_src/python/dataset/stats_test.py
index a4f57ae99..2bce939dc 100644
--- a/grain/_src/python/dataset/stats_test.py
+++ b/grain/_src/python/dataset/stats_test.py
@@ -15,6 +15,7 @@
 import collections
 import contextlib
 import functools
+import gc
 import platform
 import re
 import sys
@@ -639,5 +640,57 @@ def test_dataset_visualization_with_output_dir_none(self):
     self.assertIsInstance(s, stats._DefaultStats)
 
 
+class IterWeakRefRegistryTest(absltest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    stats._iter_weakref_registry.clear()
+
+  def tearDown(self):
+    stats._iter_weakref_registry.clear()
+    super().tearDown()
+
+  def test_drops_collected_iterators(self):
+    # Dead refs must not stay in the registry. CPython reuses addresses, so a
+    # stale ref compares equal to a new iterator and set.add skips it.
+    ds = (
+        dataset.MapDataset.range(8)
+        .map(lambda i: np.full(2, i))
+        .batch(4)
+        .to_iter_dataset(
+            options.ReadOptions(num_threads=0, prefetch_buffer_size=0)
+        )
+    )
+    iterator = None
+    for _ in range(100):
+      if iterator is not None:
+        del iterator
+        gc.collect()
+      iterator = iter(ds)
+      self.assertNotEmpty(list(iterator))
+      self.assertTrue(
+          any(
+              ref() is iterator
+              for ref in stats._iter_weakref_registry
+              if ref is not None
+          )
+      )
+
+    gc.collect()
+    dead_refs = [
+        ref
+        for ref in stats._iter_weakref_registry
+        if isinstance(ref, stats.HashableWeakRef) and ref() is None
+    ]
+    self.assertEmpty(dead_refs)
+    self.assertTrue(
+        any(
+            ref() is iterator
+            for ref in stats._iter_weakref_registry
+            if ref is not None
+        )
+    )
+
+
 if __name__ == "__main__":
   absltest.main()