gh-126366: Make native generators thread safe

Lib/test/test_free_threading/test_generators.py

@@ -0,0 +1,23 @@
+import unittest
+import concurrent.futures
+
+from unittest import TestCase
+
+from test.support import threading_helper
+
+@threading_helper.requires_working_threading()
+class TestGen(TestCase):
+    def test_generators_basic(self):
+        def gen():
+            for _ in range(5000):
+                yield
+
+
+        it = gen()
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            for _ in range(5000):
+                executor.submit(lambda: next(it))
+
+
+if __name__ == "__main__":
+    unittest.main()

Lib/test/test_generators.py

@@ -10,6 +10,7 @@
 import types
 
 from test import support
+from test.support import threading_helper
 
 try:
     import _testcapi
@@ -268,6 +269,47 @@ def loop():
         #This should not raise
         loop()
 
+    @support.requires_resource("cpu")
+    def test_generator_thread_safety(self):
+        # GH-126369: generators were not thread safe
+        from threading import Thread, Lock
+        import contextlib
+
+        def my_generator():
+            for i in range(1000000):
+                yield i
+
+
+        lock = Lock()
+        amount = 0
+        thread_count = 10
+
+        def gen_to_list(gen):
+            # Note: it's intentional to throw out the exception here. Generators
+            # aren't thread safe, so who knows what will happen. Right now, it just spams
+            # a lot of ValueError's, but that might change if we decide to make generators
+            # thread safe in the future. We're just making sure it doesn't crash.
+            with contextlib.suppress(ValueError):
+                list(gen)
+
+            with lock:
+                nonlocal amount
+                amount += 1
+
+        generator = my_generator()
+        with threading_helper.catch_threading_exception() as cm:
+            ts = [Thread(target=gen_to_list, args=(generator,)) for _ in range(thread_count)]
+            for t in ts:
+                t.start()
+
+            for t in ts:
+                t.join()
+
+            self.assertIsNone(cm.exc_value)
+
+        self.assertEqual(amount, thread_count)
+
+
 
 class ModifyUnderlyingIterableTest(unittest.TestCase):
     iterables = [

Misc/NEWS.d/next/Core_and_Builtins/2024-11-03-15-07-15.gh-issue-126366.nd_LOu.rst

@@ -0,0 +1 @@
+Fix possible data races when using :term:`generator iterator` objects concurrently.

Objects/genobject.c

@@ -175,7 +175,7 @@ gen_dealloc(PyObject *self)
 }
 
 static PySendResult
-gen_send_ex2(PyGenObject *gen, PyObject *arg, PyObject **presult,
+gen_send_ex2_lock_held(PyGenObject *gen, PyObject *arg, PyObject **presult,
              int exc, int closing)
 {
     PyThreadState *tstate = _PyThreadState_GET();
@@ -273,6 +273,17 @@ gen_send_ex2(PyGenObject *gen, PyObject *arg, PyObject **presult,
     return result ? PYGEN_RETURN : PYGEN_ERROR;
 }
 
+static inline PySendResult
+gen_send_ex2(PyGenObject *gen, PyObject *arg, PyObject **presult,
+             int exc, int closing)
+{
+    PySendResult result;
+    Py_BEGIN_CRITICAL_SECTION(gen);
+    result = gen_send_ex2_lock_held(gen, arg, presult, exc, closing);
+    Py_END_CRITICAL_SECTION();
+    return result;
+}
+
 static PySendResult
 PyGen_am_send(PyObject *self, PyObject *arg, PyObject **result)
 {

Objects/listobject.c

@@ -1362,9 +1362,9 @@ _list_extend(PyListObject *self, PyObject *iterable)
         Py_END_CRITICAL_SECTION2();
     }
     else {
-        Py_BEGIN_CRITICAL_SECTION(self);
+        Py_BEGIN_CRITICAL_SECTION2(self, iterable);
         res = list_extend_iter_lock_held(self, iterable);
-        Py_END_CRITICAL_SECTION();
+        Py_END_CRITICAL_SECTION2();
     }
     return res;
 }

Python/bytecodes.c

@@ -1104,8 +1104,10 @@ dummy_func(
                 (Py_TYPE(receiver_o) == &PyGen_Type || Py_TYPE(receiver_o) == &PyCoro_Type) &&
                 ((PyGenObject *)receiver_o)->gi_frame_state < FRAME_EXECUTING)
             {
+                _PyInterpreterFrame *gen_frame;
                 PyGenObject *gen = (PyGenObject *)receiver_o;
-                _PyInterpreterFrame *gen_frame = &gen->gi_iframe;
+                Py_BEGIN_CRITICAL_SECTION(gen);
+                gen_frame = &gen->gi_iframe;
                 STACK_SHRINK(1);
                 _PyFrame_StackPush(gen_frame, v);
                 gen->gi_frame_state = FRAME_EXECUTING;
@@ -1115,6 +1117,7 @@ dummy_func(
                 frame->return_offset = (uint16_t)(INSTRUCTION_SIZE + oparg);
                 assert(gen_frame->previous == NULL);
                 gen_frame->previous = frame;
+                Py_END_CRITICAL_SECTION();
                 DISPATCH_INLINED(gen_frame);
             }
             if (PyStackRef_Is(v, PyStackRef_None) && PyIter_Check(receiver_o)) {