bpo-32030: Rework memory allocators (#4625)

* Fix _PyMem_SetupAllocators("debug"): always restore allocators to
  the defaults, rather than only caling _PyMem_SetupDebugHooks().
* Add _PyMem_SetDefaultAllocator() helper to set the "default"
  allocator.
* Add _PyMem_GetAllocatorsName(): get the name of the allocators
* main() now uses debug hooks on memory allocators if Py_DEBUG is
  defined, rather than calling directly malloc()
* Document default memory allocators in C API documentation
* _Py_InitializeCore() now fails with a fatal user error if
  PYTHONMALLOC value is an unknown memory allocator, instead of
  failing with a fatal internal error.
* Add new tests on the PYTHONMALLOC environment variable
* Add support.with_pymalloc()
* Add the _testcapi.WITH_PYMALLOC constant and expose it as
   support.with_pymalloc().
* sysconfig.get_config_var('WITH_PYMALLOC') doesn't work on Windows, so
   replace it with support.with_pymalloc().
* pythoninfo: add _testcapi collector for pymem

Doc/c-api/memory.rst

@@ -100,9 +100,10 @@ The following function sets are wrappers to the system allocator. These
 functions are thread-safe, the :term:`GIL <global interpreter lock>` does not
 need to be held.
 
-The default raw memory block allocator uses the following functions:
-:c:func:`malloc`, :c:func:`calloc`, :c:func:`realloc` and :c:func:`free`; call
-``malloc(1)`` (or ``calloc(1, 1)``) when requesting zero bytes.
+The :ref:`default raw memory allocator <default-memory-allocators>` uses
+the following functions: :c:func:`malloc`, :c:func:`calloc`, :c:func:`realloc`
+and :c:func:`free`; call ``malloc(1)`` (or ``calloc(1, 1)``) when requesting
+zero bytes.
 
 .. versionadded:: 3.4
 
@@ -165,7 +166,8 @@ The following function sets, modeled after the ANSI C standard, but specifying
 behavior when requesting zero bytes, are available for allocating and releasing
 memory from the Python heap.
 
-By default, these functions use :ref:`pymalloc memory allocator <pymalloc>`.
+The :ref:`default memory allocator <default-memory-allocators>` uses the
+:ref:`pymalloc memory allocator <pymalloc>`.
 
 .. warning::
 
@@ -270,7 +272,8 @@ The following function sets, modeled after the ANSI C standard, but specifying
 behavior when requesting zero bytes, are available for allocating and releasing
 memory from the Python heap.
 
-By default, these functions use :ref:`pymalloc memory allocator <pymalloc>`.
+The :ref:`default object allocator <default-memory-allocators>` uses the
+:ref:`pymalloc memory allocator <pymalloc>`.
 
 .. warning::
 
@@ -326,6 +329,31 @@ By default, these functions use :ref:`pymalloc memory allocator <pymalloc>`.
    If *p* is *NULL*, no operation is performed.
 
 
+.. _default-memory-allocators:
+
+Default Memory Allocators
+=========================
+
+Default memory allocators:
+
+===============================  ====================  ==================  =====================  ====================
+Configuration                    Name                  PyMem_RawMalloc     PyMem_Malloc           PyObject_Malloc
+===============================  ====================  ==================  =====================  ====================
+Release build                    ``"pymalloc"``        ``malloc``          ``pymalloc``           ``pymalloc``
+Debug build                      ``"pymalloc_debug"``  ``malloc`` + debug  ``pymalloc`` + debug   ``pymalloc`` + debug
+Release build, without pymalloc  ``"malloc"``          ``malloc``          ``malloc``             ``malloc``
+Release build, without pymalloc  ``"malloc_debug"``    ``malloc`` + debug  ``malloc`` + debug     ``malloc`` + debug
+===============================  ====================  ==================  =====================  ====================
+
+Legend:
+
+* Name: value for :envvar:`PYTHONMALLOC` environment variable
+* ``malloc``: system allocators from the standard C library, C functions:
+  :c:func:`malloc`, :c:func:`calloc`, :c:func:`realloc` and :c:func:`free`
+* ``pymalloc``: :ref:`pymalloc memory allocator <pymalloc>`
+* "+ debug": with debug hooks installed by :c:func:`PyMem_SetupDebugHooks`
+
+
 Customize Memory Allocators
 ===========================
 
@@ -431,7 +459,8 @@ Customize Memory Allocators
    displayed if :mod:`tracemalloc` is tracing Python memory allocations and the
    memory block was traced.
 
-   These hooks are installed by default if Python is compiled in debug
+   These hooks are :ref:`installed by default <default-memory-allocators>` if
+   Python is compiled in debug
    mode. The :envvar:`PYTHONMALLOC` environment variable can be used to install
    debug hooks on a Python compiled in release mode.
 
@@ -453,9 +482,9 @@ to 512 bytes) with a short lifetime. It uses memory mappings called "arenas"
 with a fixed size of 256 KiB. It falls back to :c:func:`PyMem_RawMalloc` and
 :c:func:`PyMem_RawRealloc` for allocations larger than 512 bytes.
 
-*pymalloc* is the default allocator of the :c:data:`PYMEM_DOMAIN_MEM` (ex:
-:c:func:`PyMem_Malloc`) and :c:data:`PYMEM_DOMAIN_OBJ` (ex:
-:c:func:`PyObject_Malloc`) domains.
+*pymalloc* is the :ref:`default allocator <default-memory-allocators>` of the
+:c:data:`PYMEM_DOMAIN_MEM` (ex: :c:func:`PyMem_Malloc`) and
+:c:data:`PYMEM_DOMAIN_OBJ` (ex: :c:func:`PyObject_Malloc`) domains.
 
 The arena allocator uses the following functions:
 

Doc/using/cmdline.rst

@@ -687,6 +687,8 @@ conflict.
 
    Set the family of memory allocators used by Python:
 
+   * ``default``: use the :ref:`default memory allocators
+     <default-memory-allocators>`.
    * ``malloc``: use the :c:func:`malloc` function of the C library
      for all domains (:c:data:`PYMEM_DOMAIN_RAW`, :c:data:`PYMEM_DOMAIN_MEM`,
      :c:data:`PYMEM_DOMAIN_OBJ`).
@@ -696,20 +698,17 @@ conflict.
 
    Install debug hooks:
 
-   * ``debug``: install debug hooks on top of the default memory allocator
+   * ``debug``: install debug hooks on top of the :ref:`default memory
+     allocators <default-memory-allocators>`.
    * ``malloc_debug``: same as ``malloc`` but also install debug hooks
    * ``pymalloc_debug``: same as ``pymalloc`` but also install debug hooks
 
-   When Python is compiled in release mode, the default is ``pymalloc``. When
-   compiled in debug mode, the default is ``pymalloc_debug`` and the debug hooks
-   are used automatically.
+   See the :ref:`default memory allocators <default-memory-allocators>` and the
+   :c:func:`PyMem_SetupDebugHooks` function (install debug hooks on Python
+   memory allocators).
 
-   If Python is configured without ``pymalloc`` support, ``pymalloc`` and
-   ``pymalloc_debug`` are not available, the default is ``malloc`` in release
-   mode and ``malloc_debug`` in debug mode.
-
-   See the :c:func:`PyMem_SetupDebugHooks` function for debug hooks on Python
-   memory allocators.
+   .. versionchanged:: 3.7
+      Added the ``"default"`` allocator.
 
    .. versionadded:: 3.6
 

Include/pymem.h

@@ -21,6 +21,9 @@ PyAPI_FUNC(void) PyMem_RawFree(void *ptr);
    allocators. */
 PyAPI_FUNC(int) _PyMem_SetupAllocators(const char *opt);
 
+/* Try to get the allocators name set by _PyMem_SetupAllocators(). */
+PyAPI_FUNC(const char*) _PyMem_GetAllocatorsName(void);
+
 #ifdef WITH_PYMALLOC
 PyAPI_FUNC(int) _PyMem_PymallocEnabled(void);
 #endif
@@ -230,7 +233,12 @@ PyAPI_FUNC(void) PyMem_SetupDebugHooks(void);
 #endif
 
 #ifdef Py_BUILD_CORE
-PyAPI_FUNC(void) _PyMem_GetDefaultRawAllocator(PyMemAllocatorEx *alloc);
+/* 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. */
+PyAPI_FUNC(int) _PyMem_SetDefaultAllocator(
+    PyMemAllocatorDomain domain,
+    PyMemAllocatorEx *old_alloc);
 #endif
 
 #ifdef __cplusplus

Lib/test/pythoninfo.py

@@ -56,6 +56,14 @@ def copy_attributes(info_add, obj, name_fmt, attributes, *, formatter=None):
         info_add(name, value)
 
 
+def copy_attr(info_add, name, mod, attr_name):
+    try:
+        value = getattr(mod, attr_name)
+    except AttributeError:
+        return
+    info_add(name, value)
+
+
 def call_func(info_add, name, mod, func_name, *, formatter=None):
     try:
         func = getattr(mod, func_name)
@@ -168,11 +176,10 @@ def collect_os(info_add):
     call_func(info_add, 'os.gid', os, 'getgid')
     call_func(info_add, 'os.uname', os, 'uname')
 
-    if hasattr(os, 'getgroups'):
-        groups = os.getgroups()
-        groups = map(str, groups)
-        groups = ', '.join(groups)
-        info_add("os.groups", groups)
+    def format_groups(groups):
+        return ', '.join(map(str, groups))
+
+    call_func(info_add, 'os.groups', os, 'getgroups', formatter=format_groups)
 
     if hasattr(os, 'getlogin'):
         try:
@@ -184,11 +191,7 @@ def collect_os(info_add):
         else:
             info_add("os.login", login)
 
-    if hasattr(os, 'cpu_count'):
-        cpu_count = os.cpu_count()
-        if cpu_count:
-            info_add('os.cpu_count', cpu_count)
-
+    call_func(info_add, 'os.cpu_count', os, 'cpu_count')
     call_func(info_add, 'os.loadavg', os, 'getloadavg')
 
     # Get environment variables: filter to list
@@ -219,7 +222,9 @@ def collect_os(info_add):
     )
     for name, value in os.environ.items():
         uname = name.upper()
-        if (uname in ENV_VARS or uname.startswith(("PYTHON", "LC_"))
+        if (uname in ENV_VARS
+           # Copy PYTHON* and LC_* variables
+           or uname.startswith(("PYTHON", "LC_"))
            # Visual Studio: VS140COMNTOOLS
            or (uname.startswith("VS") and uname.endswith("COMNTOOLS"))):
             info_add('os.environ[%s]' % name, value)
@@ -313,12 +318,10 @@ def collect_time(info_add):
     )
     copy_attributes(info_add, time, 'time.%s', attributes)
 
-    if not hasattr(time, 'get_clock_info'):
-        return
-
-    for clock in ('time', 'perf_counter'):
-        tinfo = time.get_clock_info(clock)
-        info_add('time.%s' % clock, tinfo)
+    if hasattr(time, 'get_clock_info'):
+        for clock in ('time', 'perf_counter'):
+            tinfo = time.get_clock_info(clock)
+            info_add('time.%s' % clock, tinfo)
 
 
 def collect_sysconfig(info_add):
@@ -331,7 +334,6 @@ def collect_sysconfig(info_add):
         'CCSHARED',
         'CFLAGS',
         'CFLAGSFORSHARED',
-        'PY_LDFLAGS',
         'CONFIG_ARGS',
         'HOST_GNU_TYPE',
         'MACHDEP',
@@ -339,6 +341,7 @@ def collect_sysconfig(info_add):
         'OPT',
         'PY_CFLAGS',
         'PY_CFLAGS_NODIST',
+        'PY_LDFLAGS',
         'Py_DEBUG',
         'Py_ENABLE_SHARED',
         'SHELL',
@@ -422,6 +425,16 @@ def collect_decimal(info_add):
     copy_attributes(info_add, _decimal, '_decimal.%s', attributes)
 
 
+def collect_testcapi(info_add):
+    try:
+        import _testcapi
+    except ImportError:
+        return
+
+    call_func(info_add, 'pymem.allocator', _testcapi, 'pymem_getallocatorsname')
+    copy_attr(info_add, 'pymem.with_pymalloc', _testcapi, 'WITH_PYMALLOC')
+
+
 def collect_info(info):
     error = False
     info_add = info.add
@@ -444,6 +457,7 @@ def collect_info(info):
         collect_zlib,
         collect_expat,
         collect_decimal,
+        collect_testcapi,
     ):
         try:
             collect_func(info_add)

Lib/test/support/__init__.py

@@ -2848,3 +2848,8 @@ class SaveSignals:
     def restore(self):
         for signum, handler in self.handlers.items():
             self.signal.signal(signum, handler)
+
+
+def with_pymalloc():
+    import _testcapi
+    return _testcapi.WITH_PYMALLOC

Lib/test/test_capi.py

@@ -654,8 +654,7 @@ class PyMemMallocDebugTests(PyMemDebugTests):
     PYTHONMALLOC = 'malloc_debug'
 
 
-@unittest.skipUnless(sysconfig.get_config_var('WITH_PYMALLOC') == 1,
-                     'need pymalloc')
+@unittest.skipUnless(support.with_pymalloc(), 'need pymalloc')
 class PyMemPymallocDebugTests(PyMemDebugTests):
     PYTHONMALLOC = 'pymalloc_debug'
 

Lib/test/test_cmd_line.py

@@ -5,6 +5,7 @@
 import os
 import subprocess
 import sys
+import sysconfig
 import tempfile
 import unittest
 from test import support
@@ -559,10 +560,14 @@ class CmdLineTest(unittest.TestCase):
         except ImportError:
             pass
         else:
-            code = "import _testcapi; _testcapi.pymem_api_misuse()"
+            code = "import _testcapi; print(_testcapi.pymem_getallocatorsname())"
             with support.SuppressCrashReport():
                 out = self.run_xdev("-c", code, check_exitcode=False)
-            self.assertIn("Debug memory block at address p=", out)
+            if support.with_pymalloc():
+                alloc_name = "pymalloc_debug"
+            else:
+                alloc_name = "malloc_debug"
+            self.assertEqual(out, alloc_name)
 
         try:
             import faulthandler
@@ -573,6 +578,49 @@ class CmdLineTest(unittest.TestCase):
             out = self.run_xdev("-c", code)
             self.assertEqual(out, "True")
 
+    def check_pythonmalloc(self, env_var, name):
+        code = 'import _testcapi; print(_testcapi.pymem_getallocatorsname())'
+        env = dict(os.environ)
+        if env_var is not None:
+            env['PYTHONMALLOC'] = env_var
+        else:
+            env.pop('PYTHONMALLOC', None)
+        args = (sys.executable, '-c', code)
+        proc = subprocess.run(args,
+                              stdout=subprocess.PIPE,
+                              stderr=subprocess.STDOUT,
+                              universal_newlines=True,
+                              env=env)
+        self.assertEqual(proc.stdout.rstrip(), name)
+        self.assertEqual(proc.returncode, 0)
+
+    def test_pythonmalloc(self):
+        # Test the PYTHONMALLOC environment variable
+        pydebug = hasattr(sys, "gettotalrefcount")
+        pymalloc = support.with_pymalloc()
+        if pymalloc:
+            default_name = 'pymalloc_debug' if pydebug else 'pymalloc'
+            default_name_debug = 'pymalloc_debug'
+        else:
+            default_name = 'malloc_debug' if pydebug else 'malloc'
+            default_name_debug = 'malloc_debug'
+
+        tests = [
+            (None, default_name),
+            ('debug', default_name_debug),
+            ('malloc', 'malloc'),
+            ('malloc_debug', 'malloc_debug'),
+        ]
+        if pymalloc:
+            tests.extend((
+                ('pymalloc', 'pymalloc'),
+                ('pymalloc_debug', 'pymalloc_debug'),
+            ))
+
+        for env_var, name in tests:
+            with self.subTest(env_var=env_var, name=name):
+                self.check_pythonmalloc(env_var, name)
+
 
 class IgnoreEnvironmentTest(unittest.TestCase):
 

Lib/test/test_sys.py

@@ -753,8 +753,15 @@ class SysModuleTest(unittest.TestCase):
     @unittest.skipUnless(hasattr(sys, "getallocatedblocks"),
                          "sys.getallocatedblocks unavailable on this build")
     def test_getallocatedblocks(self):
+        try:
+            import _testcapi
+        except ImportError:
+            with_pymalloc = support.with_pymalloc()
+        else:
+            alloc_name = _testcapi.pymem_getallocatorsname()
+            with_pymalloc = (alloc_name in ('pymalloc', 'pymalloc_debug'))
+
         # Some sanity checks
-        with_pymalloc = sysconfig.get_config_var('WITH_PYMALLOC')
         a = sys.getallocatedblocks()
         self.assertIs(type(a), int)
         if with_pymalloc:

Modules/_testcapimodule.c

@@ -4104,6 +4104,19 @@ pymem_malloc_without_gil(PyObject *self, PyObject *args)
     Py_RETURN_NONE;
 }
 
+
+static PyObject*
+test_pymem_getallocatorsname(PyObject *self, PyObject *args)
+{
+    const char *name = _PyMem_GetAllocatorsName();
+    if (name == NULL) {
+        PyErr_SetString(PyExc_RuntimeError, "cannot get allocators name");
+        return NULL;
+    }
+    return PyUnicode_FromString(name);
+}
+
+
 static PyObject*
 pyobject_malloc_without_gil(PyObject *self, PyObject *args)
 {
@@ -4624,6 +4637,7 @@ static PyMethodDef TestMethods[] = {
     {"pymem_buffer_overflow", pymem_buffer_overflow, METH_NOARGS},
     {"pymem_api_misuse", pymem_api_misuse, METH_NOARGS},
     {"pymem_malloc_without_gil", pymem_malloc_without_gil, METH_NOARGS},
+    {"pymem_getallocatorsname", test_pymem_getallocatorsname, METH_NOARGS},
     {"pyobject_malloc_without_gil", pyobject_malloc_without_gil, METH_NOARGS},
     {"tracemalloc_track", tracemalloc_track, METH_VARARGS},
     {"tracemalloc_untrack", tracemalloc_untrack, METH_VARARGS},
@@ -5115,6 +5129,11 @@ PyInit__testcapi(void)
     PyModule_AddObject(m, "instancemethod", (PyObject *)&PyInstanceMethod_Type);
 
     PyModule_AddIntConstant(m, "the_number_three", 3);
+#ifdef WITH_PYMALLOC
+    PyModule_AddObject(m, "WITH_PYMALLOC", Py_True);
+#else
+    PyModule_AddObject(m, "WITH_PYMALLOC", Py_False);
+#endif
 
     TestError = PyErr_NewException("_testcapi.error", NULL, NULL);
     Py_INCREF(TestError);

Modules/main.c

@@ -475,11 +475,9 @@ pymain_free_impl(_PyMain *pymain)
 static void
 pymain_free(_PyMain *pymain)
 {
-    /* Force malloc() memory allocator */
-    PyMemAllocatorEx old_alloc, raw_alloc;
-    PyMem_GetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-    _PyMem_GetDefaultRawAllocator(&raw_alloc);
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &raw_alloc);
+    /* Force the allocator used by pymain_parse_cmdline_envvars() */
+    PyMemAllocatorEx old_alloc;
+    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
 
     pymain_free_impl(pymain);
 
@@ -1561,17 +1559,14 @@ pymain_parse_cmdline_envvars_impl(_PyMain *pymain)
 static int
 pymain_parse_cmdline_envvars(_PyMain *pymain)
 {
-    /* Force malloc() memory allocator */
-    PyMemAllocatorEx old_alloc, raw_alloc;
-    PyMem_GetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-    _PyMem_GetDefaultRawAllocator(&raw_alloc);
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &raw_alloc);
+    /* Force default allocator, since pymain_free() must use the same allocator
+       than this function. */
+    PyMemAllocatorEx old_alloc;
+    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
 
     int res = pymain_parse_cmdline_envvars_impl(pymain);
 
-    /* Restore the old memory allocator */
     PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-
     return res;
 }
 

Objects/obmalloc.c

@@ -26,6 +26,8 @@ static void _PyMem_DebugFree(void *ctx, void *p);
 static void _PyObject_DebugDumpAddress(const void *p);
 static void _PyMem_DebugCheckAddress(char api_id, const void *p);
 
+static void _PyMem_SetupDebugHooksDomain(PyMemAllocatorDomain domain);
+
 #if defined(__has_feature)  /* Clang */
  #if __has_feature(address_sanitizer)  /* is ASAN enabled? */
   #define ATTRIBUTE_NO_ADDRESS_SAFETY_ANALYSIS \
@@ -149,14 +151,18 @@ _PyObject_ArenaFree(void *ctx, void *ptr, size_t size)
 }
 #endif
 
+#define MALLOC_ALLOC {NULL, _PyMem_RawMalloc, _PyMem_RawCalloc, _PyMem_RawRealloc, _PyMem_RawFree}
+#ifdef WITH_PYMALLOC
+#  define PYMALLOC_ALLOC {NULL, _PyObject_Malloc, _PyObject_Calloc, _PyObject_Realloc, _PyObject_Free}
+#endif
 
-#define PYRAW_FUNCS _PyMem_RawMalloc, _PyMem_RawCalloc, _PyMem_RawRealloc, _PyMem_RawFree
+#define PYRAW_ALLOC MALLOC_ALLOC
 #ifdef WITH_PYMALLOC
-#  define PYOBJ_FUNCS _PyObject_Malloc, _PyObject_Calloc, _PyObject_Realloc, _PyObject_Free
+#  define PYOBJ_ALLOC PYMALLOC_ALLOC
 #else
-#  define PYOBJ_FUNCS PYRAW_FUNCS
+#  define PYOBJ_ALLOC MALLOC_ALLOC
 #endif
-#define PYMEM_FUNCS PYOBJ_FUNCS
+#define PYMEM_ALLOC PYOBJ_ALLOC
 
 typedef struct {
     /* We tag each block with an API ID in order to tag API violations */
@@ -168,103 +174,118 @@ static struct {
     debug_alloc_api_t mem;
     debug_alloc_api_t obj;
 } _PyMem_Debug = {
-    {'r', {NULL, PYRAW_FUNCS}},
-    {'m', {NULL, PYMEM_FUNCS}},
-    {'o', {NULL, PYOBJ_FUNCS}}
+    {'r', PYRAW_ALLOC},
+    {'m', PYMEM_ALLOC},
+    {'o', PYOBJ_ALLOC}
     };
 
-#define PYRAWDBG_FUNCS \
-    _PyMem_DebugRawMalloc, _PyMem_DebugRawCalloc, _PyMem_DebugRawRealloc, _PyMem_DebugRawFree
-#define PYDBG_FUNCS \
-    _PyMem_DebugMalloc, _PyMem_DebugCalloc, _PyMem_DebugRealloc, _PyMem_DebugFree
+#define PYDBGRAW_ALLOC \
+    {&_PyMem_Debug.raw, _PyMem_DebugRawMalloc, _PyMem_DebugRawCalloc, _PyMem_DebugRawRealloc, _PyMem_DebugRawFree}
+#define PYDBGMEM_ALLOC \
+    {&_PyMem_Debug.mem, _PyMem_DebugMalloc, _PyMem_DebugCalloc, _PyMem_DebugRealloc, _PyMem_DebugFree}
+#define PYDBGOBJ_ALLOC \
+    {&_PyMem_Debug.obj, _PyMem_DebugMalloc, _PyMem_DebugCalloc, _PyMem_DebugRealloc, _PyMem_DebugFree}
 
-static PyMemAllocatorEx _PyMem_Raw = {
 #ifdef Py_DEBUG
-    &_PyMem_Debug.raw, PYRAWDBG_FUNCS
+static PyMemAllocatorEx _PyMem_Raw = PYDBGRAW_ALLOC;
+static PyMemAllocatorEx _PyMem = PYDBGMEM_ALLOC;
+static PyMemAllocatorEx _PyObject = PYDBGOBJ_ALLOC;
 #else
-    NULL, PYRAW_FUNCS
+static PyMemAllocatorEx _PyMem_Raw = PYRAW_ALLOC;
+static PyMemAllocatorEx _PyMem = PYMEM_ALLOC;
+static PyMemAllocatorEx _PyObject = PYOBJ_ALLOC;
 #endif
-    };
 
-static PyMemAllocatorEx _PyMem = {
-#ifdef Py_DEBUG
-    &_PyMem_Debug.mem, PYDBG_FUNCS
-#else
-    NULL, PYMEM_FUNCS
-#endif
-    };
 
-static PyMemAllocatorEx _PyObject = {
-#ifdef Py_DEBUG
-    &_PyMem_Debug.obj, PYDBG_FUNCS
-#else
-    NULL, PYOBJ_FUNCS
-#endif
-    };
+static int
+pymem_set_default_allocator(PyMemAllocatorDomain domain, int debug,
+                            PyMemAllocatorEx *old_alloc)
+{
+    if (old_alloc != NULL) {
+        PyMem_GetAllocator(domain, old_alloc);
+    }
 
-void
-_PyMem_GetDefaultRawAllocator(PyMemAllocatorEx *alloc_p)
+
+    PyMemAllocatorEx new_alloc;
+    switch(domain)
+    {
+    case PYMEM_DOMAIN_RAW:
+        new_alloc = (PyMemAllocatorEx)PYRAW_ALLOC;
+        break;
+    case PYMEM_DOMAIN_MEM:
+        new_alloc = (PyMemAllocatorEx)PYMEM_ALLOC;
+        break;
+    case PYMEM_DOMAIN_OBJ:
+        new_alloc = (PyMemAllocatorEx)PYOBJ_ALLOC;
+        break;
+    default:
+        /* unknown domain */
+        return -1;
+    }
+    PyMem_SetAllocator(domain, &new_alloc);
+    if (debug) {
+        _PyMem_SetupDebugHooksDomain(domain);
+    }
+    return 0;
+}
+
+
+int
+_PyMem_SetDefaultAllocator(PyMemAllocatorDomain domain,
+                           PyMemAllocatorEx *old_alloc)
 {
 #ifdef Py_DEBUG
-    PyMemAllocatorEx alloc = {&_PyMem_Debug.raw, PYDBG_FUNCS};
+    const int debug = 1;
 #else
-    PyMemAllocatorEx alloc = {NULL, PYRAW_FUNCS};
+    const int debug = 0;
 #endif
-    *alloc_p = alloc;
+    return pymem_set_default_allocator(domain, debug, old_alloc);
 }
 
+
 int
 _PyMem_SetupAllocators(const char *opt)
 {
     if (opt == NULL || *opt == '\0') {
         /* PYTHONMALLOC is empty or is not set or ignored (-E/-I command line
-           options): use default allocators */
-#ifdef Py_DEBUG
-#  ifdef WITH_PYMALLOC
-        opt = "pymalloc_debug";
-#  else
-        opt = "malloc_debug";
-#  endif
-#else
-   /* !Py_DEBUG */
-#  ifdef WITH_PYMALLOC
-        opt = "pymalloc";
-#  else
-        opt = "malloc";
-#  endif
-#endif
+           options): use default memory allocators */
+        opt = "default";
     }
 
-    if (strcmp(opt, "debug") == 0) {
-        PyMem_SetupDebugHooks();
+    if (strcmp(opt, "default") == 0) {
+        (void)_PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, NULL);
+        (void)_PyMem_SetDefaultAllocator(PYMEM_DOMAIN_MEM, NULL);
+        (void)_PyMem_SetDefaultAllocator(PYMEM_DOMAIN_OBJ, NULL);
     }
-    else if (strcmp(opt, "malloc") == 0 || strcmp(opt, "malloc_debug") == 0)
-    {
-        PyMemAllocatorEx alloc = {NULL, PYRAW_FUNCS};
-
-        PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &alloc);
-        PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &alloc);
-        PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &alloc);
-
-        if (strcmp(opt, "malloc_debug") == 0)
-            PyMem_SetupDebugHooks();
+    else if (strcmp(opt, "debug") == 0) {
+        (void)pymem_set_default_allocator(PYMEM_DOMAIN_RAW, 1, NULL);
+        (void)pymem_set_default_allocator(PYMEM_DOMAIN_MEM, 1, NULL);
+        (void)pymem_set_default_allocator(PYMEM_DOMAIN_OBJ, 1, NULL);
     }
 #ifdef WITH_PYMALLOC
-    else if (strcmp(opt, "pymalloc") == 0
-             || strcmp(opt, "pymalloc_debug") == 0)
-    {
-        PyMemAllocatorEx raw_alloc = {NULL, PYRAW_FUNCS};
-        PyMemAllocatorEx mem_alloc = {NULL, PYMEM_FUNCS};
-        PyMemAllocatorEx obj_alloc = {NULL, PYOBJ_FUNCS};
+    else if (strcmp(opt, "pymalloc") == 0 || strcmp(opt, "pymalloc_debug") == 0) {
+        PyMemAllocatorEx malloc_alloc = MALLOC_ALLOC;
+        PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &malloc_alloc);
 
-        PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &raw_alloc);
-        PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &mem_alloc);
-        PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &obj_alloc);
+        PyMemAllocatorEx pymalloc = PYMALLOC_ALLOC;
+        PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &pymalloc);
+        PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &pymalloc);
 
-        if (strcmp(opt, "pymalloc_debug") == 0)
+        if (strcmp(opt, "pymalloc_debug") == 0) {
             PyMem_SetupDebugHooks();
+        }
     }
 #endif
+    else if (strcmp(opt, "malloc") == 0 || strcmp(opt, "malloc_debug") == 0) {
+        PyMemAllocatorEx malloc_alloc = MALLOC_ALLOC;
+        PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &malloc_alloc);
+        PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &malloc_alloc);
+        PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &malloc_alloc);
+
+        if (strcmp(opt, "malloc_debug") == 0) {
+            PyMem_SetupDebugHooks();
+        }
+    }
     else {
         /* unknown allocator */
         return -1;
@@ -272,11 +293,74 @@ _PyMem_SetupAllocators(const char *opt)
     return 0;
 }
 
-#undef PYRAW_FUNCS
-#undef PYMEM_FUNCS
-#undef PYOBJ_FUNCS
-#undef PYRAWDBG_FUNCS
-#undef PYDBG_FUNCS
+
+static int
+pymemallocator_eq(PyMemAllocatorEx *a, PyMemAllocatorEx *b)
+{
+    return (memcmp(a, b, sizeof(PyMemAllocatorEx)) == 0);
+}
+
+
+const char*
+_PyMem_GetAllocatorsName(void)
+{
+    PyMemAllocatorEx malloc_alloc = MALLOC_ALLOC;
+#ifdef WITH_PYMALLOC
+    PyMemAllocatorEx pymalloc = PYMALLOC_ALLOC;
+#endif
+
+    if (pymemallocator_eq(&_PyMem_Raw, &malloc_alloc) &&
+        pymemallocator_eq(&_PyMem, &malloc_alloc) &&
+        pymemallocator_eq(&_PyObject, &malloc_alloc))
+    {
+        return "malloc";
+    }
+#ifdef WITH_PYMALLOC
+    if (pymemallocator_eq(&_PyMem_Raw, &malloc_alloc) &&
+        pymemallocator_eq(&_PyMem, &pymalloc) &&
+        pymemallocator_eq(&_PyObject, &pymalloc))
+    {
+        return "pymalloc";
+    }
+#endif
+
+    PyMemAllocatorEx dbg_raw = PYDBGRAW_ALLOC;
+    PyMemAllocatorEx dbg_mem = PYDBGMEM_ALLOC;
+    PyMemAllocatorEx dbg_obj = PYDBGOBJ_ALLOC;
+
+    if (pymemallocator_eq(&_PyMem_Raw, &dbg_raw) &&
+        pymemallocator_eq(&_PyMem, &dbg_mem) &&
+        pymemallocator_eq(&_PyObject, &dbg_obj))
+    {
+        /* Debug hooks installed */
+        if (pymemallocator_eq(&_PyMem_Debug.raw.alloc, &malloc_alloc) &&
+            pymemallocator_eq(&_PyMem_Debug.mem.alloc, &malloc_alloc) &&
+            pymemallocator_eq(&_PyMem_Debug.obj.alloc, &malloc_alloc))
+        {
+            return "malloc_debug";
+        }
+#ifdef WITH_PYMALLOC
+        if (pymemallocator_eq(&_PyMem_Debug.raw.alloc, &malloc_alloc) &&
+            pymemallocator_eq(&_PyMem_Debug.mem.alloc, &pymalloc) &&
+            pymemallocator_eq(&_PyMem_Debug.obj.alloc, &pymalloc))
+        {
+            return "pymalloc_debug";
+        }
+#endif
+    }
+    return NULL;
+}
+
+
+#undef MALLOC_ALLOC
+#undef PYMALLOC_ALLOC
+#undef PYRAW_ALLOC
+#undef PYMEM_ALLOC
+#undef PYOBJ_ALLOC
+#undef PYDBGRAW_ALLOC
+#undef PYDBGMEM_ALLOC
+#undef PYDBGOBJ_ALLOC
+
 
 static PyObjectArenaAllocator _PyObject_Arena = {NULL,
 #ifdef MS_WINDOWS
@@ -307,40 +391,62 @@ _PyMem_PymallocEnabled(void)
 }
 #endif
 
-void
-PyMem_SetupDebugHooks(void)
+
+static void
+_PyMem_SetupDebugHooksDomain(PyMemAllocatorDomain domain)
 {
     PyMemAllocatorEx alloc;
 
-    alloc.malloc = _PyMem_DebugRawMalloc;
-    alloc.calloc = _PyMem_DebugRawCalloc;
-    alloc.realloc = _PyMem_DebugRawRealloc;
-    alloc.free = _PyMem_DebugRawFree;
+    if (domain == PYMEM_DOMAIN_RAW) {
+        if (_PyMem_Raw.malloc == _PyMem_DebugRawMalloc) {
+            return;
+        }
 
-    if (_PyMem_Raw.malloc != _PyMem_DebugRawMalloc) {
-        alloc.ctx = &_PyMem_Debug.raw;
         PyMem_GetAllocator(PYMEM_DOMAIN_RAW, &_PyMem_Debug.raw.alloc);
+        alloc.ctx = &_PyMem_Debug.raw;
+        alloc.malloc = _PyMem_DebugRawMalloc;
+        alloc.calloc = _PyMem_DebugRawCalloc;
+        alloc.realloc = _PyMem_DebugRawRealloc;
+        alloc.free = _PyMem_DebugRawFree;
         PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &alloc);
     }
+    else if (domain == PYMEM_DOMAIN_MEM) {
+        if (_PyMem.malloc == _PyMem_DebugMalloc) {
+            return;
+        }
 
-    alloc.malloc = _PyMem_DebugMalloc;
-    alloc.calloc = _PyMem_DebugCalloc;
-    alloc.realloc = _PyMem_DebugRealloc;
-    alloc.free = _PyMem_DebugFree;
-
-    if (_PyMem.malloc != _PyMem_DebugMalloc) {
-        alloc.ctx = &_PyMem_Debug.mem;
         PyMem_GetAllocator(PYMEM_DOMAIN_MEM, &_PyMem_Debug.mem.alloc);
+        alloc.ctx = &_PyMem_Debug.mem;
+        alloc.malloc = _PyMem_DebugMalloc;
+        alloc.calloc = _PyMem_DebugCalloc;
+        alloc.realloc = _PyMem_DebugRealloc;
+        alloc.free = _PyMem_DebugFree;
         PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &alloc);
     }
+    else if (domain == PYMEM_DOMAIN_OBJ)  {
+        if (_PyObject.malloc == _PyMem_DebugMalloc) {
+            return;
+        }
 
-    if (_PyObject.malloc != _PyMem_DebugMalloc) {
-        alloc.ctx = &_PyMem_Debug.obj;
         PyMem_GetAllocator(PYMEM_DOMAIN_OBJ, &_PyMem_Debug.obj.alloc);
+        alloc.ctx = &_PyMem_Debug.obj;
+        alloc.malloc = _PyMem_DebugMalloc;
+        alloc.calloc = _PyMem_DebugCalloc;
+        alloc.realloc = _PyMem_DebugRealloc;
+        alloc.free = _PyMem_DebugFree;
         PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &alloc);
     }
 }
 
+
+void
+PyMem_SetupDebugHooks(void)
+{
+    _PyMem_SetupDebugHooksDomain(PYMEM_DOMAIN_RAW);
+    _PyMem_SetupDebugHooksDomain(PYMEM_DOMAIN_MEM);
+    _PyMem_SetupDebugHooksDomain(PYMEM_DOMAIN_OBJ);
+}
+
 void
 PyMem_GetAllocator(PyMemAllocatorDomain domain, PyMemAllocatorEx *allocator)
 {

Programs/python.c

@@ -33,12 +33,9 @@ main(int argc, char **argv)
         exit(1);
     }
 
-    /* Force malloc() allocator to bootstrap Python */
-#ifdef Py_DEBUG
-    (void)_PyMem_SetupAllocators("malloc_debug");
-#  else
-    (void)_PyMem_SetupAllocators("malloc");
-#  endif
+    /* Force default allocator, to be able to release memory above
+       with a known allocator. */
+    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, NULL);
 
     argv_copy = (wchar_t **)PyMem_RawMalloc(sizeof(wchar_t*) * (argc+1));
     argv_copy2 = (wchar_t **)PyMem_RawMalloc(sizeof(wchar_t*) * (argc+1));
@@ -98,13 +95,9 @@ main(int argc, char **argv)
 
     status = Py_Main(argc, argv_copy);
 
-    /* Force again malloc() allocator to release memory blocks allocated
-       before Py_Main() */
-#ifdef Py_DEBUG
-    (void)_PyMem_SetupAllocators("malloc_debug");
-#  else
-    (void)_PyMem_SetupAllocators("malloc");
-#  endif
+    /* Py_Main() can change PyMem_RawMalloc() allocator, so restore the default
+       to release memory blocks allocated before Py_Main() */
+    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, NULL);
 
     for (i = 0; i < argc; i++) {
         PyMem_RawFree(argv_copy2[i]);

Python/pylifecycle.c

@@ -630,7 +630,7 @@ _Py_InitializeCore(const _PyCoreConfig *config)
     }
 
     if (_PyMem_SetupAllocators(core_config.allocator) < 0) {
-        return _Py_INIT_ERR("Unknown PYTHONMALLOC allocator");
+        return _Py_INIT_USER_ERR("Unknown PYTHONMALLOC allocator");
     }
 
     if (_PyRuntime.initialized) {

Python/pystate.c

@@ -35,8 +35,8 @@ to avoid the expense of doing their own locking).
 extern "C" {
 #endif
 
-_PyInitError
-_PyRuntimeState_Init(_PyRuntimeState *runtime)
+static _PyInitError
+_PyRuntimeState_Init_impl(_PyRuntimeState *runtime)
 {
     memset(runtime, 0, sizeof(*runtime));
 
@@ -59,14 +59,26 @@ _PyRuntimeState_Init(_PyRuntimeState *runtime)
     return _Py_INIT_OK();
 }
 
+_PyInitError
+_PyRuntimeState_Init(_PyRuntimeState *runtime)
+{
+    /* Force default allocator, since _PyRuntimeState_Fini() must
+       use the same allocator than this function. */
+    PyMemAllocatorEx old_alloc;
+    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
+
+    _PyInitError err = _PyRuntimeState_Init_impl(runtime);
+
+    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
+    return err;
+}
+
 void
 _PyRuntimeState_Fini(_PyRuntimeState *runtime)
 {
-    /* Use the same memory allocator than _PyRuntimeState_Init() */
-    PyMemAllocatorEx old_alloc, raw_alloc;
-    PyMem_GetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-    _PyMem_GetDefaultRawAllocator(&raw_alloc);
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &raw_alloc);
+    /* Force the allocator used by _PyRuntimeState_Init(). */
+    PyMemAllocatorEx old_alloc;
+    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
 
     if (runtime->interpreters.mutex != NULL) {
         PyThread_free_lock(runtime->interpreters.mutex);