gh-115999: Make list and tuple iteration more thread-safe.

Include/internal/pycore_range.h

@@ -13,6 +13,12 @@ typedef struct {
     long start;
     long step;
     long len;
+#ifdef Py_GIL_DISABLED
+    // Make sure multi-threaded use of a single iterator doesn't produce
+    // values past the end of the range (even though it is NOT guaranteed to
+    // uniquely produce all the values in the range!)
+    long stop;
+#endif
 } _PyRangeIterObject;
 
 #ifdef __cplusplus

Lib/test/libregrtest/tsan.py

@@ -26,6 +26,7 @@
     'test_threadsignals',
     'test_weakref',
     'test_free_threading.test_slots',
+    'test_free_threading.test_iteration',
 ]
 
 

Lib/test/test_free_threading/test_iteration.py

@@ -0,0 +1,123 @@
+import sys
+import threading
+import unittest
+from test import support
+
+# The race conditions these tests were written for only happen every now and
+# then, even with the current numbers. To find rare race conditions, bumping
+# these up will help, but it makes the test runtime highly variable under
+# free-threading. Overhead is much higher under ThreadSanitizer, but it's
+# also much better at detecting certain races, so we don't need as many
+# items/threads.
+if support.check_sanitizer(thread=True):
+    NUMITEMS = 1000
+    NUMTHREADS = 2
+else:
+    NUMITEMS = 50000
+    NUMTHREADS = 3
+NUMMUTATORS = 2
+
+class ContendedTupleIterationTest(unittest.TestCase):
+    def make_testdata(self, n):
+        return tuple(range(n))
+
+    def assert_iterator_results(self, results, expected):
+        # Most iterators are not atomic (yet?) so they can skip or duplicate
+        # items, but they should not invent new items (like the range
+        # iterator has done in the past).
+        extra_items = set(results) - set(expected)
+        self.assertEqual(set(), extra_items)
+
+    def run_threads(self, func, *args, numthreads=NUMTHREADS):
+        threads = []
+        for _ in range(numthreads):
+            t = threading.Thread(target=func, args=args)
+            t.start()
+            threads.append(t)
+        return threads
+
+    def test_iteration(self):
+        """Test iteration over a shared container"""
+        seq = self.make_testdata(NUMITEMS)
+        results = []
+        start = threading.Event()
+        def worker():
+            idx = 0
+            start.wait()
+            for item in seq:
+                idx += 1
+            results.append(idx)
+        threads = self.run_threads(worker)
+        start.set()
+        for t in threads:
+            t.join()
+        # Each thread has its own iterator, so results should be entirely predictable.
+        self.assertEqual(results, [NUMITEMS] * NUMTHREADS)
+
+    def test_shared_iterator(self):
+        """Test iteration over a shared iterator"""
+        seq = self.make_testdata(NUMITEMS)
+        it = iter(seq)
+        results = []
+        start = threading.Event()
+        def worker():
+            items = []
+            start.wait()
+            # We want a tight loop, so put items in the shared list at the end.
+            for item in it:
+                items.append(item)
+            results.extend(items)
+        threads = self.run_threads(worker)
+        start.set()
+        for t in threads:
+            t.join()
+        self.assert_iterator_results(sorted(results), seq)
+
+class ContendedListIterationTest(ContendedTupleIterationTest):
+    def make_testdata(self, n):
+        return list(range(n))
+
+    def test_iteration_while_mutating(self):
+        """Test iteration over a shared mutating container."""
+        seq = self.make_testdata(NUMITEMS)
+        results = []
+        start = threading.Event()
+        endmutate = threading.Event()
+        def mutator():
+            orig = seq[:]
+            # Make changes big enough to cause resizing of the list, with
+            # items shifted around for good measure.
+            replacement = (orig * 3)[NUMITEMS//2:]
+            start.wait()
+            while not endmutate.is_set():
+                seq[:] = replacement
+                seq[:] = orig
+        def worker():
+            items = []
+            start.wait()
+            # We want a tight loop, so put items in the shared list at the end.
+            for item in seq:
+                items.append(item)
+            results.extend(items)
+        mutators = ()
+        try:
+            threads = self.run_threads(worker)
+            mutators = self.run_threads(mutator, numthreads=NUMMUTATORS)
+            start.set()
+            for t in threads:
+                t.join()
+        finally:
+            endmutate.set()
+            for m in mutators:
+                m.join()
+        self.assert_iterator_results(results, list(seq))
+
+
+class ContendedRangeIterationTest(ContendedTupleIterationTest):
+    def make_testdata(self, n):
+        return range(n)
+
+
+class ContendedLongRangeIterationTest(ContendedTupleIterationTest):
+    def make_testdata(self, n):
+        return range(0, sys.maxsize*n, sys.maxsize)

Objects/listobject.c

@@ -344,7 +344,7 @@ list_item_impl(PyListObject *self, Py_ssize_t idx)
 static inline PyObject*
 list_get_item_ref(PyListObject *op, Py_ssize_t i)
 {
-    if (!_Py_IsOwnedByCurrentThread((PyObject *)op) && !_PyObject_GC_IS_SHARED(op)) {
+    if (!_Py_IsOwnedByCurrentThread((PyObject *)op)) {
         return list_item_impl(op, i);
     }
     // Need atomic operation for the getting size.

Objects/rangeobject.c

@@ -7,6 +7,7 @@
 #include "pycore_modsupport.h"    // _PyArg_NoKwnames()
 #include "pycore_range.h"
 #include "pycore_tuple.h"         // _PyTuple_ITEMS()
+#include "pycore_pyatomic_ft_wrappers.h"
 
 
 /* Support objects whose length is > PY_SSIZE_T_MAX.
@@ -816,10 +817,19 @@
 static PyObject *
 rangeiter_next(_PyRangeIterObject *r)
 {
-    if (r->len > 0) {
-        long result = r->start;
-        r->start = result + r->step;
-        r->len--;
+    long len = FT_ATOMIC_LOAD_LONG_RELAXED(r->len);
+    if (len > 0) {
+        long result = FT_ATOMIC_LOAD_LONG_RELAXED(r->start);
+        FT_ATOMIC_STORE_LONG_RELAXED(r->start, result + r->step);
+        FT_ATOMIC_STORE_LONG_RELAXED(r->len, len - 1);
+#ifdef Py_GIL_DISABLED
+        // Concurrent calls can cause start to be updated by another thread
+        // after our len check, so double-check that we're not past the end.
+        if ((r->step > 0 && result >= r->stop) ||
+            (r->step < 0 && result <= r->stop)) {
+            return NULL;
+        }
+#endif
         return PyLong_FromLong(result);
     }
     return NULL;
@@ -828,7 +838,7 @@
 static PyObject *
 rangeiter_len(_PyRangeIterObject *r, PyObject *Py_UNUSED(ignored))
 {
-    return PyLong_FromLong(r->len);
+    return PyLong_FromLong(FT_ATOMIC_LOAD_LONG_RELAXED(r->len));
 }
 
 PyDoc_STRVAR(length_hint_doc,
@@ -841,10 +851,14 @@
     PyObject *range;
 
     /* create a range object for pickling */
-    start = PyLong_FromLong(r->start);
+    start = PyLong_FromLong(FT_ATOMIC_LOAD_LONG_RELAXED(r->start));
     if (start == NULL)
         goto err;
+#ifdef Py_GIL_DISABLED
+    stop = PyLong_FromLong(r->stop);
+#else
     stop = PyLong_FromLong(r->start + r->len * r->step);
+#endif
     if (stop == NULL)
         goto err;
     step = PyLong_FromLong(r->step);
@@ -870,13 +884,15 @@
     long index = PyLong_AsLong(state);
     if (index == -1 && PyErr_Occurred())
         return NULL;
+    long len = FT_ATOMIC_LOAD_LONG_RELAXED(r->len);
     /* silently clip the index value */
     if (index < 0)
         index = 0;
-    else if (index > r->len)
-        index = r->len; /* exhausted iterator */
-    r->start += index * r->step;
-    r->len -= index;
+    else if (index > len)
+        index = len; /* exhausted iterator */
+    FT_ATOMIC_STORE_LONG_RELAXED(r->start,
+        FT_ATOMIC_LOAD_LONG_RELAXED(r->start) + index * r->step);
+    FT_ATOMIC_STORE_LONG_RELAXED(r->len, len - index);
     Py_RETURN_NONE;
 }
 
@@ -966,6 +982,9 @@
     it->start = start;
     it->step = step;
     it->len = len;
+#ifdef Py_GIL_DISABLED
+    it->stop = stop;
+#endif
     return (PyObject *)it;
 }
 
@@ -979,75 +998,87 @@
 static PyObject *
 longrangeiter_len(longrangeiterobject *r, PyObject *no_args)
 {
-    Py_INCREF(r->len);
-    return r->len;
+    PyObject *len;
+    Py_BEGIN_CRITICAL_SECTION(r);
+    len = Py_NewRef(r->len);
+    Py_END_CRITICAL_SECTION();
+    return len;
 }
 
 static PyObject *
 longrangeiter_reduce(longrangeiterobject *r, PyObject *Py_UNUSED(ignored))
 {
     PyObject *product, *stop=NULL;
-    PyObject *range;
+    PyObject *range, *result=NULL;
 
+    Py_BEGIN_CRITICAL_SECTION(r);
     /* create a range object for pickling.  Must calculate the "stop" value */
     product = PyNumber_Multiply(r->len, r->step);
     if (product == NULL)
-        return NULL;
+        goto fail;
     stop = PyNumber_Add(r->start, product);
     Py_DECREF(product);
     if (stop ==  NULL)
-        return NULL;
+        goto fail;
     range =  (PyObject*)make_range_object(&PyRange_Type,
                                Py_NewRef(r->start), stop, Py_NewRef(r->step));
     if (range == NULL) {
         Py_DECREF(r->start);
         Py_DECREF(stop);
         Py_DECREF(r->step);
-        return NULL;
+        goto fail;
     }
 
     /* return the result */
-    return Py_BuildValue("N(N)O", _PyEval_GetBuiltin(&_Py_ID(iter)),
-                         range, Py_None);
+    result = Py_BuildValue("N(N)O", _PyEval_GetBuiltin(&_Py_ID(iter)),
+                           range, Py_None);
+fail:
+    Py_END_CRITICAL_SECTION();
+    return result;
 }
 
 static PyObject *
 longrangeiter_setstate(longrangeiterobject *r, PyObject *state)
 {
     PyObject *zero = _PyLong_GetZero();  // borrowed reference
     int cmp;
+    PyObject *result = NULL;
 
+    Py_BEGIN_CRITICAL_SECTION(r);
     /* clip the value */
     cmp = PyObject_RichCompareBool(state, zero, Py_LT);
     if (cmp < 0)
-        return NULL;
+        goto fail;
     if (cmp > 0) {
         state = zero;
     }
     else {
         cmp = PyObject_RichCompareBool(r->len, state, Py_LT);
         if (cmp < 0)
-            return NULL;
+            goto fail;
         if (cmp > 0)
             state = r->len;
     }
     PyObject *product = PyNumber_Multiply(state, r->step);
     if (product == NULL)
-        return NULL;
+        goto fail;
     PyObject *new_start = PyNumber_Add(r->start, product);
     Py_DECREF(product);
     if (new_start == NULL)
-        return NULL;
+        goto fail;
     PyObject *new_len = PyNumber_Subtract(r->len, state);
     if (new_len == NULL) {
         Py_DECREF(new_start);
-        return NULL;
+        goto fail;
     }
     PyObject *tmp = r->start;
     r->start = new_start;
     Py_SETREF(r->len, new_len);
     Py_DECREF(tmp);
-    Py_RETURN_NONE;
+    result = Py_NewRef(Py_None);
+fail:
+    Py_END_CRITICAL_SECTION();
+    return result;
 }
 
 static PyMethodDef longrangeiter_methods[] = {
@@ -1072,21 +1103,25 @@
 static PyObject *
 longrangeiter_next(longrangeiterobject *r)
 {
+    PyObject *result = NULL;
+    Py_BEGIN_CRITICAL_SECTION(r);
     if (PyObject_RichCompareBool(r->len, _PyLong_GetZero(), Py_GT) != 1)
-        return NULL;
+        goto fail;
 
     PyObject *new_start = PyNumber_Add(r->start, r->step);
     if (new_start == NULL) {
-        return NULL;
+        goto fail;
     }
     PyObject *new_len = PyNumber_Subtract(r->len, _PyLong_GetOne());
     if (new_len == NULL) {
         Py_DECREF(new_start);
-        return NULL;
+        goto fail;
     }
-    PyObject *result = r->start;
+    result = r->start;
     r->start = new_start;
     Py_SETREF(r->len, new_len);
+fail:
+    Py_END_CRITICAL_SECTION();
     return result;
 }
 

Objects/tupleobject.c

@@ -1017,14 +1017,17 @@ tupleiter_next(PyObject *obj)
         return NULL;
     assert(PyTuple_Check(seq));
 
-    if (it->it_index < PyTuple_GET_SIZE(seq)) {
-        item = PyTuple_GET_ITEM(seq, it->it_index);
-        ++it->it_index;
+    Py_ssize_t index = FT_ATOMIC_LOAD_SSIZE_RELAXED(it->it_index);
+    if (index < PyTuple_GET_SIZE(seq)) {
+        item = PyTuple_GET_ITEM(seq, index);
+        FT_ATOMIC_STORE_SSIZE_RELAXED(it->it_index, index + 1);
         return Py_NewRef(item);
     }
 
+#ifndef Py_GIL_DISABLED
     it->it_seq = NULL;
     Py_DECREF(seq);
+#endif
     return NULL;
 }