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".

diff --git a/Include/internal/pycore_coreconfig.h b/Include/internal/pycore_coreconfig.h
index b34416bfd7d9167c9e24c1f54d272f8ee43926e5..0917a6ac2a34c068ce9a6aad22dd58c6b3cd55cc 100644 (file)
--- a/Include/internal/pycore_coreconfig.h
+++ b/Include/internal/pycore_coreconfig.h
@@ -59,7 +59,7 @@ PyAPI_FUNC(int) _PyPreConfig_AsDict(const _PyPreConfig *config,
     PyObject *dict);
 PyAPI_FUNC(_PyInitError) _PyPreConfig_ReadFromArgv(_PyPreConfig *config,
     const _PyArgv *args);
-PyAPI_FUNC(_PyInitError) _PyPreConfig_Write(const _PyPreConfig *config);
+PyAPI_FUNC(_PyInitError) _PyPreConfig_Write(_PyPreConfig *config);
 
 
 /* --- _PyCoreConfig ---------------------------------------------- */

diff --git a/Include/internal/pycore_pymem.h b/Include/internal/pycore_pymem.h
index fedc7cc119bfb5d11ad4dca9f268d06976c13f6e..1e7da87cd75c0aa9e4b7adb79634c5781a00bb7d 100644 (file)
--- a/Include/internal/pycore_pymem.h
+++ b/Include/internal/pycore_pymem.h
@@ -155,8 +155,6 @@ PyAPI_FUNC(int) _PyMem_SetDefaultAllocator(
     PyMemAllocatorDomain domain,
     PyMemAllocatorEx *old_alloc);
 
-PyAPI_FUNC(const char*) _PyMem_GetDebugAllocatorsName(void);
-
 #ifdef __cplusplus
 }
 #endif

diff --git a/Lib/test/test_embed.py b/Lib/test/test_embed.py
index 2827e8708e114d34b639b2237f0bb77cd766e179..1f236a985ad6a25f1b19dbbbfacd8cd9e2528cd7 100644 (file)
--- a/Lib/test/test_embed.py
+++ b/Lib/test/test_embed.py
@@ -336,7 +336,6 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
             'legacy_windows_fs_encoding': 0,
             'legacy_windows_stdio': 0,
         })
-    DEBUG_ALLOCATOR = 'pymalloc_debug' if support.with_pymalloc() else 'malloc_debug'
 
     # main config
     COPY_MAIN_CONFIG = (
@@ -589,7 +588,7 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
 
     def test_init_env_dev_mode(self):
         config = dict(self.INIT_ENV_CONFIG,
-                      allocator=self.DEBUG_ALLOCATOR,
+                      allocator='debug',
                       dev_mode=1)
         self.check_config("init_env_dev_mode", config)
 
@@ -597,7 +596,7 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
         config = {
             'dev_mode': 1,
             'faulthandler': 1,
-            'allocator': self.DEBUG_ALLOCATOR,
+            'allocator': 'debug',
         }
         self.check_config("init_dev_mode", config)
 

diff --git a/Modules/main.c b/Modules/main.c
index 9a2347e2b277247c51d24ee9f624c034779b8b95..14055c8101c8dc6591f1f2af628e50817920381b 100644 (file)
--- a/Modules/main.c
+++ b/Modules/main.c
@@ -289,17 +289,9 @@ _PyMainInterpreterConfig_Read(_PyMainInterpreterConfig *main_config,
 static _PyInitError
 preconfig_read_write(_PyPreConfig *config, const _PyArgv *args)
 {
-    _PyInitError err;
-
-    PyMemAllocatorEx old_alloc;
-    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-
     _PyPreConfig_GetGlobalConfig(config);
 
-    err = _PyPreConfig_ReadFromArgv(config, args);
-
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-
+    _PyInitError err = _PyPreConfig_ReadFromArgv(config, args);
     if (_Py_INIT_FAILED(err)) {
         return err;
     }
@@ -312,17 +304,9 @@ static _PyInitError
 config_read_write(_PyCoreConfig *config, const _PyArgv *args,
                   const _PyPreConfig *preconfig)
 {
-    _PyInitError err;
-
-    PyMemAllocatorEx old_alloc;
-    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-
     _PyCoreConfig_GetGlobalConfig(config);
 
-    err = _PyCoreConfig_ReadFromArgv(config, args, preconfig);
-
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-
+    _PyInitError err = _PyCoreConfig_ReadFromArgv(config, args, preconfig);
     if (_Py_INIT_FAILED(err)) {
         return err;
     }
@@ -355,7 +339,6 @@ static _PyInitError
 pymain_init(const _PyArgv *args, PyInterpreterState **interp_p)
 {
     _PyInitError err;
-    PyMemAllocatorEx old_alloc;
 
     err = _PyRuntime_Initialize();
     if (_Py_INIT_FAILED(err)) {
@@ -402,12 +385,8 @@ pymain_init(const _PyArgv *args, PyInterpreterState **interp_p)
     err = _Py_INIT_OK();
 
 done:
-    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-
     _PyPreConfig_Clear(preconfig);
     _PyCoreConfig_Clear(config);
-
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
     return err;
 }
 

diff --git a/Objects/obmalloc.c b/Objects/obmalloc.c
index 1afbcca14f28aa7039135c336714e0c0cc3f38e5..1c2a32050f9381e0a8aab3591a42e191e8401644 100644 (file)
--- a/Objects/obmalloc.c
+++ b/Objects/obmalloc.c
@@ -221,20 +221,6 @@ static PyMemAllocatorEx _PyObject = PYOBJ_ALLOC;
 #endif
 
 
-/* Get the effective name of "debug" memory allocators,
-   as if _PyMem_GetAllocatorsName() is called after
-   _PyMem_SetupAllocators("debug"). */
-const char*
-_PyMem_GetDebugAllocatorsName(void)
-{
-#ifdef WITH_PYMALLOC
-    return "pymalloc_debug";
-#else
-    return "malloc_debug";
-#endif
-}
-
-
 static int
 pymem_set_default_allocator(PyMemAllocatorDomain domain, int debug,
                             PyMemAllocatorEx *old_alloc)

diff --git a/Python/preconfig.c b/Python/preconfig.c
index 69242032151ab60f05d9993957a6660e0384ce19..ee9dca4bdc924126feea57c98d56ba7feb2fd1ba 100644 (file)
--- 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 f72af8638402f61dd3a5a4ead7d25e2a0985d398..522a4275a51d91c228b3f4c5f2ef5790b8804355 100644 (file)
--- 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;
 }