gh-113743: Make the MRO cache thread-safe in free-threaded builds (#113930)

Makes _PyType_Lookup thread safe, including:
    Thread safety of the underlying cache.
    Make mutation of mro and type members thread safe
    Also _PyType_GetMRO and _PyType_GetBases are currently returning borrowed references which aren't safe.

2 files changed, 74 insertions, 0 deletions

diff --git a/Python/lock.c b/Python/lock.c
index f0ff117..bf01436 100644
--- a/Python/lock.c
+++ b/Python/lock.c
@@ -459,3 +459,74 @@ _PyRWMutex_Unlock(_PyRWMutex *rwmutex)
         _PyParkingLot_UnparkAll(&rwmutex->bits);
     }
 }
+
+#define SEQLOCK_IS_UPDATING(sequence) (sequence & 0x01)
+
+void _PySeqLock_LockWrite(_PySeqLock *seqlock)
+{
+    // lock the entry by setting by moving to an odd sequence number
+    uint32_t prev = _Py_atomic_load_uint32_relaxed(&seqlock->sequence);
+    while (1) {
+        if (SEQLOCK_IS_UPDATING(prev)) {
+            // Someone else is currently updating the cache
+            _Py_yield();
+            prev = _Py_atomic_load_uint32_relaxed(&seqlock->sequence);
+        }
+        else if (_Py_atomic_compare_exchange_uint32(&seqlock->sequence, &prev, prev + 1)) {
+            // We've locked the cache
+            break;
+        }
+        else {
+            _Py_yield();
+        }
+    }
+}
+
+void _PySeqLock_AbandonWrite(_PySeqLock *seqlock)
+{
+    uint32_t new_seq = seqlock->sequence - 1;
+    assert(!SEQLOCK_IS_UPDATING(new_seq));
+    _Py_atomic_store_uint32(&seqlock->sequence, new_seq);
+}
+
+void _PySeqLock_UnlockWrite(_PySeqLock *seqlock)
+{
+    uint32_t new_seq = seqlock->sequence + 1;
+    assert(!SEQLOCK_IS_UPDATING(new_seq));
+    _Py_atomic_store_uint32(&seqlock->sequence, new_seq);
+}
+
+uint32_t _PySeqLock_BeginRead(_PySeqLock *seqlock)
+{
+    uint32_t sequence = _Py_atomic_load_uint32_acquire(&seqlock->sequence);
+    while (SEQLOCK_IS_UPDATING(sequence)) {
+        _Py_yield();
+        sequence = _Py_atomic_load_uint32_acquire(&seqlock->sequence);
+    }
+
+    return sequence;
+}
+
+uint32_t _PySeqLock_EndRead(_PySeqLock *seqlock, uint32_t previous)
+{
+    // Synchronize again and validate that the entry hasn't been updated
+    // while we were readying the values.
+     if (_Py_atomic_load_uint32_acquire(&seqlock->sequence) == previous) {
+        return 1;
+     }
+
+     _Py_yield();
+     return 0;
+}
+
+uint32_t _PySeqLock_AfterFork(_PySeqLock *seqlock)
+{
+    // Synchronize again and validate that the entry hasn't been updated
+    // while we were readying the values.
+     if (SEQLOCK_IS_UPDATING(seqlock->sequence)) {
+        seqlock->sequence = 0;
+        return 1;
+     }
+
+     return 0;
+}
diff --git a/Python/pystate.c b/Python/pystate.c
index 08ec586..b1d1a08 100644
--- a/Python/pystate.c
+++ b/Python/pystate.c
@@ -395,6 +395,7 @@ _Py_COMP_DIAG_POP
         &(runtime)->atexit.mutex, \
         &(runtime)->audit_hooks.mutex, \
         &(runtime)->allocators.mutex, \
+        &(runtime)->types.type_mutex, \
     }
 
 static void
@@ -499,6 +500,8 @@ _PyRuntimeState_ReInitThreads(_PyRuntimeState *runtime)
         _PyMutex_at_fork_reinit(locks[i]);
     }
 
+    _PyTypes_AfterFork();
+
     /* bpo-42540: id_mutex is freed by _PyInterpreterState_Delete, which does
      * not force the default allocator. */
     if (_PyThread_at_fork_reinit(&runtime->interpreters.main->id_mutex) < 0) {