Switch to a sharded lock for linked list and object

python · colesbury · Apr 8, 2024 · Mar 14, 2024 · Mar 15, 2024 · Mar 15, 2024
commit 88078a686bf0e6ca509d0819289634efdf25da8f
diff --git a/Include/cpython/object.h b/Include/cpython/object.h
@@ -465,13 +465,6 @@ PyAPI_FUNC(void) _PyTrash_end(PyThreadState *tstate);
 /* Python 3.10 private API, invoked by the Py_TRASHCAN_BEGIN(). */
 PyAPI_FUNC(int) _PyTrash_cond(PyObject *op, destructor dealloc);
 
-#ifdef Py_GIL_DISABLED
-
-/* Private API, used by weakrefs in free-threaded builds */
-PyAPI_FUNC(int) _PyTrash_contains(PyObject *op);
-
-#endif
-
 #define Py_TRASHCAN_BEGIN_CONDITION(op, cond) \
     do { \
         PyThreadState *_tstate = NULL; \

diff --git a/Include/internal/pycore_critical_section.h b/Include/internal/pycore_critical_section.h
@@ -87,10 +87,12 @@ extern "C" {
 #define _Py_CRITICAL_SECTION_MASK           0x3
 
 #ifdef Py_GIL_DISABLED
-# define Py_BEGIN_CRITICAL_SECTION(op)                                  \
+# define Py_BEGIN_CRITICAL_SECTION_MU(mu)                               \
     {                                                                   \
         _PyCriticalSection _cs;                                         \
-        _PyCriticalSection_Begin(&_cs, &_PyObject_CAST(op)->ob_mutex)
+        _PyCriticalSection_Begin(&_cs, &mu)
+
+# define Py_BEGIN_CRITICAL_SECTION(op) Py_BEGIN_CRITICAL_SECTION_MU(_PyObject_CAST(op)->ob_mutex)
 
 # define Py_END_CRITICAL_SECTION()                                      \
         _PyCriticalSection_End(&_cs);                                   \
@@ -105,10 +107,13 @@ extern "C" {
         _PyCriticalSection_XEnd(&_cs_opt);                              \
     }
 
-# define Py_BEGIN_CRITICAL_SECTION2(a, b)                               \
+# define Py_BEGIN_CRITICAL_SECTION2_MU(a, b)                            \
     {                                                                   \
         _PyCriticalSection2 _cs2;                                       \
-        _PyCriticalSection2_Begin(&_cs2, &_PyObject_CAST(a)->ob_mutex, &_PyObject_CAST(b)->ob_mutex)
+        _PyCriticalSection2_Begin(&_cs2, &a, &b)
+
+# define Py_BEGIN_CRITICAL_SECTION2(a, b)                               \
+    Py_BEGIN_CRITICAL_SECTION2_MU(_PyObject_CAST(a)->ob_mutex, _PyObject_CAST(b)->ob_mutex)
 
 # define Py_END_CRITICAL_SECTION2()                                     \
         _PyCriticalSection2_End(&_cs2);                                 \
@@ -138,10 +143,12 @@ extern "C" {
 
 #else  /* !Py_GIL_DISABLED */
 // The critical section APIs are no-ops with the GIL.
+# define Py_BEGIN_CRITICAL_SECTION_MU(mu)
 # define Py_BEGIN_CRITICAL_SECTION(op)
 # define Py_END_CRITICAL_SECTION()
 # define Py_XBEGIN_CRITICAL_SECTION(op)
 # define Py_XEND_CRITICAL_SECTION()
+# define Py_BEGIN_CRITICAL_SECTION2_MU(a, b)
 # define Py_BEGIN_CRITICAL_SECTION2(a, b)
 # define Py_END_CRITICAL_SECTION2()
 # define _Py_CRITICAL_SECTION_ASSERT_MUTEX_LOCKED(mutex)

@@ -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

@@ -8,16 +8,25 @@ extern "C" {
 #  error "this header requires Py_BUILD_CORE define"
 #endif
 
+#include "pycore_critical_section.h" // Py_BEGIN_CRITICAL_SECTION()
 #include "pycore_lock.h"
 #include "pycore_object.h"           // _Py_REF_IS_MERGED()
 
 #ifdef Py_GIL_DISABLED
+
 typedef struct _PyWeakRefClearState {
+    // Protects the weakref's wr_object and wr_callback.
+    // Protects the cleared flag below.
     PyMutex mutex;
-    _PyOnceFlag once;
+
+    // Set if the weakref has been cleared.
     int cleared;
+
     Py_ssize_t refcount;
 } _PyWeakRefClearState;
+
+#define WEAKREF_LIST_LOCK(obj) _PyInterpreterState_GET()->weakref_locks[((uintptr_t) obj) % NUM_WEAKREF_LIST_LOCKS]
+
 #endif
 
 static inline int _is_dead(PyObject *obj)
@@ -36,23 +45,12 @@ static inline int _is_dead(PyObject *obj)
 #endif
 }
 
-#if defined(Py_GIL_DISABLED)
-#    define LOCK_WR_OBJECT(weakref) PyMutex_Lock(&(weakref)->clear_state->mutex);
-#    define UNLOCK_WR_OBJECT(weakref) PyMutex_Unlock(&(weakref)->clear_state->mutex);
-#else
-#    define LOCK_WR_OBJECT(weakref)
-#    define UNLOCK_WR_OBJECT(weakref)
-#endif
-
-// NB: In free-threaded builds nothing between the LOCK/UNLOCK calls below can
-// suspend.
-
 static inline PyObject* _PyWeakref_GET_REF(PyObject *ref_obj)
 {
     assert(PyWeakref_Check(ref_obj));
     PyObject *ret = NULL;
     PyWeakReference *ref = _Py_CAST(PyWeakReference*, ref_obj);
-    LOCK_WR_OBJECT(ref)
+    Py_BEGIN_CRITICAL_SECTION_MU(ref->clear_state->mutex);
     PyObject *obj = ref->wr_object;
 
     if (obj == Py_None) {
@@ -68,7 +66,7 @@ static inline PyObject* _PyWeakref_GET_REF(PyObject *ref_obj)
 #endif
     ret = Py_NewRef(obj);
 end:
-    UNLOCK_WR_OBJECT(ref);
+    Py_END_CRITICAL_SECTION();
     return ret;
 }
 
@@ -77,7 +75,7 @@ static inline int _PyWeakref_IS_DEAD(PyObject *ref_obj)
     assert(PyWeakref_Check(ref_obj));
     int ret = 0;
     PyWeakReference *ref = _Py_CAST(PyWeakReference*, ref_obj);
-    LOCK_WR_OBJECT(ref);
+    Py_BEGIN_CRITICAL_SECTION_MU(ref->clear_state->mutex);
     PyObject *obj = ref->wr_object;
     if (obj == Py_None) {
         // clear_weakref() was called
@@ -87,14 +85,16 @@ static inline int _PyWeakref_IS_DEAD(PyObject *ref_obj)
         // See _PyWeakref_GET_REF() for the rationale of this test
         ret = _is_dead(obj);
     }
-    UNLOCK_WR_OBJECT(ref);
+    Py_END_CRITICAL_SECTION();
     return ret;
 }
 
-// NB: In free-threaded builds the referenced object must be live for the
-// duration of the call.
+// NB: In free-threaded builds the weakref list lock for the referenced object
+// must be held around calls to this function.
 extern Py_ssize_t _PyWeakref_GetWeakrefCount(PyWeakReference *head);
 
+extern Py_ssize_t _PyWeakref_GetWeakrefCountThreadsafe(PyObject *obj);
+
 // Clear all the weak references to obj but leave their callbacks uncalled and
 // intact.
 extern void _PyWeakref_ClearWeakRefsExceptCallbacks(PyObject *obj);

@@ -30,9 +30,7 @@ _weakref_getweakrefcount_impl(PyObject *module, PyObject *object)
     if (!_PyType_SUPPORTS_WEAKREFS(Py_TYPE(object))) {
         return 0;
     }
-    PyWeakReference **list = GET_WEAKREFS_LISTPTR(object);
-    Py_ssize_t count = _PyWeakref_GetWeakrefCount(*list);
-    return count;
+    return _PyWeakref_GetWeakrefCountThreadsafe(object);
 }
 
 
@@ -94,7 +92,7 @@ _weakref_getweakrefs_impl(PyObject *module, PyObject *object)
     }
 
     PyWeakReference **list = GET_WEAKREFS_LISTPTR(object);
-    Py_ssize_t count = _PyWeakref_GetWeakrefCount(*list);
+    Py_ssize_t count = _PyWeakref_GetWeakrefCountThreadsafe(object);
 
     PyObject *result = PyList_New(count);
     if (result == NULL) {
@@ -103,7 +101,7 @@ _weakref_getweakrefs_impl(PyObject *module, PyObject *object)
 
 #ifdef Py_GIL_DISABLED
     Py_ssize_t num_added = 0;
-    Py_BEGIN_CRITICAL_SECTION(object);
+    Py_BEGIN_CRITICAL_SECTION_MU(WEAKREF_LIST_LOCK(object));
     PyWeakReference *current = *list;
     // Weakrefs may be added or removed since the count was computed.
     while (num_added < count && current != NULL) {

diff --git a/Objects/object.c b/Objects/object.c
@@ -2835,25 +2835,6 @@ _PyTrash_cond(PyObject *op, destructor dealloc)
     return Py_TYPE(op)->tp_dealloc == dealloc;
 }
 
-#ifdef Py_GIL_DISABLED
-
-int
-_PyTrash_contains(PyObject *op)
-{
-     PyThreadState *tstate = _PyThreadState_GET();
-    struct _py_trashcan *trash = _PyTrash_get_state(tstate);
-    PyObject *cur = trash->delete_later;
-    while (cur) {
-        if (cur == op) {
-            return 1;
-        }
-        cur = (PyObject *) cur->ob_tid;
-    }
-    return 0;
-}
-
-#endif
-
 
 void _Py_NO_RETURN
 _PyObject_AssertFailed(PyObject *obj, const char *expr, const char *msg,

@@ -2188,15 +2188,7 @@ subtype_dealloc(PyObject *self)
            finalizers since they might rely on part of the object
            being finalized that has already been destroyed. */
         if (type->tp_weaklistoffset && !base->tp_weaklistoffset) {
-            /* Modeled after GET_WEAKREFS_LISTPTR().
-
-               This is never triggered for static types so we can avoid the
-               (slightly) more costly _PyObject_GET_WEAKREFS_LISTPTR(). */
-            PyWeakReference **list = \
-                _PyObject_GET_WEAKREFS_LISTPTR_FROM_OFFSET(self);
-            while (*list) {
-                _PyWeakref_ClearRef(*list);
-            }
+            _PyWeakref_ClearWeakRefsExceptCallbacks(self);
         }
     }