use recursive lock

Author: eendebakpt

Lib/test/test_threading.py

@@ -2457,5 +2457,35 @@ def test_iter_locked(self):
             for arg in [1, None, True, sys]:
                 self.assertRaises(TypeError, threading.iter_locked, arg)
 
+     def test_iter_locked_recursive_generator(self):
+            def g():
+                yield next(it)
+
+            it = threading.iter_locked(g())
+            with self.assertRaises(ValueError):
+                next(it)
+
+     def test_iter_locked_recursive_iterator(self):
+        class evil_it:
+            def __init__(self):
+                self.it = iter(range(100))
+                self.it2 = None
+
+            def __iter__(self):
+                return self
+
+            def __next__(self):
+                if self.it2 is not None:
+                    it2, self.it2 = self.it2, None
+                    return next(it2), next(it2)
+                return next(self.it)
+
+        a = evil_it()
+        il = threading.iter_locked(a)
+        assert next(il) == 0
+        a.it2 = il
+        assert next(il) == (1, 2)
+
+
 if __name__ == "__main__":
     unittest.main()

Modules/_threadmodule.c

@@ -759,7 +759,7 @@ static PyType_Spec ThreadHandle_Type_spec = {
 typedef struct {
     PyObject_HEAD
     PyObject *it;
-    PyMutex lock;
+    _PyRecursiveMutex lock;
 } iter_locked_object;
 
 #define iter_locked_object_CAST(op)    ((iter_locked_object *)(op))
@@ -787,7 +787,7 @@ _thread_iter_locked_impl(PyTypeObject *type, PyObject *iterable)
         return NULL;
     }
     il->it = it;
-    il->lock = (PyMutex){0};
+    il->lock = (_PyRecursiveMutex){0};
 
     return (PyObject *)il;
 }
@@ -818,15 +818,14 @@ iter_locked_next(PyObject *op)
     iter_locked_object *lz = iter_locked_object_CAST(op);
     PyObject *result = NULL;
 
-    // we cannot use Py_BEGIN_CRITICAL_SECTION as it is not available in the normal build
-    PyMutex_Lock(&(lz->lock));
+    _PyRecursiveMutex_Lock(&(lz->lock));
 
     int v = PyIter_NextItem(lz->it, &result);
     if (v == -1) {
         /* Note:  StopIteration is already cleared by PyIter_Next() */
         /* If PyErr_Occurred() we will also return NULL*/
     }
-    PyMutex_Unlock(&(lz->lock));
+    _PyRecursiveMutex_Unlock(&(lz->lock));
 
     return result;
 }