gh-154130: Fix double decref of the dict in shared dict iterators

Lib/test/test_free_threading/test_dict.py

@@ -493,5 +493,50 @@ def race(iter_times: int,
         CHURN_SIZE = 512
         race(ITER_TIMES, CHURN_SIZE, APPENDER_THREADS, REPLACER_THREADS)
 
+    def test_racing_shared_iterator_exhaustion(self):
+        # gh-154130: exhausting a dict iterator shared between threads must
+        # not release the iterator's reference to the dict more than once.
+        makers = [
+            lambda d: iter(d),
+            lambda d: iter(d.values()),
+            lambda d: iter(d.items()),
+            lambda d: reversed(d),
+            lambda d: reversed(d.items()),
+        ]
+        for make in makers:
+            cell = [make({i: i for i in range(8)})]
+
+            def worker():
+                for _ in range(5000):
+                    it = cell[0]
+                    try:
+                        next(it)
+                    except StopIteration:
+                        cell[0] = make({i: i for i in range(8)})
+                    except RuntimeError:
+                        # Concurrent next() calls may observe each other's
+                        # updates as a modification; only crashes are bugs.
+                        pass
+
+            with self.subTest(make=make):
+                threading_helper.run_concurrently(worker, nthreads=4)
+
+    def test_exhausted_iterator_stays_exhausted(self):
+        d = {1: 1}
+        for it in (iter(d), iter(d.items()), reversed(d)):
+            self.assertEqual(len(list(it)), 1)
+            d[2] = 2
+            self.assertRaises(StopIteration, next, it)
+            self.assertEqual(it.__length_hint__(), 0)
+            del d[2]
+
+    def test_reversed_empty_dict_iterator_not_exhausted(self):
+        # A reverse iterator over an empty dict starts at position -1 but is
+        # not exhausted, so a later size change must still be detected.
+        d = {}
+        it = reversed(d)
+        d[1] = 1
+        self.assertRaises(RuntimeError, next, it)
+
 if __name__ == "__main__":
     unittest.main()

Objects/dictobject.c

@@ -5705,6 +5705,28 @@ static PyMethodDef dictiter_methods[] = {
     {NULL,              NULL}           /* sentinel */
 };
 
+/* Mark the iterator as exhausted.  In the free-threaded build an iterator
+   may be advanced by several threads at once, each holding a borrowed
+   pointer to di_dict, so the iterator keeps its reference to the dict until
+   it is deallocated (as list and tuple iterators do) and records exhaustion
+   by setting di_pos to DICTITER_EXHAUSTED.  This is distinct from -1, the
+   starting position of a reverse iterator over an empty dict. */
+#ifdef Py_GIL_DISABLED
+#define DICTITER_EXHAUSTED PY_SSIZE_T_MIN
+#endif
+
+static void
+dictiter_set_exhausted(dictiterobject *di, PyDictObject *d)
+{
+#ifdef Py_GIL_DISABLED
+    _Py_atomic_store_ssize_relaxed(&di->di_pos, DICTITER_EXHAUSTED);
+    _Py_atomic_store_ssize_relaxed(&di->len, 0);
+#else
+    di->di_dict = NULL;
+    Py_DECREF(d);
+#endif
+}
+
 #ifdef Py_GIL_DISABLED
 
 static int
@@ -5970,16 +5992,18 @@ dictiter_iternextitem_lock_held(PyDictObject *d, PyObject *self,
     assert (PyAnyDict_Check(d));
     ASSERT_DICT_LOCKED(d);
 
+    i = FT_ATOMIC_LOAD_SSIZE_RELAXED(di->di_pos);
+    if (i < 0) {
+        return -1;  /* exhausted */
+    }
+
     if (di->di_used != d->ma_used) {
         PyErr_SetString(PyExc_RuntimeError,
                         "dictionary changed size during iteration");
         di->di_used = -1; /* Make this state sticky */
         return -1;
     }
 
-    i = FT_ATOMIC_LOAD_SSIZE_RELAXED(di->di_pos);
-
-    assert(i >= 0);
     if (_PyDict_HasSplitTable(d)) {
         if (i >= d->ma_used)
             goto fail;
@@ -6030,8 +6054,7 @@ dictiter_iternextitem_lock_held(PyDictObject *d, PyObject *self,
     return 0;
 
 fail:
-    di->di_dict = NULL;
-    Py_DECREF(d);
+    dictiter_set_exhausted(di, d);
     return -1;
 }
 
@@ -6075,6 +6098,11 @@ dictiter_iternext_threadsafe(PyDictObject *d, PyObject *self,
 
     assert (PyAnyDict_Check(d));
 
+    i = _Py_atomic_load_ssize_relaxed(&di->di_pos);
+    if (i == DICTITER_EXHAUSTED) {
+        return -1;
+    }
+
     if (di->di_used != _Py_atomic_load_ssize_relaxed(&d->ma_used)) {
         PyErr_SetString(PyExc_RuntimeError,
                         "dictionary changed size during iteration");
@@ -6084,9 +6112,7 @@ dictiter_iternext_threadsafe(PyDictObject *d, PyObject *self,
 
     ensure_shared_on_read(d);
 
-    i = _Py_atomic_load_ssize_relaxed(&di->di_pos);
     k = _Py_atomic_load_ptr_acquire(&d->ma_keys);
-    assert(i >= 0);
     if (_PyDict_HasSplitTable(d)) {
         PyDictValues *values = _Py_atomic_load_ptr_consume(&d->ma_values);
         if (values == NULL) {
@@ -6165,8 +6191,7 @@ dictiter_iternext_threadsafe(PyDictObject *d, PyObject *self,
                     "dictionary keys changed during iteration");
 
 fail:
-    di->di_dict = NULL;
-    Py_DECREF(d);
+    dictiter_set_exhausted(di, d);
     return -1;
 
 try_locked:
@@ -6268,6 +6293,12 @@ dictreviter_iter_lock_held(PyDictObject *d, PyObject *self)
     assert (PyAnyDict_Check(d));
     ASSERT_DICT_LOCKED(d);
 
+#ifdef Py_GIL_DISABLED
+    if (di->di_pos == DICTITER_EXHAUSTED) {
+        return NULL;
+    }
+#endif
+
     if (di->di_used != d->ma_used) {
         PyErr_SetString(PyExc_RuntimeError,
                          "dictionary changed size during iteration");
@@ -6354,8 +6385,7 @@ dictreviter_iter_lock_held(PyDictObject *d, PyObject *self)
     }
 
 fail:
-    di->di_dict = NULL;
-    Py_DECREF(d);
+    dictiter_set_exhausted(di, d);
     return NULL;
 }