gh-111926: Make weakrefs thread-safe in free-threaded builds (#117168)

Most mutable data is protected by a striped lock that is keyed on the
referenced object's address. The weakref's hash is protected using the
weakref's per-object lock.
 
Note that this only affects free-threaded builds. Apart from some minor
refactoring, the added code is all either gated by `ifdef`s or is a no-op
(e.g. `Py_BEGIN_CRITICAL_SECTION`).

Author: mpage

Diff for: Include/cpython/weakrefobject.h

@@ -30,6 +30,14 @@ struct _PyWeakReference {
     PyWeakReference *wr_prev;
     PyWeakReference *wr_next;
     vectorcallfunc vectorcall;
+
+#ifdef Py_GIL_DISABLED
+    /* Pointer to the lock used when clearing in free-threaded builds.
+     * Normally this can be derived from wr_object, but in some cases we need
+     * to lock after wr_object has been set to Py_None.
+     */
+    struct _PyMutex *weakrefs_lock;
+#endif
 };
 
 Py_DEPRECATED(3.13) static inline PyObject* PyWeakref_GET_OBJECT(PyObject *ref_obj)

Diff for: Include/internal/pycore_interp.h

@@ -59,6 +59,12 @@ struct _stoptheworld_state {
     PyThreadState *requester; // Thread that requested the pause (may be NULL).
 };
 
+#ifdef Py_GIL_DISABLED
+// This should be prime but otherwise the choice is arbitrary. A larger value
+// increases concurrency at the expense of memory.
+#  define NUM_WEAKREF_LIST_LOCKS 127
+#endif
+
 /* cross-interpreter data registry */
 
 /* Tracks some rare events per-interpreter, used by the optimizer to turn on/off
@@ -203,6 +209,7 @@ struct _is {
 #if defined(Py_GIL_DISABLED)
     struct _mimalloc_interp_state mimalloc;
     struct _brc_state brc;  // biased reference counting state
+    PyMutex weakref_locks[NUM_WEAKREF_LIST_LOCKS];
 #endif
 
     // Per-interpreter state for the obmalloc allocator.  For the main

Diff for: Include/internal/pycore_object.h

@@ -426,7 +426,7 @@ _Py_TryIncRefShared(PyObject *op)
 
 /* Tries to incref the object op and ensures that *src still points to it. */
 static inline int
-_Py_TryIncref(PyObject **src, PyObject *op)
+_Py_TryIncrefCompare(PyObject **src, PyObject *op)
 {
     if (_Py_TryIncrefFast(op)) {
         return 1;
@@ -452,7 +452,7 @@ _Py_XGetRef(PyObject **ptr)
         if (value == NULL) {
             return value;
         }
-        if (_Py_TryIncref(ptr, value)) {
+        if (_Py_TryIncrefCompare(ptr, value)) {
             return value;
         }
     }
@@ -467,7 +467,7 @@ _Py_TryXGetRef(PyObject **ptr)
     if (value == NULL) {
         return value;
     }
-    if (_Py_TryIncref(ptr, value)) {
+    if (_Py_TryIncrefCompare(ptr, value)) {
         return value;
     }
     return NULL;
@@ -506,8 +506,42 @@ _Py_XNewRefWithLock(PyObject *obj)
     return _Py_NewRefWithLock(obj);
 }
 
+static inline void
+_PyObject_SetMaybeWeakref(PyObject *op)
+{
+    if (_Py_IsImmortal(op)) {
+        return;
+    }
+    for (;;) {
+        Py_ssize_t shared = _Py_atomic_load_ssize_relaxed(&op->ob_ref_shared);
+        if ((shared & _Py_REF_SHARED_FLAG_MASK) != 0) {
+            // Nothing to do if it's in WEAKREFS, QUEUED, or MERGED states.
+            return;
+        }
+        if (_Py_atomic_compare_exchange_ssize(
+                &op->ob_ref_shared, &shared, shared | _Py_REF_MAYBE_WEAKREF)) {
+            return;
+        }
+    }
+}
+
 #endif
 
+/* Tries to incref op and returns 1 if successful or 0 otherwise. */
+static inline int
+_Py_TryIncref(PyObject *op)
+{
+#ifdef Py_GIL_DISABLED
+    return _Py_TryIncrefFast(op) || _Py_TryIncRefShared(op);
+#else
+    if (Py_REFCNT(op) > 0) {
+        Py_INCREF(op);
+        return 1;
+    }
+    return 0;
+#endif
+}
+
 #ifdef Py_REF_DEBUG
 extern void _PyInterpreterState_FinalizeRefTotal(PyInterpreterState *);
 extern void _Py_FinalizeRefTotal(_PyRuntimeState *);

Diff for: Include/internal/pycore_pyatomic_ft_wrappers.h

@@ -20,18 +20,23 @@ extern "C" {
 #endif
 
 #ifdef Py_GIL_DISABLED
+#define FT_ATOMIC_LOAD_PTR(value) _Py_atomic_load_ptr(&value)
 #define FT_ATOMIC_LOAD_SSIZE(value) _Py_atomic_load_ssize(&value)
 #define FT_ATOMIC_LOAD_SSIZE_RELAXED(value) \
     _Py_atomic_load_ssize_relaxed(&value)
+#define FT_ATOMIC_STORE_PTR(value, new_value) \
+    _Py_atomic_store_ptr(&value, new_value)
 #define FT_ATOMIC_STORE_PTR_RELAXED(value, new_value) \
     _Py_atomic_store_ptr_relaxed(&value, new_value)
 #define FT_ATOMIC_STORE_PTR_RELEASE(value, new_value) \
     _Py_atomic_store_ptr_release(&value, new_value)
 #define FT_ATOMIC_STORE_SSIZE_RELAXED(value, new_value) \
     _Py_atomic_store_ssize_relaxed(&value, new_value)
 #else
+#define FT_ATOMIC_LOAD_PTR(value) value
 #define FT_ATOMIC_LOAD_SSIZE(value) value
 #define FT_ATOMIC_LOAD_SSIZE_RELAXED(value) value
+#define FT_ATOMIC_STORE_PTR(value, new_value) value = new_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_SSIZE_RELAXED(value, new_value) value = new_value

Diff for: Include/internal/pycore_weakref.h

@@ -9,7 +9,35 @@ extern "C" {
 #endif
 
 #include "pycore_critical_section.h" // Py_BEGIN_CRITICAL_SECTION()
+#include "pycore_lock.h"
 #include "pycore_object.h"           // _Py_REF_IS_MERGED()
+#include "pycore_pyatomic_ft_wrappers.h"
+
+#ifdef Py_GIL_DISABLED
+
+#define WEAKREF_LIST_LOCK(obj) \
+    _PyInterpreterState_GET()  \
+        ->weakref_locks[((uintptr_t)obj) % NUM_WEAKREF_LIST_LOCKS]
+
+// Lock using the referenced object
+#define LOCK_WEAKREFS(obj) \
+    PyMutex_LockFlags(&WEAKREF_LIST_LOCK(obj), _Py_LOCK_DONT_DETACH)
+#define UNLOCK_WEAKREFS(obj) PyMutex_Unlock(&WEAKREF_LIST_LOCK(obj))
+
+// Lock using a weakref
+#define LOCK_WEAKREFS_FOR_WR(wr) \
+    PyMutex_LockFlags(wr->weakrefs_lock, _Py_LOCK_DONT_DETACH)
+#define UNLOCK_WEAKREFS_FOR_WR(wr) PyMutex_Unlock(wr->weakrefs_lock)
+
+#else
+
+#define LOCK_WEAKREFS(obj)
+#define UNLOCK_WEAKREFS(obj)
+
+#define LOCK_WEAKREFS_FOR_WR(wr)
+#define UNLOCK_WEAKREFS_FOR_WR(wr)
+
+#endif
 
 static inline int _is_dead(PyObject *obj)
 {
@@ -30,53 +58,64 @@ static inline int _is_dead(PyObject *obj)
 static inline PyObject* _PyWeakref_GET_REF(PyObject *ref_obj)
 {
     assert(PyWeakref_Check(ref_obj));
-    PyObject *ret = NULL;
-    Py_BEGIN_CRITICAL_SECTION(ref_obj);
     PyWeakReference *ref = _Py_CAST(PyWeakReference*, ref_obj);
-    PyObject *obj = ref->wr_object;
 
+    PyObject *obj = FT_ATOMIC_LOAD_PTR(ref->wr_object);
     if (obj == Py_None) {
         // clear_weakref() was called
-        goto end;
+        return NULL;
     }
 
-    if (_is_dead(obj)) {
-        goto end;
+    LOCK_WEAKREFS(obj);
+#ifdef Py_GIL_DISABLED
+    if (ref->wr_object == Py_None) {
+        // clear_weakref() was called
+        UNLOCK_WEAKREFS(obj);
+        return NULL;
     }
-#if !defined(Py_GIL_DISABLED)
-    assert(Py_REFCNT(obj) > 0);
 #endif
-    ret = Py_NewRef(obj);
-end:
-    Py_END_CRITICAL_SECTION();
-    return ret;
+    if (_Py_TryIncref(obj)) {
+        UNLOCK_WEAKREFS(obj);
+        return obj;
+    }
+    UNLOCK_WEAKREFS(obj);
+    return NULL;
 }
 
 static inline int _PyWeakref_IS_DEAD(PyObject *ref_obj)
 {
     assert(PyWeakref_Check(ref_obj));
     int ret = 0;
-    Py_BEGIN_CRITICAL_SECTION(ref_obj);
     PyWeakReference *ref = _Py_CAST(PyWeakReference*, ref_obj);
-    PyObject *obj = ref->wr_object;
+    PyObject *obj = FT_ATOMIC_LOAD_PTR(ref->wr_object);
     if (obj == Py_None) {
         // clear_weakref() was called
         ret = 1;
     }
     else {
+        LOCK_WEAKREFS(obj);
         // See _PyWeakref_GET_REF() for the rationale of this test
+#ifdef Py_GIL_DISABLED
+        ret = (ref->wr_object == Py_None) || _is_dead(obj);
+#else
         ret = _is_dead(obj);
+#endif
+        UNLOCK_WEAKREFS(obj);
     }
-    Py_END_CRITICAL_SECTION();
     return ret;
 }
 
-extern Py_ssize_t _PyWeakref_GetWeakrefCount(PyWeakReference *head);
+extern Py_ssize_t _PyWeakref_GetWeakrefCount(PyObject *obj);
+
+// Clear all the weak references to obj but leave their callbacks uncalled and
+// intact.
+extern void _PyWeakref_ClearWeakRefsExceptCallbacks(PyObject *obj);
 
 extern void _PyWeakref_ClearRef(PyWeakReference *self);
 
+PyAPI_FUNC(int) _PyWeakref_IsDead(PyObject *weakref);
+
 #ifdef __cplusplus
 }
 #endif
 #endif /* !Py_INTERNAL_WEAKREF_H */
-

Diff for: Lib/test/test_sys.py

@@ -1708,11 +1708,15 @@ class newstyleclass(object): pass
         # TODO: add check that forces layout of unicodefields
         # weakref
         import weakref
-        check(weakref.ref(int), size('2Pn3P'))
+        if support.Py_GIL_DISABLED:
+            expected = size('2Pn4P')
+        else:
+            expected = size('2Pn3P')
+        check(weakref.ref(int), expected)
         # weakproxy
         # XXX
         # weakcallableproxy
-        check(weakref.proxy(int), size('2Pn3P'))
+        check(weakref.proxy(int), expected)
 
     def check_slots(self, obj, base, extra):
         expected = sys.getsizeof(base) + struct.calcsize(extra)

Diff for: Lib/test/test_weakref.py

@@ -1907,6 +1907,25 @@ def test_threaded_weak_valued_consistency(self):
                 self.assertEqual(len(d), 1)
                 o = None  # lose ref
 
+    @support.cpython_only
+    def test_weak_valued_consistency(self):
+        # A single-threaded, deterministic repro for issue #28427: old keys
+        # should not remove new values from WeakValueDictionary. This relies on
+        # an implementation detail of CPython's WeakValueDictionary (its
+        # underlying dictionary of KeyedRefs) to reproduce the issue.
+        d = weakref.WeakValueDictionary()
+        with support.disable_gc():
+            d[10] = RefCycle()
+            # Keep the KeyedRef alive after it's replaced so that GC will invoke
+            # the callback.
+            wr = d.data[10]
+            # Replace the value with something that isn't cyclic garbage
+            o = RefCycle()
+            d[10] = o
+            # Trigger GC, which will invoke the callback for `wr`
+            gc.collect()
+            self.assertEqual(len(d), 1)
+
     def check_threaded_weak_dict_copy(self, type_, deepcopy):
         # `type_` should be either WeakKeyDictionary or WeakValueDictionary.
         # `deepcopy` should be either True or False.

Diff for: Modules/_sqlite/blob.c

@@ -4,7 +4,6 @@
 
 #include "blob.h"
 #include "util.h"
-#include "pycore_weakref.h"       // _PyWeakref_GET_REF()
 
 #define clinic_state() (pysqlite_get_state_by_type(Py_TYPE(self)))
 #include "clinic/blob.c.h"
@@ -102,8 +101,8 @@ pysqlite_close_all_blobs(pysqlite_Connection *self)
 {
     for (int i = 0; i < PyList_GET_SIZE(self->blobs); i++) {
         PyObject *weakref = PyList_GET_ITEM(self->blobs, i);
-        PyObject *blob = _PyWeakref_GET_REF(weakref);
-        if (blob == NULL) {
+        PyObject *blob;
+        if (!PyWeakref_GetRef(weakref, &blob)) {
             continue;
         }
         close_blob((pysqlite_Blob *)blob);

Diff for: Modules/_sqlite/connection.c

@@ -38,7 +38,7 @@
 #include "pycore_modsupport.h"    // _PyArg_NoKeywords()
 #include "pycore_pyerrors.h"      // _PyErr_ChainExceptions1()
 #include "pycore_pylifecycle.h"   // _Py_IsInterpreterFinalizing()
-#include "pycore_weakref.h"       // _PyWeakref_IS_DEAD()
+#include "pycore_weakref.h"
 
 #include <stdbool.h>
 
@@ -1065,7 +1065,7 @@ static void _pysqlite_drop_unused_cursor_references(pysqlite_Connection* self)
 
     for (Py_ssize_t i = 0; i < PyList_Size(self->cursors); i++) {
         PyObject* weakref = PyList_GetItem(self->cursors, i);
-        if (_PyWeakref_IS_DEAD(weakref)) {
+        if (_PyWeakref_IsDead(weakref)) {
             continue;
         }
         if (PyList_Append(new_list, weakref) != 0) {

Diff for: Modules/_ssl.c

@@ -29,7 +29,6 @@
 #include "pycore_fileutils.h"     // _PyIsSelectable_fd()
 #include "pycore_pyerrors.h"      // _PyErr_ChainExceptions1()
 #include "pycore_time.h"          // _PyDeadline_Init()
-#include "pycore_weakref.h"       // _PyWeakref_GET_REF()
 
 /* Include symbols from _socket module */
 #include "socketmodule.h"
@@ -392,8 +391,8 @@ typedef enum {
 // Return a borrowed reference.
 static inline PySocketSockObject* GET_SOCKET(PySSLSocket *obj) {
     if (obj->Socket) {
-        PyObject *sock = _PyWeakref_GET_REF(obj->Socket);
-        if (sock != NULL) {
+        PyObject *sock;
+        if (PyWeakref_GetRef(obj->Socket, &sock)) {
             // GET_SOCKET() returns a borrowed reference
             Py_DECREF(sock);
         }
@@ -2205,8 +2204,8 @@ PySSL_get_owner(PySSLSocket *self, void *c)
     if (self->owner == NULL) {
         Py_RETURN_NONE;
     }
-    PyObject *owner = _PyWeakref_GET_REF(self->owner);
-    if (owner == NULL) {
+    PyObject *owner;
+    if (!PyWeakref_GetRef(self->owner, &owner)) {
         Py_RETURN_NONE;
     }
     return owner;
@@ -4433,9 +4432,9 @@ _servername_callback(SSL *s, int *al, void *args)
      * will be passed. If both do not exist only then the C-level object is
      * passed. */
     if (ssl->owner)
-        ssl_socket = _PyWeakref_GET_REF(ssl->owner);
+        PyWeakref_GetRef(ssl->owner, &ssl_socket);
     else if (ssl->Socket)
-        ssl_socket = _PyWeakref_GET_REF(ssl->Socket);
+        PyWeakref_GetRef(ssl->Socket, &ssl_socket);
     else
         ssl_socket = Py_NewRef(ssl);
 

Diff for: Modules/_ssl/debughelpers.c

@@ -28,12 +28,12 @@ _PySSL_msg_callback(int write_p, int version, int content_type,
 
     PyObject *ssl_socket;  /* ssl.SSLSocket or ssl.SSLObject */
     if (ssl_obj->owner)
-        ssl_socket = _PyWeakref_GET_REF(ssl_obj->owner);
+        PyWeakref_GetRef(ssl_obj->owner, &ssl_socket);
     else if (ssl_obj->Socket)
-        ssl_socket = _PyWeakref_GET_REF(ssl_obj->Socket);
+        PyWeakref_GetRef(ssl_obj->Socket, &ssl_socket);
     else
         ssl_socket = (PyObject *)Py_NewRef(ssl_obj);
-    assert(ssl_socket != NULL);  // _PyWeakref_GET_REF() can return NULL
+    assert(ssl_socket != NULL);  // PyWeakref_GetRef() can return NULL
 
     /* assume that OpenSSL verifies all payload and buf len is of sufficient
        length */