Add iter_locked

Author: eendebakpt

Doc/library/threading.rst

@@ -144,6 +144,20 @@ This module defines the following functions:
    of the result, even when terminated.
 
 
+.. function:: locked_iter(iterable)
+
+   Make an iterator thread-safe.
+
+   Roughly equivalent to::
+
+        class locked_iter(Iterator):
+            def __init__(self, it):
+                self._it = iter(it)
+                self._lock = Lock()
+            def __next__(self):
+                with self._lock:
+                    return next(self._it)
+
 .. function:: main_thread()
 
    Return the main :class:`Thread` object.  In normal conditions, the

Lib/test/test_free_threading/test_threading_iter_locked.py

@@ -0,0 +1,81 @@
+import unittest
+from threading import Thread, Barrier, iter_locked
+from test.support import threading_helper
+
+
+threading_helper.requires_working_threading(module=True)
+
+class non_atomic_iterator:
+
+    def __init__(self, it):
+        self.it = iter(it)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        a = next(self.it)
+        b = next(self.it)
+        return a, b
+
+def count():
+    i = 0
+    while True:
+        i += 1
+        yield i
+
+class iter_lockedThreading(unittest.TestCase):
+
+    @threading_helper.reap_threads
+    def test_iter_locked(self):
+        number_of_threads = 10
+        number_of_iterations = 10
+        barrier = Barrier(number_of_threads)
+        def work(it):
+            while True:
+                try:
+                    a, b = next(it)
+                    assert a + 1 == b
+                except StopIteration:
+                    break
+
+        data = tuple(range(400))
+        for it in range(number_of_iterations):
+            iter_locked_iterator = iter_locked(non_atomic_iterator(data,))
+            worker_threads = []
+            for ii in range(number_of_threads):
+                worker_threads.append(
+                    Thread(target=work, args=[iter_locked_iterator]))
+
+            with threading_helper.start_threads(worker_threads):
+                pass
+
+            barrier.reset()
+
+    @threading_helper.reap_threads
+    def test_iter_locked_generator(self):
+        number_of_threads = 5
+        number_of_iterations = 4
+        barrier = Barrier(number_of_threads)
+        def work(it):
+            barrier.wait()
+            for _ in range(1_000):
+                try:
+                    next(it)
+                except StopIteration:
+                    break
+
+        for it in range(number_of_iterations):
+            generator = iter_locked(count())
+            worker_threads = []
+            for ii in range(number_of_threads):
+                worker_threads.append(
+                    Thread(target=work, args=[generator]))
+
+            with threading_helper.start_threads(worker_threads):
+                pass
+
+            barrier.reset()
+
+if __name__ == "__main__":
+    unittest.main()

Lib/test/test_threading.py

@@ -2416,6 +2416,20 @@ def run_last():
         self.assertIn("RuntimeError: can't register atexit after shutdown",
                 err.decode())
 
+class LockedIterTests(unittest.TestCase):
+
+     def test_locked_iter(self):
+            for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
+                for g in (G, I, Ig, S, L, R):
+                    seq = list(g(s))
+                    expected = seq
+                    actual = list(serialize(g(s)))
+                    self.assertEqual(actual, expected)
+                self.assertRaises(TypeError, serialize, X(s))
+                self.assertRaises(TypeError, serialize, N(s))
+                self.assertRaises(ZeroDivisionError, list, serialize(E(s)))
+            for arg in [1, True, sys]:
+                self.assertRaises(TypeError, serialize, arg)
 
 if __name__ == "__main__":
     unittest.main()

Modules/_threadmodule.c

@@ -19,7 +19,6 @@
 #  include <signal.h>             // SIGINT
 #endif
 
-#include "clinic/_threadmodule.c.h"
 
 // ThreadError is just an alias to PyExc_RuntimeError
 #define ThreadError PyExc_RuntimeError
@@ -30,6 +29,7 @@ static struct PyModuleDef thread_module;
 // Module state
 typedef struct {
     PyTypeObject *excepthook_type;
+    PyTypeObject *iter_locked_type;
     PyTypeObject *lock_type;
     PyTypeObject *local_type;
     PyTypeObject *local_dummy_type;
@@ -48,6 +48,17 @@ get_thread_state(PyObject *module)
     return (thread_module_state *)state;
 }
 
+static inline thread_module_state *
+find_state_by_type(PyTypeObject *tp)
+{
+    PyObject *mod = PyType_GetModuleByDef(tp, &thread_module);
+    assert(mod != NULL);
+    return get_thread_state(mod);
+}
+
+#define clinic_state() (find_state_by_type(type))
+#include "clinic/_threadmodule.c.h"
+#undef clinic_state
 
 #ifdef MS_WINDOWS
 typedef HRESULT (WINAPI *PF_GET_THREAD_DESCRIPTION)(HANDLE, PCWSTR*);
@@ -59,8 +70,10 @@ static PF_SET_THREAD_DESCRIPTION pSetThreadDescription = NULL;
 
 /*[clinic input]
 module _thread
+class _thread.iter_locked "iter_locked_object *" "clinic_state()->iter_locked_type"
+
 [clinic start generated code]*/
-/*[clinic end generated code: output=da39a3ee5e6b4b0d input=be8dbe5cc4b16df7]*/
+/*[clinic end generated code: output=da39a3ee5e6b4b0d input=6c78d729dec7bf7e]*/
 
 
 // _ThreadHandle type
@@ -731,6 +744,99 @@ static PyType_Spec ThreadHandle_Type_spec = {
     ThreadHandle_Type_slots,
 };
 
+/* iter_locked object **************************************************************/
+
+typedef struct {
+    PyObject_HEAD
+    PyObject *it;
+} iter_locked_object;
+
+#define iter_locked_object_CAST(op)    ((iter_locked_object *)(op))
+
+/*[clinic input]
+@classmethod
+_thread.iter_locked.__new__
+    iterable: object
+    /
+Make an iterator thread-safe.
+[clinic start generated code]*/
+
+static PyObject *
+_thread_iter_locked_impl(PyTypeObject *type, PyObject *iterable)
+/*[clinic end generated code: output=4a8ad5a25f7c09ba input=ae6124177726e809]*/
+{
+    /* Get iterator. */
+    PyObject *it = PyObject_GetIter(iterable);
+    if (it == NULL)
+        return NULL;
+
+    iter_locked_object *lz = (iter_locked_object *)type->tp_alloc(type, 0);
+    lz->it = it;
+
+    return (PyObject *)lz;
+}
+
+static void
+iter_locked_dealloc(PyObject *op)
+{
+    iter_locked_object *lz = iter_locked_object_CAST(op);
+    PyTypeObject *tp = Py_TYPE(lz);
+    PyObject_GC_UnTrack(lz);
+    Py_XDECREF(lz->it);
+    tp->tp_free(lz);
+    Py_DECREF(tp);
+}
+
+static int
+iter_locked_traverse(PyObject *op, visitproc visit, void *arg)
+{
+    iter_locked_object *lz = iter_locked_object_CAST(op);
+    Py_VISIT(Py_TYPE(lz));
+    Py_VISIT(lz->it);
+    return 0;
+}
+
+static PyObject *
+iter_locked_next(PyObject *op)
+{
+    iter_locked_object *lz = iter_locked_object_CAST(op);
+    PyObject *result = NULL;
+
+    Py_BEGIN_CRITICAL_SECTION(lz->it);  // or lock on op?
+    PyObject *it = lz->it;
+    if (it != NULL) {
+        result = PyIter_Next(lz->it);
+        if (result == NULL) {
+            /* Note:  StopIteration is already cleared by PyIter_Next() */
+            if (PyErr_Occurred())
+                return NULL;
+            Py_CLEAR(lz->it);
+        }
+    }
+    Py_END_CRITICAL_SECTION();
+    return result;
+}
+
+static PyType_Slot iter_locked_slots[] = {
+    {Py_tp_dealloc, iter_locked_dealloc},
+    {Py_tp_getattro, PyObject_GenericGetAttr},
+    {Py_tp_doc, (void *)_thread_iter_locked__doc__},
+    {Py_tp_traverse, iter_locked_traverse},
+    {Py_tp_iter, PyObject_SelfIter},
+    {Py_tp_iternext, iter_locked_next},
+    {Py_tp_new, _thread_iter_locked},
+    {Py_tp_free, PyObject_GC_Del},
+    {0, NULL},
+};
+
+static PyType_Spec iter_locked_spec = {
+    .name = "threading.iter_locked",
+    .basicsize = sizeof(iter_locked_object),
+    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
+              Py_TPFLAGS_IMMUTABLETYPE),
+    .slots = iter_locked_slots,
+};
+
 /* Lock objects */
 
 typedef struct {
@@ -2631,6 +2737,15 @@ thread_module_exec(PyObject *module)
         return -1;
     }
 
+    // iter_locked
+    state->iter_locked_type = (PyTypeObject *)PyType_FromModuleAndSpec(module, &iter_locked_spec, NULL);
+    if (state->iter_locked_type == NULL) {
+        return -1;
+    }
+    if (PyModule_AddType(module, state->iter_locked_type) < 0) {
+        return -1;
+    }
+
     // Lock
     state->lock_type = (PyTypeObject *)PyType_FromModuleAndSpec(module, &lock_type_spec, NULL);
     if (state->lock_type == NULL) {
@@ -2739,6 +2854,7 @@ thread_module_traverse(PyObject *module, visitproc visit, void *arg)
 {
     thread_module_state *state = get_thread_state(module);
     Py_VISIT(state->excepthook_type);
+    Py_CLEAR(state->iter_locked_type);
     Py_VISIT(state->lock_type);
     Py_VISIT(state->local_type);
     Py_VISIT(state->local_dummy_type);
@@ -2751,6 +2867,7 @@ thread_module_clear(PyObject *module)
 {
     thread_module_state *state = get_thread_state(module);
     Py_CLEAR(state->excepthook_type);
+    Py_CLEAR(state->iter_locked_type);
     Py_CLEAR(state->lock_type);
     Py_CLEAR(state->local_type);
     Py_CLEAR(state->local_dummy_type);