Add _PyInterpreterConfig.use_main_obmalloc.

python · ericsnowcurrently · Apr 24, 2023 · Oct 6, 2022 · Oct 7, 2022 · Feb 7, 2023
commit 4dc087dea276b3f5c0b8b98b79614661c777971d
@@ -244,6 +244,7 @@ PyAPI_FUNC(PyStatus) PyConfig_SetWideStringList(PyConfig *config,
 /* --- PyInterpreterConfig ------------------------------------ */
 
 typedef struct {
+    int use_main_obmalloc;
     int allow_fork;
     int allow_exec;
     int allow_threads;
@@ -253,6 +254,7 @@ typedef struct {
 
 #define _PyInterpreterConfig_INIT \
     { \
+        .use_main_obmalloc = 0, \
         .allow_fork = 0, \
         .allow_exec = 0, \
         .allow_threads = 1, \
@@ -262,6 +264,7 @@ typedef struct {
 
 #define _PyInterpreterConfig_LEGACY_INIT \
     { \
+        .use_main_obmalloc = 1, \
         .allow_fork = 1, \
         .allow_exec = 1, \
         .allow_threads = 1, \

@@ -11,6 +11,10 @@ is available in a given context.  For example, forking the process
 might not be allowed in the current interpreter (i.e. os.fork() would fail).
 */
 
+/* Set if the interpreter share obmalloc runtime state
+   with the main interpreter. */
+#define Py_RTFLAGS_USE_MAIN_OBMALLOC (1UL << 5)
+
 /* Set if import should check a module for subinterpreter support. */
 #define Py_RTFLAGS_MULTI_INTERP_EXTENSIONS (1UL << 8)
 

@@ -1215,20 +1215,23 @@ def test_configured_settings(self):
         """
         import json
 
+        OBMALLOC = 1<<5
         EXTENSIONS = 1<<8
         THREADS = 1<<10
         DAEMON_THREADS = 1<<11
         FORK = 1<<15
         EXEC = 1<<16
 
-        features = ['fork', 'exec', 'threads', 'daemon_threads', 'extensions']
+        features = ['obmalloc', 'fork', 'exec', 'threads', 'daemon_threads',
+                    'extensions']
         kwlist = [f'allow_{n}' for n in features]
+        kwlist[0] = 'use_main_obmalloc'
         kwlist[-1] = 'check_multi_interp_extensions'
         for config, expected in {
-            (True, True, True, True, True):
-                FORK | EXEC | THREADS | DAEMON_THREADS | EXTENSIONS,
-            (False, False, False, False, False): 0,
-            (False, False, True, False, True): THREADS | EXTENSIONS,
+            (True, True, True, True, True, True):
+                OBMALLOC | FORK | EXEC | THREADS | DAEMON_THREADS | EXTENSIONS,
+            (False, False, False, False, False, False): 0,
+            (False, False, False, True, False, True): THREADS | EXTENSIONS,
         }.items():
             kwargs = dict(zip(kwlist, config))
             expected = {
@@ -1261,13 +1264,15 @@ def test_overridden_setting_extensions_subinterp_check(self):
         """
         import json
 
+        OBMALLOC = 1<<5
         EXTENSIONS = 1<<8
         THREADS = 1<<10
         DAEMON_THREADS = 1<<11
         FORK = 1<<15
         EXEC = 1<<16
-        BASE_FLAGS = FORK | EXEC | THREADS | DAEMON_THREADS
+        BASE_FLAGS = OBMALLOC | FORK | EXEC | THREADS | DAEMON_THREADS
         base_kwargs = {
+            'use_main_obmalloc': True,
             'allow_fork': True,
             'allow_exec': True,
             'allow_threads': True,

diff --git a/Lib/test/test_embed.py b/Lib/test/test_embed.py
@@ -1656,6 +1656,7 @@ def test_init_use_frozen_modules(self):
                                        api=API_PYTHON, env=env)
 
     def test_init_main_interpreter_settings(self):
+        OBMALLOC = 1<<5
         EXTENSIONS = 1<<8
         THREADS = 1<<10
         DAEMON_THREADS = 1<<11
@@ -1664,7 +1665,7 @@ def test_init_main_interpreter_settings(self):
         expected = {
             # All optional features should be enabled.
             'feature_flags':
-                FORK | EXEC | THREADS | DAEMON_THREADS,
+                OBMALLOC | FORK | EXEC | THREADS | DAEMON_THREADS,
         }
         out, err = self.run_embedded_interpreter(
             'test_init_main_interpreter_settings',

diff --git a/Lib/test/test_threading.py b/Lib/test/test_threading.py
@@ -1343,6 +1343,7 @@ def func():
             import test.support
             test.support.run_in_subinterp_with_config(
                 {subinterp_code!r},
+                use_main_obmalloc=True,
                 allow_fork=True,
                 allow_exec=True,
                 allow_threads={allowed},

diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
@@ -1482,6 +1482,7 @@ static PyObject *
 run_in_subinterp_with_config(PyObject *self, PyObject *args, PyObject *kwargs)
 {
     const char *code;
+    int use_main_obmalloc = -1;
     int allow_fork = -1;
     int allow_exec = -1;
     int allow_threads = -1;
@@ -1493,19 +1494,25 @@ run_in_subinterp_with_config(PyObject *self, PyObject *args, PyObject *kwargs)
     PyCompilerFlags cflags = {0};
 
     static char *kwlist[] = {"code",
+                             "use_main_obmalloc",
                              "allow_fork",
                              "allow_exec",
                              "allow_threads",
                              "allow_daemon_threads",
                              "check_multi_interp_extensions",
                              NULL};
     if (!PyArg_ParseTupleAndKeywords(args, kwargs,
-                    "s$ppppp:run_in_subinterp_with_config", kwlist,
-                    &code, &allow_fork, &allow_exec,
+                    "s$pppppp:run_in_subinterp_with_config", kwlist,
+                    &code, &use_main_obmalloc,
+                    &allow_fork, &allow_exec,
                     &allow_threads, &allow_daemon_threads,
                     &check_multi_interp_extensions)) {
         return NULL;
     }
+    if (use_main_obmalloc < 0) {
+        PyErr_SetString(PyExc_ValueError, "missing use_main_obmalloc");
+        return NULL;
+    }
     if (allow_fork < 0) {
         PyErr_SetString(PyExc_ValueError, "missing allow_fork");
         return NULL;
@@ -1532,6 +1539,7 @@ run_in_subinterp_with_config(PyObject *self, PyObject *args, PyObject *kwargs)
     PyThreadState_Swap(NULL);
 
     const _PyInterpreterConfig config = {
+        .use_main_obmalloc = use_main_obmalloc,
         .allow_fork = allow_fork,
         .allow_exec = allow_exec,
         .allow_threads = allow_threads,

@@ -546,11 +546,19 @@ pycore_init_runtime(_PyRuntimeState *runtime,
 }
 
 
-static void
+static PyStatus
 init_interp_settings(PyInterpreterState *interp, const _PyInterpreterConfig *config)
 {
     assert(interp->feature_flags == 0);
 
+    if (config->use_main_obmalloc) {
+        interp->feature_flags |= Py_RTFLAGS_USE_MAIN_OBMALLOC;
+    }
+    else if (!config->check_multi_interp_extensions) {
+        return _PyStatus_ERR("per-interpreter obmalloc does not support "
+                             "single-phase init extension modules");
+    }
+
     if (config->allow_fork) {
         interp->feature_flags |= Py_RTFLAGS_FORK;
     }
@@ -569,6 +577,8 @@ init_interp_settings(PyInterpreterState *interp, const _PyInterpreterConfig *con
     if (config->check_multi_interp_extensions) {
         interp->feature_flags |= Py_RTFLAGS_MULTI_INTERP_EXTENSIONS;
     }
+
+    return _PyStatus_OK();
 }
 
 
@@ -621,7 +631,10 @@ pycore_create_interpreter(_PyRuntimeState *runtime,
     }
 
     const _PyInterpreterConfig config = _PyInterpreterConfig_LEGACY_INIT;
-    init_interp_settings(interp, &config);
+    status = init_interp_settings(interp, &config);
+    if (_PyStatus_EXCEPTION(status)) {
+        return status;
+    }
 
     PyThreadState *tstate = _PyThreadState_New(interp);
     if (tstate == NULL) {
@@ -2031,7 +2044,10 @@ new_interpreter(PyThreadState **tstate_p, const _PyInterpreterConfig *config)
         goto error;
     }
 
-    init_interp_settings(interp, config);
+    status = init_interp_settings(interp, config);
+    if (_PyStatus_EXCEPTION(status)) {
+        goto error;
+    }
 
     status = init_interp_create_gil(tstate);
     if (_PyStatus_EXCEPTION(status)) {