iterator safe from multi-crit deadlocks

Author: tom-pytel

Modules/arraymodule.c

@@ -3199,53 +3199,46 @@ array_iter(arrayobject *ao)
         return NULL;
 
     it->ao = (arrayobject*)Py_NewRef(ao);
-    it->index = 0;
+    it->index = 0;  // -1 indicates exhausted
     it->getitem = ao->ob_descr->getitem;
     PyObject_GC_Track(it);
     return (PyObject *)it;
 }
 
 static PyObject *
-arrayiter_next_lock_held(arrayiterobject *it)
+arrayiter_next(arrayiterobject *it)
 {
-    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(it);
-    arrayobject *ao;
-
+    Py_ssize_t index = FT_ATOMIC_LOAD_SSIZE_RELAXED(it->index);
+    if (index < 0) {
+        return NULL;
+    }
+    PyObject *ret;
 #ifndef NDEBUG
     array_state *state = find_array_state_by_type(Py_TYPE(it));
     assert(PyObject_TypeCheck(it, state->ArrayIterType));
+    assert(array_Check(it->ao, state));
 #endif
-    ao = it->ao;
-    if (ao == NULL) {
-        return NULL;
+
+    Py_BEGIN_CRITICAL_SECTION(it->ao);
+    if (index < Py_SIZE(it->ao)) {
+        ret = (*it->getitem)(it->ao, index);
     }
-    PyObject *ret = NULL;
-    Py_BEGIN_CRITICAL_SECTION(ao);
-#ifndef NDEBUG
-    assert(array_Check(ao, state));
-#endif
-    if (it->index < Py_SIZE(ao)) {
-        ret = (*it->getitem)(ao, it->index++);
+    else {
+        ret = NULL;
     }
     Py_END_CRITICAL_SECTION();
+
     if (ret != NULL) {
-        return ret;
+        FT_ATOMIC_STORE_SSIZE_RELAXED(it->index, index + 1);
+    }
+    else {
+        FT_ATOMIC_STORE_SSIZE_RELAXED(it->index, -1);
+#ifndef Py_GIL_DISABLED
+        arrayobject *ao = it->ao;
+        it->ao = NULL;
+        Py_DECREF(ao);
+#endif
     }
-    it->ao = NULL;
-    Py_DECREF(ao);
-    return NULL;
-}
-
-static PyObject *
-arrayiter_next(arrayiterobject *it)
-{
-    PyObject *ret;
-    assert(it != NULL);
-
-    Py_BEGIN_CRITICAL_SECTION(it);
-    ret = arrayiter_next_lock_held(it);
-    Py_END_CRITICAL_SECTION();
-
     return ret;
 }
 
@@ -3269,7 +3262,6 @@ arrayiter_traverse(arrayiterobject *it, visitproc visit, void *arg)
 }
 
 /*[clinic input]
-@critical_section
 array.arrayiterator.__reduce__
 
     cls: defining_class
@@ -3280,19 +3272,27 @@ Return state information for pickling.
 
 static PyObject *
 array_arrayiterator___reduce___impl(arrayiterobject *self, PyTypeObject *cls)
-/*[clinic end generated code: output=4b032417a2c8f5e6 input=61ad213fe49ae0f7]*/
+/*[clinic end generated code: output=4b032417a2c8f5e6 input=ac64e65a87ad452e]*/
 {
     array_state *state = get_array_state_by_class(cls);
     assert(state != NULL);
     PyObject *func = _PyEval_GetBuiltin(state->str_iter);
-    if (self->ao == NULL) {
-        return Py_BuildValue("N(())", func);
+    PyObject *ret = NULL;
+    Py_ssize_t index = FT_ATOMIC_LOAD_SSIZE_RELAXED(self->index);
+    if (index >= 0) {
+        Py_BEGIN_CRITICAL_SECTION(self->ao);
+        if (index <= Py_SIZE(self->ao)) {
+            ret = Py_BuildValue("N(O)n", func, self->ao, index);
+        }
+        Py_END_CRITICAL_SECTION();
+    }
+    if (ret == NULL) {
+        ret = Py_BuildValue("N(())", func);
     }
-    return Py_BuildValue("N(O)n", func, self->ao, self->index);
+    return ret;
 }
 
 /*[clinic input]
-@critical_section
 array.arrayiterator.__setstate__
 
     state: object
@@ -3302,24 +3302,26 @@ Set state information for unpickling.
 [clinic start generated code]*/
 
 static PyObject *
-array_arrayiterator___setstate___impl(arrayiterobject *self, PyObject *state)
-/*[clinic end generated code: output=d7837ae4ac1fd8b9 input=8d8dc7ce40b9c1f7]*/
+array_arrayiterator___setstate__(arrayiterobject *self, PyObject *state)
+/*[clinic end generated code: output=397da9904e443cbe input=f47d5ceda19e787b]*/
 {
     Py_ssize_t index = PyLong_AsSsize_t(state);
     if (index == -1 && PyErr_Occurred()) {
         return NULL;
     }
-    arrayobject *ao = self->ao;
-    if (ao != NULL) {
-        Py_BEGIN_CRITICAL_SECTION(ao);
-        if (index < 0) {
-            index = 0;
+    if (FT_ATOMIC_LOAD_SSIZE_RELAXED(self->index) >= 0) {
+        Py_BEGIN_CRITICAL_SECTION(self->ao);
+        if (index < -1) {
+            index = -1;
         }
-        else if (index > Py_SIZE(ao)) {
-            index = Py_SIZE(ao); /* iterator exhausted */
+        else {
+            Py_ssize_t size = Py_SIZE(self->ao);
+            if (index > size) {
+                index = size; /* iterator at end */
+            }
         }
+        FT_ATOMIC_STORE_SSIZE_RELAXED(self->index, index);
         Py_END_CRITICAL_SECTION();
-        self->index = index;
     }
     Py_RETURN_NONE;
 }