pymem: remove uses of _PyMem_SetDefaultAllocator during finalization

There is a data race (even with the GIL) between the calls to
_PyMem_SetDefaultAllocator and calls to PyMem_RawMalloc from other
threads.

Author: colesbury

Include/internal/pycore_pymem.h

@@ -92,6 +92,12 @@ PyAPI_FUNC(int) _PyMem_GetAllocatorName(
    PYMEM_ALLOCATOR_NOT_SET does nothing. */
 PyAPI_FUNC(int) _PyMem_SetupAllocators(PyMemAllocatorName allocator);
 
+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 char * _PyMem_DefaultRawStrdup(const char *);
+extern wchar_t * _PyMem_DefaultRawWcsdup(const wchar_t *str);
 
 #ifdef __cplusplus
 }

Modules/getpath.c

@@ -6,7 +6,6 @@
 #include "pycore_initconfig.h"
 #include "pycore_fileutils.h"
 #include "pycore_pathconfig.h"
-#include "pycore_pymem.h"         // _PyMem_SetDefaultAllocator()
 #include <wchar.h>
 
 #ifdef MS_WINDOWS

Objects/obmalloc.c

@@ -73,6 +73,82 @@ _PyMem_RawFree(void *Py_UNUSED(ctx), void *ptr)
     free(ptr);
 }
 
+/* runtime default raw allocator */
+
+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
+}
+
+char *
+_PyMem_DefaultRawStrdup(const char *str)
+{
+    assert(str != NULL);
+    size_t size = strlen(str) + 1;
+    char *copy = _PyMem_DefaultRawMalloc(size);
+    if (copy == NULL) {
+        return NULL;
+    }
+    memcpy(copy, str, size);
+    return copy;
+}
+
+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;
+}
+
 #define MALLOC_ALLOC {NULL, _PyMem_RawMalloc, _PyMem_RawCalloc, _PyMem_RawRealloc, _PyMem_RawFree}
 
 /* the default object allocator */

Python/import.c

@@ -9,7 +9,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_Audit()
 #include "marshal.h"              // PyMarshal_ReadObjectFromString()
@@ -228,30 +228,23 @@ _PyImport_Init(void)
     if (_PyRuntime.imports.inittab != NULL) {
         return _PyStatus_ERR("global import state already initialized");
     }
-    PyStatus status = _PyStatus_OK();
 
     size_t size;
     for (size = 0; PyImport_Inittab[size].name != NULL; size++)
         ;
     size++;
 
-    /* Force default raw memory allocator to get a known allocator to be able
+    /* Use 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);
 
     /* Make the copy. */
-    struct _inittab *copied = PyMem_RawMalloc(size * sizeof(struct _inittab));
+    struct _inittab *copied = _PyMem_DefaultRawMalloc(size * sizeof(struct _inittab));
     if (copied == NULL) {
-        status = PyStatus_NoMemory();
-        goto done;
+        return PyStatus_NoMemory();
     }
     memcpy(copied, PyImport_Inittab, size * sizeof(struct _inittab));
     _PyRuntime.imports.inittab = copied;
-
-done:
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-    return status;
+    return _PyStatus_OK();
 }
 
 static inline void _extensions_cache_clear(void);
@@ -266,32 +259,22 @@ _PyImport_Fini(void)
     }
 
     /* Use the same memory allocator as _PyImport_Init(). */
-    PyMemAllocatorEx old_alloc;
-    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-
     /* Free memory allocated by _PyImport_Init() */
     struct _inittab *inittab = _PyRuntime.imports.inittab;
     _PyRuntime.imports.inittab = NULL;
-    PyMem_RawFree(inittab);
-
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
+    _PyMem_DefaultRawFree(inittab);
 }
 
 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);
 }
 
 /* Helper for sys */
@@ -2657,14 +2640,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;
@@ -2678,9 +2658,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;
 }