gh-132070: Fix PyObject_Realloc thread-safety in free threaded Python

The PyObject header reference count fields must be initialized using
atomic operations because they may be concurrently read by another
thread (e.g., from _Py_TryIncref).

Author: colesbury

Objects/obmalloc.c

@@ -307,8 +307,42 @@ _PyObject_MiRealloc(void *ctx, void *ptr, size_t nbytes)
 {
 #ifdef Py_GIL_DISABLED
     _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET();
+    // Implement our own realloc logic so that we can copy PyObject header
+    // in a thread-safe way.
+    size_t size = mi_usable_size(ptr);
+    if (nbytes <= size && nbytes >= (size / 2) && nbytes > 0) {
+        return ptr;
+    }
+
     mi_heap_t *heap = tstate->mimalloc.current_object_heap;
-    return mi_heap_realloc(heap, ptr, nbytes);
+    void* newp = mi_heap_malloc(heap, nbytes);
+    if (newp == NULL) {
+        return NULL;
+    }
+
+    // Free threaded Python allows safe access to the PyObject reference count
+    // fields for a period of time after the object is freed (see InternalDocs/qsbr.md).
+    // These fields are typically initialized by PyObject_Init() using relaxed
+    // atomic stores. We need to copy these fields in a thread-safe way here.
+    // We use the "debug_offset" to determine how many bytes to copy -- it
+    // includes the PyObject header and plus any extra pre-headers.
+    size_t offset = heap->debug_offset;
+    assert(offset % sizeof(void*) == 0);
+
+    size_t copy_size = (size < nbytes ? size : nbytes);
+    if (copy_size >= offset) {
+        for (size_t i = 0; i != offset; i += sizeof(void*)) {
+            void *word;
+            memcpy(&word, (char*)ptr + i, sizeof(void*));
+            _Py_atomic_store_ptr_relaxed((void**)((char*)newp + i), word);
+        }
+        _mi_memcpy((char*)newp + offset, (char*)ptr + offset, copy_size - offset);
+    }
+    else {
+        _mi_memcpy(newp, ptr, copy_size);
+    }
+    mi_free(ptr);
+    return newp;
 #else
     return mi_realloc(ptr, nbytes);
 #endif

Tools/tsan/suppressions_free_threading.txt

@@ -16,7 +16,3 @@ race_top:_PyObject_TryGetInstanceAttribute
 
 # https://gist.github.com/mpage/6962e8870606cfc960e159b407a0cb40
 thread:pthread_create
-
-# PyObject_Realloc internally does memcpy which isn't atomic so can race
-# with non-locking reads. See #132070
-race:PyObject_Realloc