[mypyc] Make native attribute writes faster with free threading

mypyc/codegen/emitclass.py

@@ -1190,21 +1190,22 @@ def generate_getter(cl: ClassIR, attr: str, rtype: RType, emitter: Emitter) -> N
     attr_expr = f"self->{attr_field}"
 
     if IS_FREE_THREADED and is_simple_refcounted_pointer(rtype):
-        # In free-threaded builds, load the attribute and take a new reference
-        # atomically to avoid a use-after-free race with a concurrent setter.
+        # In free-threaded builds, load the attribute and take a new reference with
+        # an optimistic validated incref to avoid racing with a concurrent setter.
         # CPy_GetAttrRef returns NULL if the attribute is undefined (NULL field),
         # which is exactly the error/undefined value for a 'PyObject *' field.
         #
         # Final attributes are never rebound (no setter), so there is no concurrent
         # writer to race with: a plain load + incref is safe. Use the cheaper
-        # CPy_GetAttrRefFinal, which skips the try-incref and _Py_NewRefWithLock
-        # slow path entirely (an unconditional Py_INCREF needs no maybe-weakref).
+        # CPy_GetAttrRefFinal, which skips the try-incref and the locked
+        # _Py_XNewRefWithLock fallback entirely (an unconditional Py_INCREF needs no
+        # maybe-weakref).
         # This getter is generated per defining class, so a direct membership test
         # matches the read-only getset table above (no need to walk the MRO).
         if attr in cl.final_attributes:
             getattr_ref = f"CPy_GetAttrRefFinal((PyObject **)&{attr_expr})"
         else:
-            getattr_ref = f"CPy_GetAttrRef((PyObject **)&{attr_expr})"
+            getattr_ref = f"CPy_GetAttrRef((PyObject *)self, (PyObject **)&{attr_expr})"
         emitter.emit_line(f"PyObject *retval = {getattr_ref};")
         emitter.emit_line("if (unlikely(retval == NULL)) {")
         emitter.emit_line("PyErr_SetString(PyExc_AttributeError,")
@@ -1253,11 +1254,12 @@ def generate_setter(cl: ClassIR, attr: str, rtype: RType, emitter: Emitter) -> N
         emitter.emit_line("}")
 
     if IS_FREE_THREADED and is_simple_refcounted_pointer(rtype):
-        # In free-threaded builds, publish the new value atomically via
-        # CPy_SetAttrRef so a concurrent reader (see CPy_GetAttrRef) never sees a
-        # torn pointer or a freed old value. CPy_SetAttrRef steals its value and
-        # reclaims the old one, so we cast/type-check the incoming value, take a
-        # new reference (the setter only borrows 'value'), then hand it over.
+        # In free-threaded builds, publish the new value via CPy_SetAttrRef, which
+        # takes the owner's critical section so a concurrent reader (see
+        # CPy_GetAttrRef) can always secure a reference to the value it observes,
+        # even though the old value is decrefed right away. CPy_SetAttrRef steals its
+        # value and reclaims the old one, so we cast/type-check the incoming value,
+        # take a new reference (the setter only borrows 'value'), then hand it over.
         # A NULL value deletes the attribute (reclaims the old value, stores NULL).
         if deletable:
             emitter.emit_line("if (value != NULL) {")
@@ -1267,10 +1269,14 @@ def generate_setter(cl: ClassIR, attr: str, rtype: RType, emitter: Emitter) -> N
             emitter.emit_cast("value", "tmp", rtype, declare_dest=True)
             emitter.emit_lines("if (!tmp)", "    return -1;")
         emitter.emit_inc_ref("tmp", rtype)
-        emitter.emit_line(f"CPy_SetAttrRef((PyObject **)&self->{attr_field}, tmp);")
+        emitter.emit_line(
+            f"CPy_SetAttrRef((PyObject *)self, (PyObject **)&self->{attr_field}, tmp);"
+        )
         if deletable:
             emitter.emit_line("} else {")
-            emitter.emit_line(f"CPy_SetAttrRef((PyObject **)&self->{attr_field}, NULL);")
+            emitter.emit_line(
+                f"CPy_SetAttrRef((PyObject *)self, (PyObject **)&self->{attr_field}, NULL);"
+            )
             emitter.emit_line("}")
         emitter.emit_line("return 0;")
         emitter.emit_line("}")

mypyc/codegen/emitfunc.py

@@ -407,10 +407,12 @@ def emit_load_attr_take_ref(
         """Emit the load of a native attribute into 'dest', taking a new reference.
 
         On free-threaded builds, reading a single reference-counted 'PyObject *' field
-        and taking a new reference must be done atomically to avoid a use-after-free
-        race with a concurrent setter. CPy_GetAttrRef performs the load and incref
-        atomically and returns a new reference (or NULL if undefined), so callers must
-        NOT emit a separate inc_ref. Return True in that case so the caller can skip it.
+        and taking a new reference must be synchronized to avoid a use-after-free race
+        with a concurrent setter. CPy_GetAttrRef uses an optimistic validated incref
+        with a locked fallback and returns a new reference (or NULL if undefined, which
+        on a free-threaded build includes an attribute deleted by another thread while
+        the read was in flight), so callers must NOT emit a separate inc_ref. Return
+        True in that case so the caller can skip it.
 
         Final attributes are never rebound (no setter), so there is no concurrent writer
         and no use-after-free window; an owned read uses the cheaper CPy_GetAttrRefFinal
@@ -425,7 +427,9 @@ def emit_load_attr_take_ref(
         if use_get_attr_ref and cl.is_final_attr(op.attr):
             self.emitter.emit_line(f"{dest} = CPy_GetAttrRefFinal((PyObject **)&{attr_expr});")
         elif use_get_attr_ref:
-            self.emitter.emit_line(f"{dest} = CPy_GetAttrRef((PyObject **)&{attr_expr});")
+            self.emitter.emit_line(
+                f"{dest} = CPy_GetAttrRef((PyObject *){obj}, (PyObject **)&{attr_expr});"
+            )
         else:
             self.emitter.emit_line(f"{dest} = {attr_expr};")
         return use_get_attr_ref
@@ -580,19 +584,23 @@ def visit_set_attr(self, op: SetAttr) -> None:
             self.emitter.emit_error_check(tmp, ret_type, f"{dest} = 0;")
         elif IS_FREE_THREADED and is_simple_refcounted_pointer(attr_rtype):
             # In free-threaded builds, publishing a single reference-counted
-            # 'PyObject *' field must be atomic so a concurrent reader (see
-            # CPy_GetAttrRef) never observes a torn pointer or a freed value.
-            # Both helpers steal the reference to src.
+            # 'PyObject *' field must be synchronized with concurrent readers (see
+            # CPy_GetAttrRef): the store is atomic, and CPy_SetAttrRef additionally
+            # takes the owner's critical section, which is what lets it decref the
+            # old value right away. Both helpers steal the reference to src.
             attr_expr = self.get_attr_expr(obj, op, decl_cl)
             if op.is_init:
-                # The attribute is known to be previously undefined (NULL), so
-                # there is no old value to reclaim; a relaxed store suffices
-                # (self's later publication provides the release barrier -- see
-                # CPy_InitAttrRef).
+                # The attribute is known to be previously undefined (NULL) and self
+                # can't have leaked yet, so there is no old value to reclaim and no
+                # competing writer; a relaxed store suffices (self's later publication
+                # provides the release barrier -- see CPy_InitAttrRef).
                 self.emitter.emit_line(f"CPy_InitAttrRef((PyObject **)&{attr_expr}, {src});")
             else:
-                # Atomically swap in the new value and reclaim the old one.
-                self.emitter.emit_line(f"CPy_SetAttrRef((PyObject **)&{attr_expr}, {src});")
+                # Replace the value under the owner's critical section and reclaim
+                # the old one with a normal decref.
+                self.emitter.emit_line(
+                    f"CPy_SetAttrRef((PyObject *){obj}, (PyObject **)&{attr_expr}, {src});"
+                )
             if op.error_kind == ERR_FALSE:
                 self.emitter.emit_line(f"{dest} = 1;")
         else:

mypyc/lib-rt/mypyc_util.h

@@ -17,10 +17,13 @@
 
 #if defined(__clang__) || defined(__GNUC__)
 #define CPy_NOINLINE __attribute__((noinline))
+#define CPy_COLD __attribute__((cold))
 #elif defined(_MSC_VER)
 #define CPy_NOINLINE __declspec(noinline)
+#define CPy_COLD
 #else
 #define CPy_NOINLINE
+#define CPy_COLD
 #endif
 
 #ifndef Py_GIL_DISABLED

mypyc/lib-rt/pythonsupport.c

@@ -6,23 +6,33 @@
 #include "pythonsupport.h"
 
 #ifdef Py_GIL_DISABLED
-// Cold slow path of CPy_GetAttrRef (declared in pythonsupport.h). Reached only
-// when the inline fast-path try-incref fails: the value is owned by another
-// thread, so taking a reference requires an atomic shared-refcount operation.
-// Kept out-of-line so the inline fast path stays small.
+// Cold slow path of CPy_GetAttrRef (declared in pythonsupport.h). First try to
+// acquire a shared reference without locking, then validate that the field still
+// contains v. If validation fails, drop the provisional reference. If the shared
+// incref or validation fails, take the owner's critical section, reload the field,
+// and take a reference while the value cannot be replaced. _Py_XNewRefWithLock
+// also sets maybe-weakref lazily, so later cross-thread reads generally use the
+// lock-free shared-refcount path. Returns NULL if the field is NULL on reload.
 //
-// First try the lock-free shared-refcount CAS. If the value has not had
-// maybe-weakref set yet (for example, it was published by CPy_InitAttrRef), force
-// a cross-thread reference via _Py_NewRefWithLock, which cannot fail and sets
-// maybe-weakref so subsequent reads take the fast path. The value was already
-// observed in the field by CPy_GetAttrRef; CPy_SetAttrRef's QSBR-delayed decref
-// keeps any replaced value alive long enough for this reader.
+// The reload is deliberately relaxed, where CPython 3.14 uses a consume load in the
+// equivalent locked read (_PyObject_TryGetInstanceAttribute): acquiring the owner's
+// critical section already pairs with CPy_SetAttrRef's release store, and a value
+// published by CPy_InitAttrRef is ordered by the publication of the owner itself
+// (see CPy_InitAttrRef).
 CPy_NOINLINE
-PyObject *CPy_GetAttrRefSlow(PyObject *v) {
+CPy_COLD
+PyObject *CPy_GetAttrRefSlow(PyObject *v, PyObject *owner, PyObject **field) {
     if (_Py_TryIncRefShared(v)) {
-        return v;
+        if (v == (PyObject *)_Py_atomic_load_ptr(field)) {
+            return v;
+        }
+        Py_DECREF(v);
     }
-    return _Py_NewRefWithLock(v);  // sets maybe-weakref; cannot fail
+    PyObject *result;
+    Py_BEGIN_CRITICAL_SECTION(owner);
+    result = _Py_XNewRefWithLock((PyObject *)_Py_atomic_load_ptr_relaxed(field));
+    Py_END_CRITICAL_SECTION();
+    return result;
 }
 #endif