gh-133931: Introduce _PyObject_XSetRefDelayed to replace Py_XSETREF

Include/internal/pycore_pymem.h

@@ -100,6 +100,17 @@ static inline void _PyObject_XDecRefDelayed(PyObject *obj)
 }
 #endif
 
+#ifdef Py_GIL_DISABLED
+// Same as `Py_XSETREF` but in free-threading, it stores the object atomically
+// and queues the old object to be decrefed at a safe point using QSBR.
+PyAPI_FUNC(void) _PyObject_XSetRefDelayed(PyObject **p_obj, PyObject *obj);
+#else
+static inline void _PyObject_XSetRefDelayed(PyObject **p_obj, PyObject *obj)
+{
+    Py_XSETREF(*p_obj, obj);
+}
+#endif
+
 // Periodically process delayed free requests.
 extern void _PyMem_ProcessDelayed(PyThreadState *tstate);
 

Lib/test/test_free_threading/test_generators.py

@@ -0,0 +1,51 @@
+import concurrent.futures
+import unittest
+from threading import Barrier
+from unittest import TestCase
+import random
+import time
+
+from test.support import threading_helper, Py_GIL_DISABLED
+
+threading_helper.requires_working_threading(module=True)
+
+
+def random_sleep():
+    delay_us = random.randint(50, 100)
+    time.sleep(delay_us * 1e-6)
+
+def random_string():
+    return ''.join(random.choice('0123456789ABCDEF') for _ in range(10))
+
+def set_gen_name(g, b):
+    b.wait()
+    random_sleep()
+    g.__name__ = random_string()
+    return g.__name__
+
+def set_gen_qualname(g, b):
+    b.wait()
+    random_sleep()
+    g.__qualname__ = random_string()
+    return g.__qualname__
+
+
+@unittest.skipUnless(Py_GIL_DISABLED, "Enable only in FT build")
+class TestFTGenerators(TestCase):
+    NUM_THREADS = 4
+
+    def concurrent_write_with_func(self, func):
+        gen = (x for x in range(42))
+        for j in range(1000):
+            with concurrent.futures.ThreadPoolExecutor(max_workers=self.NUM_THREADS) as executor:
+                b = Barrier(self.NUM_THREADS)
+                futures = {executor.submit(func, gen, b): i for i in range(self.NUM_THREADS)}
+                for fut in concurrent.futures.as_completed(futures):
+                    gen_name = fut.result()
+                    self.assertEqual(len(gen_name), 10)
+
+    def test_concurrent_write(self):
+        with self.subTest(func=set_gen_name):
+            self.concurrent_write_with_func(func=set_gen_name)
+        with self.subTest(func=set_gen_qualname):
+            self.concurrent_write_with_func(func=set_gen_qualname)

Objects/genobject.c

@@ -10,6 +10,7 @@
 #include "pycore_gc.h"            // _PyGC_CLEAR_FINALIZED()
 #include "pycore_genobject.h"     // _PyGen_SetStopIterationValue()
 #include "pycore_interpframe.h"   // _PyFrame_GetCode()
+#include "pycore_pymem.h"         // _PyObject_XSetRefDelayed()
 #include "pycore_modsupport.h"    // _PyArg_CheckPositional()
 #include "pycore_object.h"        // _PyObject_GC_UNTRACK()
 #include "pycore_opcode_utils.h"  // RESUME_AFTER_YIELD_FROM
@@ -704,7 +705,8 @@ static PyObject *
 gen_get_name(PyObject *self, void *Py_UNUSED(ignored))
 {
     PyGenObject *op = _PyGen_CAST(self);
-    return Py_NewRef(op->gi_name);
+    PyObject *name = FT_ATOMIC_LOAD_PTR_RELAXED(op->gi_name);
+    return Py_NewRef(name);
 }
 
 static int
@@ -718,15 +720,19 @@ gen_set_name(PyObject *self, PyObject *value, void *Py_UNUSED(ignored))
                         "__name__ must be set to a string object");
         return -1;
     }
-    Py_XSETREF(op->gi_name, Py_NewRef(value));
+    Py_BEGIN_CRITICAL_SECTION(self);
+    // To prevent use-after-free from other threads that reference the gi_name.
+    _PyObject_XSetRefDelayed(&op->gi_name, Py_NewRef(value));
+    Py_END_CRITICAL_SECTION();
     return 0;
 }
 
 static PyObject *
 gen_get_qualname(PyObject *self, void *Py_UNUSED(ignored))
 {
     PyGenObject *op = _PyGen_CAST(self);
-    return Py_NewRef(op->gi_qualname);
+    PyObject *qualname = FT_ATOMIC_LOAD_PTR_RELAXED(op->gi_qualname);
+    return Py_NewRef(qualname);
 }
 
 static int
@@ -740,7 +746,10 @@ gen_set_qualname(PyObject *self, PyObject *value, void *Py_UNUSED(ignored))
                         "__qualname__ must be set to a string object");
         return -1;
     }
-    Py_XSETREF(op->gi_qualname, Py_NewRef(value));
+    Py_BEGIN_CRITICAL_SECTION(self);
+    // To prevent use-after-free from other threads that reference the gi_qualname.
+    _PyObject_XSetRefDelayed(&op->gi_qualname, Py_NewRef(value));
+    Py_END_CRITICAL_SECTION();
     return 0;
 }
 

Objects/obmalloc.c

@@ -1231,6 +1231,21 @@ _PyObject_XDecRefDelayed(PyObject *ptr)
 }
 #endif
 
+#ifdef Py_GIL_DISABLED
+void
+_PyObject_XSetRefDelayed(PyObject **ptr, PyObject *value)
+{
+    PyObject *old = FT_ATOMIC_LOAD_PTR_ACQUIRE(*ptr);
+    FT_ATOMIC_STORE_PTR_RELEASE(*ptr, value);
+    if (old == NULL) {
+        return;
+    }
+    if (!_Py_IsImmortal(old)) {
+         _PyObject_XDecRefDelayed(old);
+    }
+}
+#endif
+
 static struct _mem_work_chunk *
 work_queue_first(struct llist_node *head)
 {

Objects/typeobject.c

@@ -3967,13 +3967,8 @@ _PyObject_SetDict(PyObject *obj, PyObject *value)
         return -1;
     }
     Py_BEGIN_CRITICAL_SECTION(obj);
-    PyObject *olddict = *dictptr;
-    FT_ATOMIC_STORE_PTR_RELEASE(*dictptr, Py_NewRef(value));
-#ifdef Py_GIL_DISABLED
-    _PyObject_XDecRefDelayed(olddict);
-#else
-    Py_XDECREF(olddict);
-#endif
+    // To prevent use-after-free from other threads that reference the __dict__
+    _PyObject_XSetRefDelayed(dictptr, Py_NewRef(value));
     Py_END_CRITICAL_SECTION();
     return 0;
 }