bpo-36763: Add PyMemAllocatorName (GH-13387)

* Add PyMemAllocatorName enum
* _PyPreConfig.allocator type becomes PyMemAllocatorName, instead of
  char*
* Remove _PyPreConfig_Clear()
* Add _PyMem_GetAllocatorName()
* Rename _PyMem_GetAllocatorsName() to
  _PyMem_GetCurrentAllocatorName()
* Remove _PyPreConfig_SetAllocator(): just call
  _PyMem_SetupAllocators() directly, we don't have do reallocate the
  configuration with the new allocator anymore!
* _PyPreConfig_Write() parameter becomes const, as it should be in
  the first place!

12 files changed, 119 insertions, 112 deletions

diff --git a/Include/cpython/coreconfig.h b/Include/cpython/coreconfig.h
index 1ba2663c966..dca41341dfe 100644
--- a/Include/cpython/coreconfig.h
+++ b/Include/cpython/coreconfig.h
@@ -120,7 +120,9 @@ typedef struct {
     int utf8_mode;
 
     int dev_mode;           /* Development mode. PYTHONDEVMODE, -X dev */
-    char *allocator;        /* Memory allocator: PYTHONMALLOC */
+
+    /* Memory allocator: PYTHONMALLOC env var */
+    PyMemAllocatorName allocator;
 } _PyPreConfig;
 
 #ifdef MS_WINDOWS
@@ -137,7 +139,7 @@ typedef struct {
         .isolated = -1, \
         .use_environment = -1, \
         .dev_mode = -1, \
-        .allocator = NULL}
+        .allocator = PYMEM_ALLOCATOR_NOT_SET}
 
 
 /* --- _PyCoreConfig ---------------------------------------------- */
diff --git a/Include/cpython/pymem.h b/Include/cpython/pymem.h
index bd66506639a..79f063b1217 100644
--- a/Include/cpython/pymem.h
+++ b/Include/cpython/pymem.h
@@ -11,12 +11,8 @@ PyAPI_FUNC(void *) PyMem_RawCalloc(size_t nelem, size_t elsize);
 PyAPI_FUNC(void *) PyMem_RawRealloc(void *ptr, size_t new_size);
 PyAPI_FUNC(void) PyMem_RawFree(void *ptr);
 
-/* Configure the Python memory allocators. Pass NULL to use default
-   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);
+PyAPI_FUNC(const char*) _PyMem_GetCurrentAllocatorName(void);
 
 PyAPI_FUNC(void *) PyMem_Calloc(size_t nelem, size_t elsize);
 
@@ -41,6 +37,19 @@ typedef enum {
     PYMEM_DOMAIN_OBJ
 } PyMemAllocatorDomain;
 
+typedef enum {
+    PYMEM_ALLOCATOR_NOT_SET = 0,
+    PYMEM_ALLOCATOR_DEFAULT = 1,
+    PYMEM_ALLOCATOR_DEBUG = 2,
+    PYMEM_ALLOCATOR_MALLOC = 3,
+    PYMEM_ALLOCATOR_MALLOC_DEBUG = 4,
+#ifdef WITH_PYMALLOC
+    PYMEM_ALLOCATOR_PYMALLOC = 5,
+    PYMEM_ALLOCATOR_PYMALLOC_DEBUG = 6,
+#endif
+} PyMemAllocatorName;
+
+
 typedef struct {
     /* user context passed as the first argument to the 4 functions */
     void *ctx;
diff --git a/Include/internal/pycore_coreconfig.h b/Include/internal/pycore_coreconfig.h
index d48904e482a..ccb7948ef4d 100644
--- a/Include/internal/pycore_coreconfig.h
+++ b/Include/internal/pycore_coreconfig.h
@@ -88,7 +88,6 @@ PyAPI_FUNC(_PyInitError) _PyPreCmdline_Read(_PyPreCmdline *cmdline,
 
 /* --- _PyPreConfig ----------------------------------------------- */
 
-PyAPI_FUNC(void) _PyPreConfig_Clear(_PyPreConfig *config);
 PyAPI_FUNC(int) _PyPreConfig_Copy(_PyPreConfig *config,
     const _PyPreConfig *config2);
 PyAPI_FUNC(PyObject*) _PyPreConfig_AsDict(const _PyPreConfig *config);
@@ -96,7 +95,7 @@ PyAPI_FUNC(void) _PyCoreConfig_GetCoreConfig(_PyPreConfig *config,
     const _PyCoreConfig *core_config);
 PyAPI_FUNC(_PyInitError) _PyPreConfig_Read(_PyPreConfig *config,
     const _PyArgv *args);
-PyAPI_FUNC(_PyInitError) _PyPreConfig_Write(_PyPreConfig *config);
+PyAPI_FUNC(_PyInitError) _PyPreConfig_Write(const _PyPreConfig *config);
 
 
 /* --- _PyCoreConfig ---------------------------------------------- */
diff --git a/Include/internal/pycore_pymem.h b/Include/internal/pycore_pymem.h
index 20f3b5e4006..dcc492af017 100644
--- a/Include/internal/pycore_pymem.h
+++ b/Include/internal/pycore_pymem.h
@@ -179,6 +179,15 @@ static inline int _PyMem_IsPtrFreed(void *ptr)
 #endif
 }
 
+PyAPI_FUNC(int) _PyMem_GetAllocatorName(
+    const char *name,
+    PyMemAllocatorName *allocator);
+
+/* Configure the Python memory allocators.
+   Pass PYMEM_ALLOCATOR_DEFAULT to use default allocators.
+   PYMEM_ALLOCATOR_NOT_SET does nothing. */
+PyAPI_FUNC(int) _PyMem_SetupAllocators(PyMemAllocatorName allocator);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/Lib/test/test_embed.py b/Lib/test/test_embed.py
index 4012a389d75..92cc405859c 100644
--- a/Lib/test/test_embed.py
+++ b/Lib/test/test_embed.py
@@ -13,6 +13,9 @@ import textwrap
 
 
 MS_WINDOWS = (os.name == 'nt')
+PYMEM_ALLOCATOR_NOT_SET = 0
+PYMEM_ALLOCATOR_DEBUG = 2
+PYMEM_ALLOCATOR_MALLOC = 3
 
 
 class EmbeddingTestsMixin:
@@ -272,7 +275,7 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
     # Mark config which should be get by get_default_config()
     GET_DEFAULT_CONFIG = object()
     DEFAULT_PRE_CONFIG = {
-        'allocator': None,
+        'allocator': PYMEM_ALLOCATOR_NOT_SET,
         'coerce_c_locale': 0,
         'coerce_c_locale_warn': 0,
         'utf8_mode': 0,
@@ -564,7 +567,7 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
 
     def test_init_from_config(self):
         preconfig = {
-            'allocator': 'malloc',
+            'allocator': PYMEM_ALLOCATOR_MALLOC,
             'utf8_mode': 1,
         }
         config = {
@@ -608,7 +611,7 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
         self.check_config("init_from_config", config, preconfig)
 
     INIT_ENV_PRECONFIG = {
-        'allocator': 'malloc',
+        'allocator': PYMEM_ALLOCATOR_MALLOC,
     }
     INIT_ENV_CONFIG = {
         'use_hash_seed': 1,
@@ -633,7 +636,7 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
 
     def test_init_env_dev_mode(self):
         preconfig = dict(self.INIT_ENV_PRECONFIG,
-                      allocator='debug')
+                      allocator=PYMEM_ALLOCATOR_DEBUG)
         config = dict(self.INIT_ENV_CONFIG,
                       dev_mode=1,
                       warnoptions=['default'])
@@ -641,7 +644,7 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
 
     def test_init_env_dev_mode_alloc(self):
         preconfig = dict(self.INIT_ENV_PRECONFIG,
-                         allocator='malloc')
+                         allocator=PYMEM_ALLOCATOR_MALLOC)
         config = dict(self.INIT_ENV_CONFIG,
                       dev_mode=1,
                       warnoptions=['default'])
@@ -649,7 +652,7 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
 
     def test_init_dev_mode(self):
         preconfig = {
-            'allocator': 'debug',
+            'allocator': PYMEM_ALLOCATOR_DEBUG,
         }
         config = {
             'faulthandler': 1,
diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
index 04d75ace6ec..59b42df279c 100644
--- a/Modules/_testcapimodule.c
+++ b/Modules/_testcapimodule.c
@@ -4297,7 +4297,7 @@ pymem_malloc_without_gil(PyObject *self, PyObject *args)
 static PyObject*
 test_pymem_getallocatorsname(PyObject *self, PyObject *args)
 {
-    const char *name = _PyMem_GetAllocatorsName();
+    const char *name = _PyMem_GetCurrentAllocatorName();
     if (name == NULL) {
         PyErr_SetString(PyExc_RuntimeError, "cannot get allocators name");
         return NULL;
diff --git a/Objects/obmalloc.c b/Objects/obmalloc.c
index bd15bcf1363..f54856dcfe7 100644
--- a/Objects/obmalloc.c
+++ b/Objects/obmalloc.c
@@ -268,26 +268,65 @@ _PyMem_SetDefaultAllocator(PyMemAllocatorDomain domain,
 
 
 int
-_PyMem_SetupAllocators(const char *opt)
+_PyMem_GetAllocatorName(const char *name, PyMemAllocatorName *allocator)
 {
-    if (opt == NULL || *opt == '\0') {
+    if (name == NULL || *name == '\0') {
         /* PYTHONMALLOC is empty or is not set or ignored (-E/-I command line
-           options): use default memory allocators */
-        opt = "default";
+           nameions): use default memory allocators */
+        *allocator = PYMEM_ALLOCATOR_DEFAULT;
     }
+    else if (strcmp(name, "default") == 0) {
+        *allocator = PYMEM_ALLOCATOR_DEFAULT;
+    }
+    else if (strcmp(name, "debug") == 0) {
+        *allocator = PYMEM_ALLOCATOR_DEBUG;
+    }
+#ifdef WITH_PYMALLOC
+    else if (strcmp(name, "pymalloc") == 0) {
+        *allocator = PYMEM_ALLOCATOR_PYMALLOC;
+    }
+    else if (strcmp(name, "pymalloc_debug") == 0) {
+        *allocator = PYMEM_ALLOCATOR_PYMALLOC_DEBUG;
+    }
+#endif
+    else if (strcmp(name, "malloc") == 0) {
+        *allocator = PYMEM_ALLOCATOR_MALLOC;
+    }
+    else if (strcmp(name, "malloc_debug") == 0) {
+        *allocator = PYMEM_ALLOCATOR_MALLOC_DEBUG;
+    }
+    else {
+        /* unknown allocator */
+        return -1;
+    }
+    return 0;
+}
+
 
-    if (strcmp(opt, "default") == 0) {
+int
+_PyMem_SetupAllocators(PyMemAllocatorName allocator)
+{
+    switch (allocator) {
+    case PYMEM_ALLOCATOR_NOT_SET:
+        /* do nothing */
+        break;
+
+    case PYMEM_ALLOCATOR_DEFAULT:
         (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, "debug") == 0) {
+        break;
+
+    case PYMEM_ALLOCATOR_DEBUG:
         (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);
-    }
+        break;
+
 #ifdef WITH_PYMALLOC
-    else if (strcmp(opt, "pymalloc") == 0 || strcmp(opt, "pymalloc_debug") == 0) {
+    case PYMEM_ALLOCATOR_PYMALLOC:
+    case PYMEM_ALLOCATOR_PYMALLOC_DEBUG:
+    {
         PyMemAllocatorEx malloc_alloc = MALLOC_ALLOC;
         PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &malloc_alloc);
 
@@ -295,22 +334,28 @@ _PyMem_SetupAllocators(const char *opt)
         PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &pymalloc);
         PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &pymalloc);
 
-        if (strcmp(opt, "pymalloc_debug") == 0) {
+        if (allocator == PYMEM_ALLOCATOR_PYMALLOC_DEBUG) {
             PyMem_SetupDebugHooks();
         }
+        break;
     }
 #endif
-    else if (strcmp(opt, "malloc") == 0 || strcmp(opt, "malloc_debug") == 0) {
+
+    case PYMEM_ALLOCATOR_MALLOC:
+    case PYMEM_ALLOCATOR_MALLOC_DEBUG:
+    {
         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) {
+        if (allocator == PYMEM_ALLOCATOR_MALLOC_DEBUG) {
             PyMem_SetupDebugHooks();
         }
+        break;
     }
-    else {
+
+    default:
         /* unknown allocator */
         return -1;
     }
@@ -326,7 +371,7 @@ pymemallocator_eq(PyMemAllocatorEx *a, PyMemAllocatorEx *b)
 
 
 const char*
-_PyMem_GetAllocatorsName(void)
+_PyMem_GetCurrentAllocatorName(void)
 {
     PyMemAllocatorEx malloc_alloc = MALLOC_ALLOC;
 #ifdef WITH_PYMALLOC
diff --git a/Programs/_testembed.c b/Programs/_testembed.c
index b1b7c6e0ffc..3327c8ce663 100644
--- a/Programs/_testembed.c
+++ b/Programs/_testembed.c
@@ -379,7 +379,7 @@ static int test_init_from_config(void)
     _PyPreConfig preconfig = _PyPreConfig_INIT;
 
     putenv("PYTHONMALLOC=malloc_debug");
-    preconfig.allocator = "malloc";
+    preconfig.allocator = PYMEM_ALLOCATOR_MALLOC;
 
     putenv("PYTHONUTF8=0");
     Py_UTF8Mode = 0;
diff --git a/Python/coreconfig.c b/Python/coreconfig.c
index c20ae2151c9..634891ed214 100644
--- a/Python/coreconfig.c
+++ b/Python/coreconfig.c
@@ -2021,26 +2021,22 @@ core_read_precmdline(_PyCoreConfig *config, _PyPreCmdline *precmdline)
     _PyPreConfig preconfig = _PyPreConfig_INIT;
     if (_PyPreConfig_Copy(&preconfig, &_PyRuntime.preconfig) < 0) {
         err = _Py_INIT_NO_MEMORY();
-        goto done;
+        return err;
     }
 
     _PyCoreConfig_GetCoreConfig(&preconfig, config);
 
     err = _PyPreCmdline_Read(precmdline, &preconfig);
     if (_Py_INIT_FAILED(err)) {
-        goto done;
+        return err;
     }
 
     if (_PyPreCmdline_SetCoreConfig(precmdline, config) < 0) {
         err = _Py_INIT_NO_MEMORY();
-        goto done;
+        return err;
     }
 
-    err = _Py_INIT_OK();
-
-done:
-    _PyPreConfig_Clear(&preconfig);
-    return err;
+    return _Py_INIT_OK();
 }
 
 
diff --git a/Python/preconfig.c b/Python/preconfig.c
index 2d0ace5df29..2bbf8e6fb7f 100644
--- a/Python/preconfig.c
+++ b/Python/preconfig.c
@@ -260,19 +260,9 @@ _PyPreCmdline_Read(_PyPreCmdline *cmdline,
 
 /* --- _PyPreConfig ----------------------------------------------- */
 
-void
-_PyPreConfig_Clear(_PyPreConfig *config)
-{
-    PyMem_RawFree(config->allocator);
-    config->allocator = NULL;
-}
-
-
 int
 _PyPreConfig_Copy(_PyPreConfig *config, const _PyPreConfig *config2)
 {
-    _PyPreConfig_Clear(config);
-
 #define COPY_ATTR(ATTR) config->ATTR = config2->ATTR
 #define COPY_STR_ATTR(ATTR) \
     do { \
@@ -293,7 +283,7 @@ _PyPreConfig_Copy(_PyPreConfig *config, const _PyPreConfig *config2)
     COPY_ATTR(legacy_windows_fs_encoding);
 #endif
     COPY_ATTR(utf8_mode);
-    COPY_STR_ATTR(allocator);
+    COPY_ATTR(allocator);
 
 #undef COPY_ATTR
 #undef COPY_STR_ATTR
@@ -341,7 +331,7 @@ _PyPreConfig_AsDict(const _PyPreConfig *config)
     SET_ITEM_INT(legacy_windows_fs_encoding);
 #endif
     SET_ITEM_INT(dev_mode);
-    SET_ITEM_STR(allocator);
+    SET_ITEM_INT(allocator);
     return dict;
 
 fail:
@@ -616,25 +606,21 @@ preconfig_init_coerce_c_locale(_PyPreConfig *config)
 static _PyInitError
 preconfig_init_allocator(_PyPreConfig *config)
 {
-    if (config->allocator == NULL) {
+    if (config->allocator == PYMEM_ALLOCATOR_NOT_SET) {
         /* bpo-34247. The PYTHONMALLOC environment variable has the priority
            over PYTHONDEV env var and "-X dev" command line option.
            For example, PYTHONMALLOC=malloc PYTHONDEVMODE=1 sets the memory
            allocators to "malloc" (and not to "debug"). */
-        const char *allocator = _Py_GetEnv(config->use_environment, "PYTHONMALLOC");
-        if (allocator) {
-            config->allocator = _PyMem_RawStrdup(allocator);
-            if (config->allocator == NULL) {
-                return _Py_INIT_NO_MEMORY();
+        const char *envvar = _Py_GetEnv(config->use_environment, "PYTHONMALLOC");
+        if (envvar) {
+            if (_PyMem_GetAllocatorName(envvar, &config->allocator) < 0) {
+                return _Py_INIT_ERR("PYTHONMALLOC: unknown allocator");
             }
         }
     }
 
-    if (config->dev_mode && config->allocator == NULL) {
-        config->allocator = _PyMem_RawStrdup("debug");
-        if (config->allocator == NULL) {
-            return _Py_INIT_NO_MEMORY();
-        }
+    if (config->dev_mode && config->allocator == PYMEM_ALLOCATOR_NOT_SET) {
+        config->allocator = PYMEM_ALLOCATOR_DEBUG;
     }
     return _Py_INIT_OK();
 }
@@ -815,7 +801,6 @@ done:
         setlocale(LC_CTYPE, init_ctype_locale);
         PyMem_RawFree(init_ctype_locale);
     }
-    _PyPreConfig_Clear(&save_config);
     Py_UTF8Mode = init_utf8_mode ;
 #ifdef MS_WINDOWS
     Py_LegacyWindowsFSEncodingFlag = init_legacy_encoding;
@@ -825,40 +810,6 @@ done:
 }
 
 
-static _PyInitError
-_PyPreConfig_SetAllocator(_PyPreConfig *config)
-{
-    assert(!_PyRuntime.core_initialized);
-
-    PyMemAllocatorEx old_alloc;
-    PyMem_GetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-
-    if (_PyMem_SetupAllocators(config->allocator) < 0) {
-        return _Py_INIT_ERR("Unknown PYTHONMALLOC allocator");
-    }
-
-    /* Copy the pre-configuration with the new allocator */
-    _PyPreConfig config2 = _PyPreConfig_INIT;
-    if (_PyPreConfig_Copy(&config2, config) < 0) {
-        _PyPreConfig_Clear(&config2);
-        PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-        return _Py_INIT_NO_MEMORY();
-    }
-
-    /* Free the old config and replace config with config2. Since config now
-       owns the data, don't free config2. */
-    PyMemAllocatorEx new_alloc;
-    PyMem_GetAllocator(PYMEM_DOMAIN_RAW, &new_alloc);
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-    _PyPreConfig_Clear(config);
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &new_alloc);
-
-    *config = config2;
-
-    return _Py_INIT_OK();
-}
-
-
 /* Write the pre-configuration:
 
    - set the memory allocators
@@ -872,7 +823,7 @@ _PyPreConfig_SetAllocator(_PyPreConfig *config)
    Do nothing if called after Py_Initialize(): ignore the new
    pre-configuration. */
 _PyInitError
-_PyPreConfig_Write(_PyPreConfig *config)
+_PyPreConfig_Write(const _PyPreConfig *config)
 {
     if (_PyRuntime.core_initialized) {
         /* bpo-34008: Calling this functions after Py_Initialize() ignores
@@ -880,10 +831,9 @@ _PyPreConfig_Write(_PyPreConfig *config)
         return _Py_INIT_OK();
     }
 
-    if (config->allocator != NULL) {
-        _PyInitError err = _PyPreConfig_SetAllocator(config);
-        if (_Py_INIT_FAILED(err)) {
-            return err;
+    if (config->allocator != PYMEM_ALLOCATOR_NOT_SET) {
+        if (_PyMem_SetupAllocators(config->allocator) < 0) {
+            return _Py_INIT_ERR("Unknown PYTHONMALLOC allocator");
         }
     }
 
diff --git a/Python/pylifecycle.c b/Python/pylifecycle.c
index e89152637fe..eecb439a11d 100644
--- a/Python/pylifecycle.c
+++ b/Python/pylifecycle.c
@@ -706,26 +706,22 @@ _Py_PreInitializeFromPyArgv(const _PyPreConfig *src_config, const _PyArgv *args)
     if (src_config) {
         if (_PyPreConfig_Copy(&config, src_config) < 0) {
             err = _Py_INIT_NO_MEMORY();
-            goto done;
+            return err;
         }
     }
 
     err = _PyPreConfig_Read(&config, args);
     if (_Py_INIT_FAILED(err)) {
-        goto done;
+        return err;
     }
 
     err = _PyPreConfig_Write(&config);
     if (_Py_INIT_FAILED(err)) {
-        goto done;
+        return err;
     }
 
     runtime->pre_initialized = 1;
-    err = _Py_INIT_OK();
-
-done:
-    _PyPreConfig_Clear(&config);
-    return err;
+    return _Py_INIT_OK();
 }
 
 
diff --git a/Python/pystate.c b/Python/pystate.c
index 67315756ab7..8c906ce87ad 100644
--- a/Python/pystate.c
+++ b/Python/pystate.c
@@ -106,8 +106,6 @@ _PyRuntimeState_Fini(_PyRuntimeState *runtime)
         runtime->xidregistry.mutex = NULL;
     }
 
-    _PyPreConfig_Clear(&runtime->preconfig);
-
     PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
 }