gh-120321: Make gi_frame_state transitions atomic in FT build

This makes generator frame state transitions atomic in the free
threading build, which avoids segfaults when trying to execute
a generator from multiple threads concurrently.

There are still a few operations that aren't thread-safe and may crash
if performed concurrently on the same generator/coroutine:

 * Accessing gi_yieldfrom/cr_await/ag_await
 * Accessing gi_frame/cr_frame/ag_frame
 * Async generator operations

Author: colesbury

Include/cpython/pyatomic.h

@@ -523,6 +523,9 @@ _Py_atomic_store_uintptr_release(uintptr_t *obj, uintptr_t value);
 static inline void
 _Py_atomic_store_ssize_release(Py_ssize_t *obj, Py_ssize_t value);
 
+static inline void
+_Py_atomic_store_int8_release(int8_t *obj, int8_t value);
+
 static inline void
 _Py_atomic_store_int_release(int *obj, int value);
 

Include/cpython/pyatomic_gcc.h

@@ -572,6 +572,10 @@ static inline void
 _Py_atomic_store_int_release(int *obj, int value)
 { __atomic_store_n(obj, value, __ATOMIC_RELEASE); }
 
+static inline void
+_Py_atomic_store_int8_release(int8_t *obj, int8_t value)
+{ __atomic_store_n(obj, value, __ATOMIC_RELEASE); }
+
 static inline void
 _Py_atomic_store_ssize_release(Py_ssize_t *obj, Py_ssize_t value)
 { __atomic_store_n(obj, value, __ATOMIC_RELEASE); }

Include/cpython/pyatomic_msc.h

@@ -1066,6 +1066,19 @@ _Py_atomic_store_int_release(int *obj, int value)
 #endif
 }
 
+static inline void
+_Py_atomic_store_int8_release(int8_t *obj, int8_t value)
+{
+#if defined(_M_X64) || defined(_M_IX86)
+    *(int8_t volatile *)obj = value;
+#elif defined(_M_ARM64)
+    _Py_atomic_ASSERT_ARG_TYPE(unsigned __int8);
+    __stlr8((unsigned __int8 volatile *)obj, (unsigned __int8)value);
+#else
+#  error "no implementation of _Py_atomic_store_int8_release"
+#endif
+}
+
 static inline void
 _Py_atomic_store_ssize_release(Py_ssize_t *obj, Py_ssize_t value)
 {

Include/cpython/pyatomic_std.h

@@ -1023,6 +1023,14 @@ _Py_atomic_store_int_release(int *obj, int value)
                           memory_order_release);
 }
 
+static inline void
+_Py_atomic_store_int8_release(int8_t *obj, int8_t value)
+{
+    _Py_USING_STD;
+    atomic_store_explicit((_Atomic(int8_t)*)obj, value,
+                          memory_order_release);
+}
+
 static inline void
 _Py_atomic_store_uint_release(unsigned int *obj, unsigned int value)
 {

Include/internal/pycore_frame.h

@@ -41,20 +41,6 @@ struct _frame {
 extern PyFrameObject* _PyFrame_New_NoTrack(PyCodeObject *code);
 
 
-/* other API */
-
-typedef enum _framestate {
-    FRAME_CREATED = -3,
-    FRAME_SUSPENDED = -2,
-    FRAME_SUSPENDED_YIELD_FROM = -1,
-    FRAME_EXECUTING = 0,
-    FRAME_COMPLETED = 1,
-    FRAME_CLEARED = 4
-} PyFrameState;
-
-#define FRAME_STATE_SUSPENDED(S) ((S) == FRAME_SUSPENDED || (S) == FRAME_SUSPENDED_YIELD_FROM)
-#define FRAME_STATE_FINISHED(S) ((S) >= FRAME_COMPLETED)
-
 #ifdef __cplusplus
 }
 #endif

Include/internal/pycore_genobject.h

@@ -10,8 +10,21 @@ extern "C" {
 
 #include "pycore_interpframe_structs.h" // _PyGenObject
 
+#include <stdbool.h>              // bool
 #include <stddef.h>               // offsetof()
+#include "pycore_object.h"        // _PyObject_IsUniquelyReferenced()
 
+typedef enum _framestate {
+    FRAME_CREATED = -3,
+    FRAME_SUSPENDED = -2,
+    FRAME_SUSPENDED_YIELD_FROM = -1,
+    FRAME_EXECUTING = 0,
+    FRAME_COMPLETED = 1,
+    FRAME_CLEARED = 4
+} PyFrameState;
+
+#define FRAME_STATE_SUSPENDED(S) ((S) == FRAME_SUSPENDED || (S) == FRAME_SUSPENDED_YIELD_FROM)
+#define FRAME_STATE_FINISHED(S) ((S) >= FRAME_COMPLETED)
 
 static inline
 PyGenObject *_PyGen_GetGeneratorFromFrame(_PyInterpreterFrame *frame)
@@ -21,7 +34,30 @@ PyGenObject *_PyGen_GetGeneratorFromFrame(_PyInterpreterFrame *frame)
     return (PyGenObject *)(((char *)frame) - offset_in_gen);
 }
 
-PyAPI_FUNC(PyObject *)_PyGen_yf(PyGenObject *);
+// Mark the generator as executing. Returns true if the state was changed,
+// false if it was already executing or finished.
+static inline bool
+_PyGen_SetExecuting(PyGenObject *gen)
+{
+#ifdef Py_GIL_DISABLED
+    if (!_PyObject_IsUniquelyReferenced((PyObject *)gen)) {
+        int8_t frame_state = _Py_atomic_load_int8_relaxed(&gen->gi_frame_state);
+        while (frame_state < FRAME_EXECUTING) {
+            if (_Py_atomic_compare_exchange_int8(&gen->gi_frame_state,
+                                                 &frame_state,
+                                                 FRAME_EXECUTING)) {
+                return true;
+            }
+        }
+    }
+#endif
+    if (gen->gi_frame_state < FRAME_EXECUTING) {
+        gen->gi_frame_state = FRAME_EXECUTING;
+        return true;
+    }
+    return false;
+}
+
 extern void _PyGen_Finalize(PyObject *self);
 
 // Export for '_asyncio' shared extension

Include/internal/pycore_pyatomic_ft_wrappers.h

@@ -39,6 +39,8 @@ extern "C" {
     _Py_atomic_load_uint8(&value)
 #define FT_ATOMIC_STORE_UINT8(value, new_value) \
     _Py_atomic_store_uint8(&value, new_value)
+#define FT_ATOMIC_LOAD_INT8_RELAXED(value) \
+    _Py_atomic_load_int8_relaxed(&value)
 #define FT_ATOMIC_LOAD_UINT8_RELAXED(value) \
     _Py_atomic_load_uint8_relaxed(&value)
 #define FT_ATOMIC_LOAD_UINT16_RELAXED(value) \
@@ -53,6 +55,10 @@ extern "C" {
     _Py_atomic_store_ptr_release(&value, new_value)
 #define FT_ATOMIC_STORE_UINTPTR_RELEASE(value, new_value) \
     _Py_atomic_store_uintptr_release(&value, new_value)
+#define FT_ATOMIC_STORE_INT8_RELAXED(value, new_value) \
+    _Py_atomic_store_int8_relaxed(&value, new_value)
+#define FT_ATOMIC_STORE_INT8_RELEASE(value, new_value) \
+    _Py_atomic_store_int8_release(&value, new_value)
 #define FT_ATOMIC_STORE_SSIZE_RELAXED(value, new_value) \
     _Py_atomic_store_ssize_relaxed(&value, new_value)
 #define FT_ATOMIC_STORE_UINT8_RELAXED(value, new_value) \
@@ -129,13 +135,16 @@ extern "C" {
 #define FT_ATOMIC_LOAD_PTR_RELAXED(value) value
 #define FT_ATOMIC_LOAD_UINT8(value) value
 #define FT_ATOMIC_STORE_UINT8(value, new_value) value = new_value
+#define FT_ATOMIC_LOAD_INT8_RELAXED(value) value
 #define FT_ATOMIC_LOAD_UINT8_RELAXED(value) value
 #define FT_ATOMIC_LOAD_UINT16_RELAXED(value) value
 #define FT_ATOMIC_LOAD_UINT32_RELAXED(value) value
 #define FT_ATOMIC_LOAD_ULONG_RELAXED(value) value
 #define FT_ATOMIC_STORE_PTR_RELAXED(value, new_value) value = new_value
 #define FT_ATOMIC_STORE_PTR_RELEASE(value, new_value) value = new_value
 #define FT_ATOMIC_STORE_UINTPTR_RELEASE(value, new_value) value = new_value
+#define FT_ATOMIC_STORE_INT8_RELAXED(value, new_value) value = new_value
+#define FT_ATOMIC_STORE_INT8_RELEASE(value, new_value) value = new_value
 #define FT_ATOMIC_STORE_SSIZE_RELAXED(value, new_value) value = new_value
 #define FT_ATOMIC_STORE_UINT8_RELAXED(value, new_value) value = new_value
 #define FT_ATOMIC_STORE_UINT16_RELAXED(value, new_value) value = new_value

Include/internal/pycore_tstate.h

@@ -113,6 +113,10 @@ typedef struct _PyThreadStateImpl {
     // When >1, code objects do not immortalize their non-string constants.
     int suppress_co_const_immortalization;
 
+    // Last known frame state for generators/coroutines in this thread
+    // Used by genobject.c
+    int8_t gen_last_frame_state;
+
 #ifdef Py_STATS
      // per-thread stats, will be merged into interp->pystats_struct
      PyStats *pystats_struct; // allocated by _PyStats_ThreadInit()