gh-111926: Make weakrefs thread-safe in free-threaded builds

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)

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

Include/internal/pycore_object.h

@@ -427,7 +427,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;
@@ -453,7 +453,7 @@ _Py_XGetRef(PyObject **ptr)
  if (value == NULL){
  return value;
  }
- if (_Py_TryIncref(ptr, value)){
+ if (_Py_TryIncrefCompare(ptr, value)){
  return value;
  }
  }
@@ -468,7 +468,7 @@ _Py_TryXGetRef(PyObject **ptr)
  if (value == NULL){
  return value;
  }
- if (_Py_TryIncref(ptr, value)){
+ if (_Py_TryIncrefCompare(ptr, value)){
  return value;
  }
  return NULL;
@@ -507,8 +507,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 *);

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

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 */
-

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)

Lib/test/test_weakref.py

@@ -1865,6 +1865,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.

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);

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){

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);
 

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 */