pythongh-111924: Fix data races when swapping allocators (pythongh-130287)

CPython current temporarily changes `PYMEM_DOMAIN_RAW` to the default
allocator during initialization and shutdown. The motivation is to
ensure that core runtime structures are allocated and freed using the
same allocator. However, modifying the current allocator changes global
state and is not thread-safe even with the GIL. Other threads may be
allocating or freeing objects use PYMEM_DOMAIN_RAW; they are not
required to hold the GIL to call PyMem_RawMalloc/PyMem_RawFree.

This adds new internal-only functions like `_PyMem_DefaultRawMalloc`
that aren't affected by calls to `PyMem_SetAllocator()`, so they're
appropriate for Python runtime initialization and finalization. Use
these calls in places where we previously swapped to the default raw
allocator.

Author: colesbury

Include/internal/pycore_pymem.h

@@ -55,13 +55,6 @@ struct _Py_mem_interp_free_queue {
     struct llist_node head;  // queue of _mem_work_chunk items
 };
 
-/* Set the memory allocator of the specified domain to the default.
-   Save the old allocator into *old_alloc if it's non-NULL.
-   Return on success, or return -1 if the domain is unknown. */
-extern int _PyMem_SetDefaultAllocator(
-    PyMemAllocatorDomain domain,
-    PyMemAllocatorEx *old_alloc);
-
 /* Special bytes broadcast into debug memory blocks at appropriate times.
    Strings of these are unlikely to be valid addresses, floats, ints or
    7-bit ASCII.
@@ -113,6 +106,13 @@ extern int _PyMem_GetAllocatorName(
    PYMEM_ALLOCATOR_NOT_SET does nothing. */
 extern int _PyMem_SetupAllocators(PyMemAllocatorName allocator);
 
+// Default raw memory allocator that is not affected by PyMem_SetAllocator()
+extern void *_PyMem_DefaultRawMalloc(size_t);
+extern void *_PyMem_DefaultRawCalloc(size_t, size_t);
+extern void *_PyMem_DefaultRawRealloc(void *, size_t);
+extern void _PyMem_DefaultRawFree(void *);
+extern wchar_t *_PyMem_DefaultRawWcsdup(const wchar_t *str);
+
 /* Is the debug allocator enabled? */
 extern int _PyMem_DebugEnabled(void);
 
@@ -132,7 +132,6 @@ static inline void _PyObject_XDecRefDelayed(PyObject *obj)
 // Periodically process delayed free requests.
 extern void _PyMem_ProcessDelayed(PyThreadState *tstate);
 
-
 // Periodically process delayed free requests when the world is stopped.
 // Notify of any objects whic should be freeed.
 typedef void (*delayed_dealloc_cb)(PyObject *, void *);

Objects/obmalloc.c

@@ -344,6 +344,68 @@ void _PyMem_DebugFree(void *ctx, void *p);
 #define PYDBGOBJ_ALLOC \
     {&_PyRuntime.allocators.debug.obj, _PyMem_DebugMalloc, _PyMem_DebugCalloc, _PyMem_DebugRealloc, _PyMem_DebugFree}
 
+/* default raw allocator (not swappable) */
+
+void *
+_PyMem_DefaultRawMalloc(size_t size)
+{
+#ifdef Py_DEBUG
+    return _PyMem_DebugRawMalloc(&_PyRuntime.allocators.debug.raw, size);
+#else
+    return _PyMem_RawMalloc(NULL, size);
+#endif
+}
+
+void *
+_PyMem_DefaultRawCalloc(size_t nelem, size_t elsize)
+{
+#ifdef Py_DEBUG
+    return _PyMem_DebugRawCalloc(&_PyRuntime.allocators.debug.raw, nelem, elsize);
+#else
+    return _PyMem_RawCalloc(NULL, nelem, elsize);
+#endif
+}
+
+void *
+_PyMem_DefaultRawRealloc(void *ptr, size_t size)
+{
+#ifdef Py_DEBUG
+    return _PyMem_DebugRawRealloc(&_PyRuntime.allocators.debug.raw, ptr, size);
+#else
+    return _PyMem_RawRealloc(NULL, ptr, size);
+#endif
+}
+
+void
+_PyMem_DefaultRawFree(void *ptr)
+{
+#ifdef Py_DEBUG
+    _PyMem_DebugRawFree(&_PyRuntime.allocators.debug.raw, ptr);
+#else
+    _PyMem_RawFree(NULL, ptr);
+#endif
+}
+
+wchar_t*
+_PyMem_DefaultRawWcsdup(const wchar_t *str)
+{
+    assert(str != NULL);
+
+    size_t len = wcslen(str);
+    if (len > (size_t)PY_SSIZE_T_MAX / sizeof(wchar_t) - 1) {
+        return NULL;
+    }
+
+    size_t size = (len + 1) * sizeof(wchar_t);
+    wchar_t *str2 = _PyMem_DefaultRawMalloc(size);
+    if (str2 == NULL) {
+        return NULL;
+    }
+
+    memcpy(str2, str, size);
+    return str2;
+}
+
 /* the low-level virtual memory allocator */
 
 #ifdef WITH_PYMALLOC
@@ -492,17 +554,6 @@ static const int pydebug = 1;
 static const int pydebug = 0;
 #endif
 
-int
-_PyMem_SetDefaultAllocator(PyMemAllocatorDomain domain,
-                           PyMemAllocatorEx *old_alloc)
-{
-    PyMutex_Lock(&ALLOCATORS_MUTEX);
-    int res = set_default_allocator_unlocked(domain, pydebug, old_alloc);
-    PyMutex_Unlock(&ALLOCATORS_MUTEX);
-    return res;
-}
-
-
 int
 _PyMem_GetAllocatorName(const char *name, PyMemAllocatorName *allocator)
 {

Python/import.c

@@ -13,7 +13,7 @@
 #include "pycore_pyerrors.h"      // _PyErr_SetString()
 #include "pycore_pyhash.h"        // _Py_KeyedHash()
 #include "pycore_pylifecycle.h"
-#include "pycore_pymem.h"         // _PyMem_SetDefaultAllocator()
+#include "pycore_pymem.h"         // _PyMem_DefaultRawFree()
 #include "pycore_pystate.h"       // _PyInterpreterState_GET()
 #include "pycore_sysmodule.h"     // _PySys_ClearAttrString()
 #include "pycore_time.h"          // _PyTime_AsMicroseconds()
@@ -2387,14 +2387,11 @@ PyImport_ExtendInittab(struct _inittab *newtab)
 
     /* Force default raw memory allocator to get a known allocator to be able
        to release the memory in _PyImport_Fini2() */
-    PyMemAllocatorEx old_alloc;
-    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-
     /* Allocate new memory for the combined table */
     p = NULL;
     if (i + n <= SIZE_MAX / sizeof(struct _inittab) - 1) {
         size_t size = sizeof(struct _inittab) * (i + n + 1);
-        p = PyMem_RawRealloc(inittab_copy, size);
+        p = _PyMem_DefaultRawRealloc(inittab_copy, size);
     }
     if (p == NULL) {
         res = -1;
@@ -2408,9 +2405,7 @@ PyImport_ExtendInittab(struct _inittab *newtab)
     }
     memcpy(p + i, newtab, (n + 1) * sizeof(struct _inittab));
     PyImport_Inittab = inittab_copy = p;
-
 done:
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
     return res;
 }
 
@@ -2445,7 +2440,7 @@ init_builtin_modules_table(void)
     size++;
 
     /* Make the copy. */
-    struct _inittab *copied = PyMem_RawMalloc(size * sizeof(struct _inittab));
+    struct _inittab *copied = _PyMem_DefaultRawMalloc(size * sizeof(struct _inittab));
     if (copied == NULL) {
         return -1;
     }
@@ -2459,7 +2454,7 @@ fini_builtin_modules_table(void)
 {
     struct _inittab *inittab = INITTAB;
     INITTAB = NULL;
-    PyMem_RawFree(inittab);
+    _PyMem_DefaultRawFree(inittab);
 }
 
 PyObject *
@@ -3977,22 +3972,10 @@ _PyImport_Init(void)
     if (INITTAB != NULL) {
         return _PyStatus_ERR("global import state already initialized");
     }
-
-    PyStatus status = _PyStatus_OK();
-
-    /* Force default raw memory allocator to get a known allocator to be able
-       to release the memory in _PyImport_Fini() */
-    PyMemAllocatorEx old_alloc;
-    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-
     if (init_builtin_modules_table() != 0) {
-        status = PyStatus_NoMemory();
-        goto done;
+        return PyStatus_NoMemory();
     }
-
-done:
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-    return status;
+    return _PyStatus_OK();
 }
 
 void
@@ -4003,31 +3986,19 @@ _PyImport_Fini(void)
     // ever dlclose() the module files?
     _extensions_cache_clear_all();
 
-    /* Use the same memory allocator as _PyImport_Init(). */
-    PyMemAllocatorEx old_alloc;
-    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-
     /* Free memory allocated by _PyImport_Init() */
     fini_builtin_modules_table();
-
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
 }
 
 void
 _PyImport_Fini2(void)
 {
-    /* Use the same memory allocator than PyImport_ExtendInittab(). */
-    PyMemAllocatorEx old_alloc;
-    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-
     // Reset PyImport_Inittab
     PyImport_Inittab = _PyImport_Inittab;
 
     /* Free memory allocated by PyImport_ExtendInittab() */
-    PyMem_RawFree(inittab_copy);
+    _PyMem_DefaultRawFree(inittab_copy);
     inittab_copy = NULL;
-
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
 }
 
 

Python/initconfig.c

@@ -7,7 +7,7 @@
 #include "pycore_pathconfig.h"    // _Py_path_config
 #include "pycore_pyerrors.h"      // _PyErr_GetRaisedException()
 #include "pycore_pylifecycle.h"   // _Py_PreInitializeFromConfig()
-#include "pycore_pymem.h"         // _PyMem_SetDefaultAllocator()
+#include "pycore_pymem.h"         // _PyMem_DefaultRawMalloc()
 #include "pycore_pystate.h"       // _PyThreadState_GET()
 #include "pycore_pystats.h"       // _Py_StatsOn()
 #include "pycore_sysmodule.h"     // _PySys_SetIntMaxStrDigits()
@@ -619,53 +619,87 @@ _PyWideStringList_CheckConsistency(const PyWideStringList *list)
 #endif   /* Py_DEBUG */
 
 
-void
-_PyWideStringList_Clear(PyWideStringList *list)
+static void
+_PyWideStringList_ClearEx(PyWideStringList *list,
+                          bool use_default_allocator)
 {
     assert(_PyWideStringList_CheckConsistency(list));
     for (Py_ssize_t i=0; i < list->length; i++) {
-        PyMem_RawFree(list->items[i]);
+        if (use_default_allocator) {
+            _PyMem_DefaultRawFree(list->items[i]);
+        }
+        else {
+            PyMem_RawFree(list->items[i]);
+        }
+    }
+    if (use_default_allocator) {
+        _PyMem_DefaultRawFree(list->items);
+    }
+    else {
+        PyMem_RawFree(list->items);
     }
-    PyMem_RawFree(list->items);
     list->length = 0;
     list->items = NULL;
 }
 
+void
+_PyWideStringList_Clear(PyWideStringList *list)
+{
+    _PyWideStringList_ClearEx(list, false);
+}
 
-int
-_PyWideStringList_Copy(PyWideStringList *list, const PyWideStringList *list2)
+static int
+_PyWideStringList_CopyEx(PyWideStringList *list,
+                         const PyWideStringList *list2,
+                         bool use_default_allocator)
 {
     assert(_PyWideStringList_CheckConsistency(list));
     assert(_PyWideStringList_CheckConsistency(list2));
 
     if (list2->length == 0) {
-        _PyWideStringList_Clear(list);
+        _PyWideStringList_ClearEx(list, use_default_allocator);
         return 0;
     }
 
     PyWideStringList copy = _PyWideStringList_INIT;
 
     size_t size = list2->length * sizeof(list2->items[0]);
-    copy.items = PyMem_RawMalloc(size);
+    if (use_default_allocator) {
+        copy.items = _PyMem_DefaultRawMalloc(size);
+    }
+    else {
+        copy.items = PyMem_RawMalloc(size);
+    }
     if (copy.items == NULL) {
         return -1;
     }
 
     for (Py_ssize_t i=0; i < list2->length; i++) {
-        wchar_t *item = _PyMem_RawWcsdup(list2->items[i]);
+        wchar_t *item;
+        if (use_default_allocator) {
+            item = _PyMem_DefaultRawWcsdup(list2->items[i]);
+        }
+        else {
+            item = _PyMem_RawWcsdup(list2->items[i]);
+        }
         if (item == NULL) {
-            _PyWideStringList_Clear(&copy);
+            _PyWideStringList_ClearEx(&copy, use_default_allocator);
             return -1;
         }
         copy.items[i] = item;
         copy.length = i + 1;
     }
 
-    _PyWideStringList_Clear(list);
+    _PyWideStringList_ClearEx(list, use_default_allocator);
     *list = copy;
     return 0;
 }
 
+int
+_PyWideStringList_Copy(PyWideStringList *list, const PyWideStringList *list2)
+{
+    return _PyWideStringList_CopyEx(list, list2, false);
+}
 
 PyStatus
 PyWideStringList_Insert(PyWideStringList *list,
@@ -789,30 +823,18 @@ _PyWideStringList_AsTuple(const PyWideStringList *list)
 void
 _Py_ClearArgcArgv(void)
 {
-    PyMemAllocatorEx old_alloc;
-    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-
-    _PyWideStringList_Clear(&_PyRuntime.orig_argv);
-
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
+    _PyWideStringList_ClearEx(&_PyRuntime.orig_argv, true);
 }
 
 
 static int
 _Py_SetArgcArgv(Py_ssize_t argc, wchar_t * const *argv)
 {
     const PyWideStringList argv_list = {.length = argc, .items = (wchar_t **)argv};
-    int res;
-
-    PyMemAllocatorEx old_alloc;
-    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
 
     // XXX _PyRuntime.orig_argv only gets cleared by Py_Main(),
     // so it currently leaks for embedders.
-    res = _PyWideStringList_Copy(&_PyRuntime.orig_argv, &argv_list);
-
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-    return res;
+    return _PyWideStringList_CopyEx(&_PyRuntime.orig_argv, &argv_list, true);
 }