gh-120321: Make gen.gi_frame.clear() thread-safe

Author: colesbury

Include/internal/pycore_genobject.h

@@ -22,7 +22,7 @@ PyGenObject *_PyGen_GetGeneratorFromFrame(_PyInterpreterFrame *frame)
 }
 
 PyAPI_FUNC(PyObject *)_PyGen_yf(PyGenObject *);
-extern void _PyGen_Finalize(PyObject *self);
+extern int _PyGen_ClearFrame(PyGenObject *self);
 
 // Export for '_asyncio' shared extension
 PyAPI_FUNC(int) _PyGen_SetStopIterationValue(PyObject *);

Objects/frameobject.c

@@ -2016,30 +2016,20 @@ frame_clear_impl(PyFrameObject *self)
 {
     if (self->f_frame->owner == FRAME_OWNED_BY_GENERATOR) {
         PyGenObject *gen = _PyGen_GetGeneratorFromFrame(self->f_frame);
-        if (gen->gi_frame_state == FRAME_EXECUTING) {
-            goto running;
-        }
-        if (FRAME_STATE_SUSPENDED(gen->gi_frame_state)) {
-            goto suspended;
+        if (_PyGen_ClearFrame(gen) < 0) {
+            return NULL;
         }
-        _PyGen_Finalize((PyObject *)gen);
     }
     else if (self->f_frame->owner == FRAME_OWNED_BY_THREAD) {
-        goto running;
+        PyErr_SetString(PyExc_RuntimeError,
+                        "cannot clear an executing frame");
+        return NULL;
     }
     else {
         assert(self->f_frame->owner == FRAME_OWNED_BY_FRAME_OBJECT);
         (void)frame_tp_clear((PyObject *)self);
     }
     Py_RETURN_NONE;
-running:
-    PyErr_SetString(PyExc_RuntimeError,
-                    "cannot clear an executing frame");
-    return NULL;
-suspended:
-    PyErr_SetString(PyExc_RuntimeError,
-                    "cannot clear a suspended frame");
-    return NULL;
 }
 
 /*[clinic input]

Objects/genobject.c

@@ -91,8 +91,8 @@ gen_traverse(PyObject *self, visitproc visit, void *arg)
     return 0;
 }
 
-void
-_PyGen_Finalize(PyObject *self)
+static void
+gen_finalize(PyObject *self)
 {
     PyGenObject *gen = (PyGenObject *)self;
 
@@ -160,6 +160,34 @@ gen_clear_frame(PyGenObject *gen)
     _PyErr_ClearExcState(&gen->gi_exc_state);
 }
 
+int
+_PyGen_ClearFrame(PyGenObject *gen)
+{
+    int8_t frame_state = FT_ATOMIC_LOAD_INT8_RELAXED(gen->gi_frame_state);
+    do {
+        if (FRAME_STATE_FINISHED(frame_state)) {
+            return 0;
+        }
+        else if (frame_state == FRAME_EXECUTING) {
+            PyErr_SetString(PyExc_RuntimeError,
+                            "cannot clear an executing frame");
+            return -1;
+        }
+        else if (FRAME_STATE_SUSPENDED(frame_state)) {
+            PyErr_SetString(PyExc_RuntimeError,
+                            "cannot clear an suspended frame");
+            return -1;
+        }
+        assert(frame_state == FRAME_CREATED);
+    } while (!_Py_GEN_TRY_SET_FRAME_STATE(gen, frame_state, FRAME_CLEARED));
+
+    if (_PyGen_GetCode(gen)->co_flags & CO_COROUTINE) {
+        _PyErr_WarnUnawaitedCoroutine((PyObject *)gen);
+    }
+    gen_clear_frame(gen);
+    return 0;
+}
+
 static void
 gen_dealloc(PyObject *self)
 {
@@ -1006,7 +1034,7 @@ PyTypeObject PyGen_Type = {
     0,                                          /* tp_weaklist */
     0,                                          /* tp_del */
     0,                                          /* tp_version_tag */
-    _PyGen_Finalize,                            /* tp_finalize */
+    gen_finalize,                               /* tp_finalize */
 };
 
 static PyObject *
@@ -1336,7 +1364,7 @@ PyTypeObject PyCoro_Type = {
     0,                                          /* tp_weaklist */
     0,                                          /* tp_del */
     0,                                          /* tp_version_tag */
-    _PyGen_Finalize,                            /* tp_finalize */
+    gen_finalize,                               /* tp_finalize */
 };
 
 static void
@@ -1762,7 +1790,7 @@ PyTypeObject PyAsyncGen_Type = {
     0,                                          /* tp_weaklist */
     0,                                          /* tp_del */
     0,                                          /* tp_version_tag */
-    _PyGen_Finalize,                            /* tp_finalize */
+    gen_finalize,                               /* tp_finalize */
 };