bpo-36142: Add _PyPreConfig_SetAllocator() (GH-12187)

* _PyPreConfig_Write() now reallocates the pre-configuration with the
  new memory allocator.
* It is no longer needed to force the "default raw memory allocator"
  to clear pre-configuration and core configuration. Simplify the
  code.
* _PyPreConfig_Write() now does nothing if called after
  Py_Initialize(): no longer check if the allocator is the same.
* Remove _PyMem_GetDebugAllocatorsName(): dev mode sets again
  allocator to "debug".

2 files changed, 45 insertions, 45 deletions

diff --git a/Python/preconfig.c b/Python/preconfig.c
index 6924203215..ee9dca4bdc 100644
--- a/Python/preconfig.c
+++ b/Python/preconfig.c
@@ -125,15 +125,8 @@ precmdline_clear(_PyPreCmdline *cmdline)
 void
 _PyPreConfig_Clear(_PyPreConfig *config)
 {
-#define CLEAR(ATTR) \
-    do { \
-        PyMem_RawFree(ATTR); \
-        ATTR = NULL; \
-    } while (0)
-
-    CLEAR(config->allocator);
-
-#undef CLEAR
+    PyMem_RawFree(config->allocator);
+    config->allocator = NULL;
 }
 
 
@@ -453,8 +446,7 @@ preconfig_read(_PyPreConfig *config, const _PyPreCmdline *cmdline)
 
     /* allocator */
     if (config->dev_mode && config->allocator == NULL) {
-        const char *allocator = _PyMem_GetDebugAllocatorsName();
-        config->allocator = _PyMem_RawStrdup(allocator);
+        config->allocator = _PyMem_RawStrdup("debug");
         if (config->allocator == NULL) {
             return _Py_INIT_NO_MEMORY();
         }
@@ -742,31 +734,56 @@ done:
 
 
 static _PyInitError
-_PyPreConfig_Reconfigure(const _PyPreConfig *config)
+_PyPreConfig_SetAllocator(_PyPreConfig *config)
 {
-    if (config->allocator != NULL) {
-        const char *allocator = _PyMem_GetAllocatorsName();
-        if (allocator == NULL || strcmp(config->allocator, allocator) != 0) {
-            return _Py_INIT_USER_ERR("cannot modify memory allocator "
-                                     "after first Py_Initialize()");
-        }
+    assert(!_PyRuntime.core_initialized);
+
+    PyMemAllocatorEx old_alloc;
+    PyMem_GetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
+
+    if (_PyMem_SetupAllocators(config->allocator) < 0) {
+        return _Py_INIT_USER_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.
+
+   If the memory allocator is changed, config is re-allocated with new
+   allocator. So calling _PyPreConfig_Clear(config) is safe after this call. */
 _PyInitError
-_PyPreConfig_Write(const _PyPreConfig *config)
+_PyPreConfig_Write(_PyPreConfig *config)
 {
     if (_PyRuntime.core_initialized) {
         /* bpo-34008: Calling Py_Main() after Py_Initialize() ignores
            the new configuration. */
-        return _PyPreConfig_Reconfigure(config);
+        return _Py_INIT_OK();
     }
 
     if (config->allocator != NULL) {
-        if (_PyMem_SetupAllocators(config->allocator) < 0) {
-            return _Py_INIT_USER_ERR("Unknown PYTHONMALLOC allocator");
+        _PyInitError err = _PyPreConfig_SetAllocator(config);
+        if (_Py_INIT_FAILED(err)) {
+            return err;
         }
     }
 
diff --git a/Python/pylifecycle.c b/Python/pylifecycle.c
index f72af86384..522a4275a5 100644
--- a/Python/pylifecycle.c
+++ b/Python/pylifecycle.c
@@ -716,21 +716,14 @@ _Py_InitializeCore_impl(PyInterpreterState **interp_p,
 static _PyInitError
 pyinit_preconfig(_PyPreConfig *preconfig, const _PyPreConfig *src_preconfig)
 {
-    _PyInitError err;
-    PyMemAllocatorEx old_alloc;
-
     /* Set LC_CTYPE to the user preferred locale */
     _Py_SetLocaleFromEnv(LC_CTYPE);
 
-    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-    if (_PyPreConfig_Copy(preconfig, src_preconfig) >= 0) {
-        err = _PyPreConfig_Read(preconfig);
+    if (_PyPreConfig_Copy(preconfig, src_preconfig) < 0) {
+        return _Py_INIT_ERR("failed to copy pre config");
     }
-    else {
-        err = _Py_INIT_ERR("failed to copy pre config");
-    }
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
 
+    _PyInitError err = _PyPreConfig_Read(preconfig);
     if (_Py_INIT_FAILED(err)) {
         return err;
     }
@@ -743,21 +736,15 @@ static _PyInitError
 pyinit_coreconfig(_PyCoreConfig *config, const _PyCoreConfig *src_config,
                   PyInterpreterState **interp_p)
 {
-    PyMemAllocatorEx old_alloc;
-    _PyInitError err;
 
     /* Set LC_CTYPE to the user preferred locale */
     _Py_SetLocaleFromEnv(LC_CTYPE);
 
-    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-    if (_PyCoreConfig_Copy(config, src_config) >= 0) {
-        err = _PyCoreConfig_Read(config, NULL);
-    }
-    else {
-        err = _Py_INIT_ERR("failed to copy core config");
+    if (_PyCoreConfig_Copy(config, src_config) < 0) {
+        return _Py_INIT_ERR("failed to copy core config");
     }
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
 
+    _PyInitError err = _PyCoreConfig_Read(config, NULL);
     if (_Py_INIT_FAILED(err)) {
         return err;
     }
@@ -792,7 +779,6 @@ _PyInitError
 _Py_InitializeCore(PyInterpreterState **interp_p,
                    const _PyCoreConfig *src_config)
 {
-    PyMemAllocatorEx old_alloc;
     _PyInitError err;
 
     assert(src_config != NULL);
@@ -807,10 +793,7 @@ _Py_InitializeCore(PyInterpreterState **interp_p,
     err = pyinit_coreconfig(&local_config, src_config, interp_p);
 
 done:
-    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
     _PyCoreConfig_Clear(&local_config);
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-
     return err;
 }