Make permutations thread-safe

Author: eendebakpt

Lib/test/test_free_threading/test_itertools.py

@@ -1,6 +1,6 @@
 import unittest
 from threading import Thread, Barrier
-from itertools import batched, chain, combinations_with_replacement, cycle
+from itertools import batched, chain, combinations_with_replacement, cycle, permutations
 from test.support import threading_helper
 
 
@@ -92,7 +92,7 @@ def work(it):
     @threading_helper.reap_threads
     def test_combinations_with_replacement(self):
         number_of_threads = 6
-        number_of_iterations = 50
+        number_of_iterations = 36
         data = tuple(range(2))
 
         barrier = Barrier(number_of_threads)
@@ -116,6 +116,33 @@ def work(it):
 
             barrier.reset()
 
+    @threading_helper.reap_threads
+    def test_permutations(self):
+        number_of_threads = 6
+        number_of_iterations = 36
+        data = tuple(range(4))
+
+        barrier = Barrier(number_of_threads)
+        def work(it):
+            barrier.wait()
+            while True:
+                try:
+                    next(it)
+                except StopIteration:
+                    break
+
+        for _ in range(number_of_iterations):
+            perm_iterator = permutations(data, 2)
+            worker_threads = []
+            for _ in range(number_of_threads):
+                worker_threads.append(
+                    Thread(target=work, args=[perm_iterator]))
+
+            with threading_helper.start_threads(worker_threads):
+                pass
+
+            barrier.reset()
+
 
 if __name__ == "__main__":
     unittest.main()

Modules/itertoolsmodule.c

@@ -2856,7 +2856,7 @@ permutations_traverse(PyObject *op, visitproc visit, void *arg)
 }
 
 static PyObject *
-permutations_next(PyObject *op)
+permutations_next_lock_held(PyObject *op)
 {
     permutationsobject *po = permutationsobject_CAST(op);
     PyObject *elem;
@@ -2946,6 +2946,16 @@ permutations_next(PyObject *op)
     return NULL;
 }
 
+static PyObject *
+permutations_next(PyObject *op)
+{
+    PyObject *result;
+    Py_BEGIN_CRITICAL_SECTION(op);
+    result = permutations_next_lock_held(op);
+    Py_END_CRITICAL_SECTION()
+    return result;
+}
+
 static PyMethodDef permuations_methods[] = {
     {"__sizeof__", permutations_sizeof, METH_NOARGS, sizeof_doc},
     {NULL,              NULL}   /* sentinel */