Add _PyXI_NewInterpreter() and _PyXI_EndInterpreter().

python · ericsnowcurrently · Apr 11, 2024 · Apr 8, 2024 · Apr 5, 2024 · Apr 5, 2024
commit 528c6e328d110ab81144b953ebf963c09d23f71a
@@ -313,6 +313,21 @@ PyAPI_FUNC(PyObject *) _PyXI_ApplyCapturedException(_PyXI_session *session);
 PyAPI_FUNC(int) _PyXI_HasCapturedException(_PyXI_session *session);
 
 
+/*************/
+/* other API */
+/*************/
+
+// Export for _testinternalcapi shared extension
+PyAPI_FUNC(PyInterpreterState *) _PyXI_NewInterpreter(
+    PyInterpreterConfig *config,
+    PyThreadState **p_tstate,
+    PyThreadState **p_save_tstate);
+PyAPI_FUNC(void) _PyXI_EndInterpreter(
+    PyInterpreterState *interp,
+    PyThreadState *tstate,
+    PyThreadState **p_save_tstate);
+
+
 #ifdef __cplusplus
 }
 #endif

diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c
@@ -1357,7 +1357,88 @@ dict_getitem_knownhash(PyObject *self, PyObject *args)
 }
 
 
-/* To run some code in a sub-interpreter. */
+static PyInterpreterState *
+_new_interpreter(PyObject *configobj,
+                 PyThreadState **p_tstate, PyThreadState **p_save_tstate)
+{
+    PyInterpreterConfig config;
+    if (configobj == NULL) {
+        config = (PyInterpreterConfig)_PyInterpreterConfig_INIT;
+    }
+    else {
+        PyObject *dict = PyObject_GetAttrString(configobj, "__dict__");
+        if (dict == NULL) {
+            PyErr_Format(PyExc_TypeError, "bad config %R", configobj);
+            return NULL;
+        }
+        int res = _PyInterpreterConfig_InitFromDict(&config, dict);
+        Py_DECREF(dict);
+        if (res < 0) {
+            return NULL;
+        }
+    }
+
+    return _PyXI_NewInterpreter(&config, p_tstate, p_save_tstate);
+}
+
+// This exists mostly for testing the _interpreters module, as an
+// alternative to _interpreters.create()
+static PyObject *
+create_interpreter(PyObject *self, PyObject *args, PyObject *kwargs)
+{
+    static char *kwlist[] = {"config", NULL};
+    PyObject *configobj = NULL;
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs,
+                                     "|O:create_interpreter", kwlist,
+                                     &configobj))
+    {
+        return NULL;
+    }
+
+    PyInterpreterState *interp = _new_interpreter(configobj, NULL, NULL);
+    if (interp == NULL) {
+        return NULL;
+    }
+
+    PyObject *idobj = _PyInterpreterState_GetIDObject(interp);
+    if (idobj == NULL) {
+        _PyXI_EndInterpreter(interp, NULL, NULL);
+        return NULL;
+    }
+
+    return idobj;
+}
+
+// This exists mostly for testing the _interpreters module, as an
+// alternative to _interpreters.destroy()
+static PyObject *
+destroy_interpreter(PyObject *self, PyObject *args, PyObject *kwargs)
+{
+    static char *kwlist[] = {"id", NULL};
+    PyObject *idobj = NULL;
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs,
+                                     "O:destroy_interpreter", kwlist,
+                                     &idobj))
+    {
+        return NULL;
+    }
+
+    PyInterpreterState *interp = _PyInterpreterState_LookUpIDObject(idobj);
+    if (interp == NULL) {
+        return NULL;
+    }
+
+    _PyXI_EndInterpreter(interp, NULL, NULL);
+    Py_RETURN_NONE;
+}
+
+
+/* To run some code in a sub-interpreter.
+
+Generally you can use test.support.interpreters,
+but we keep this helper as a distinct implementation.
+That's especially important for testing test.support.interpreters.
+*/
 static PyObject *
 run_in_subinterp_with_config(PyObject *self, PyObject *args, PyObject *kwargs)
 {
@@ -1371,42 +1452,18 @@ run_in_subinterp_with_config(PyObject *self, PyObject *args, PyObject *kwargs)
         return NULL;
     }
 
-    PyInterpreterConfig config;
-    PyObject *dict = PyObject_GetAttrString(configobj, "__dict__");
-    if (dict == NULL) {
-        PyErr_Format(PyExc_TypeError, "bad config %R", configobj);
-        return NULL;
-    }
-    int res = _PyInterpreterConfig_InitFromDict(&config, dict);
-    Py_DECREF(dict);
-    if (res < 0) {
+    PyThreadState *save_tstate;
+    PyThreadState *substate;
+    PyInterpreterState *interp = _new_interpreter(configobj, &substate, &save_tstate);
+    if (interp == NULL) {
         return NULL;
     }
 
-    PyThreadState *mainstate = PyThreadState_Get();
-
-    PyThreadState_Swap(NULL);
-
-    PyThreadState *substate;
-    PyStatus status = Py_NewInterpreterFromConfig(&substate, &config);
-    if (PyStatus_Exception(status)) {
-        /* Since no new thread state was created, there is no exception to
-           propagate; raise a fresh one after swapping in the old thread
-           state. */
-        PyThreadState_Swap(mainstate);
-        _PyErr_SetFromPyStatus(status);
-        PyObject *exc = PyErr_GetRaisedException();
-        PyErr_SetString(PyExc_RuntimeError, "sub-interpreter creation failed");
-        _PyErr_ChainExceptions1(exc);
-        return NULL;
-    }
-    assert(substate != NULL);
     /* only initialise 'cflags.cf_flags' to test backwards compatibility */
     PyCompilerFlags cflags = {0};
     int r = PyRun_SimpleStringFlags(code, &cflags);
-    Py_EndInterpreter(substate);
 
-    PyThreadState_Swap(mainstate);
+    _PyXI_EndInterpreter(interp, substate, &save_tstate);
 
     return PyLong_FromLong(r);
 }
@@ -1738,6 +1795,10 @@ static PyMethodDef module_functions[] = {
     {"get_object_dict_values", get_object_dict_values, METH_O},
     {"hamt", new_hamt, METH_NOARGS},
     {"dict_getitem_knownhash",  dict_getitem_knownhash,          METH_VARARGS},
+    {"create_interpreter", _PyCFunction_CAST(create_interpreter),
+     METH_VARARGS | METH_KEYWORDS},
+    {"destroy_interpreter", _PyCFunction_CAST(destroy_interpreter),
+     METH_VARARGS | METH_KEYWORDS},
     {"run_in_subinterp_with_config",
      _PyCFunction_CAST(run_in_subinterp_with_config),
      METH_VARARGS | METH_KEYWORDS},

diff --git a/Modules/_xxsubinterpretersmodule.c b/Modules/_xxsubinterpretersmodule.c
@@ -423,55 +423,6 @@ config_from_object(PyObject *configobj, PyInterpreterConfig *config)
 }
 
 
-static PyInterpreterState *
-new_interpreter(PyInterpreterConfig *config, PyObject **p_idobj,  PyThreadState **p_tstate)
-{
-    PyThreadState *save_tstate = PyThreadState_Get();
-    assert(save_tstate != NULL);
-    PyThreadState *tstate = NULL;
-    // XXX Possible GILState issues?
-    PyStatus status = Py_NewInterpreterFromConfig(&tstate, config);
-    PyThreadState_Swap(save_tstate);
-    if (PyStatus_Exception(status)) {
-        /* Since no new thread state was created, there is no exception to
-           propagate; raise a fresh one after swapping in the old thread
-           state. */
-        _PyErr_SetFromPyStatus(status);
-        return NULL;
-    }
-    assert(tstate != NULL);
-    PyInterpreterState *interp = PyThreadState_GetInterpreter(tstate);
-
-    if (p_idobj != NULL) {
-        // We create the object using the original interpreter.
-        PyObject *idobj = _PyInterpreterState_GetIDObject(interp);
-        if (idobj == NULL) {
-            goto error;
-        }
-        *p_idobj = idobj;
-    }
-
-    if (p_tstate != NULL) {
-        *p_tstate = tstate;
-    }
-    else {
-        PyThreadState_Swap(tstate);
-        PyThreadState_Clear(tstate);
-        PyThreadState_Swap(save_tstate);
-        PyThreadState_Delete(tstate);
-    }
-
-    return interp;
-
-error:
-    // XXX Possible GILState issues?
-    save_tstate = PyThreadState_Swap(tstate);
-    Py_EndInterpreter(tstate);
-    PyThreadState_Swap(save_tstate);
-    return NULL;
-}
-
-
 static int
 _run_script(PyObject *ns, const char *codestr, Py_ssize_t codestrlen, int flags)
 {
@@ -600,8 +551,7 @@ interp_create(PyObject *self, PyObject *args, PyObject *kwds)
         return NULL;
     }
 
-    PyObject *idobj = NULL;
-    PyInterpreterState *interp = new_interpreter(&config, &idobj, NULL);
+    PyInterpreterState *interp = _PyXI_NewInterpreter(&config, NULL, NULL);
     if (interp == NULL) {
         // XXX Move the chained exception to interpreters.create()?
         PyObject *exc = PyErr_GetRaisedException();
@@ -611,6 +561,12 @@ interp_create(PyObject *self, PyObject *args, PyObject *kwds)
         return NULL;
     }
 
+    PyObject *idobj = _PyInterpreterState_GetIDObject(interp);
+    if (idobj == NULL) {
+        _PyXI_EndInterpreter(interp, NULL, NULL);
+        return NULL;
+    }
+
     if (reqrefs) {
         // Decref to 0 will destroy the interpreter.
         _PyInterpreterState_RequireIDRef(interp, 1);
@@ -672,12 +628,7 @@ interp_destroy(PyObject *self, PyObject *args, PyObject *kwds)
     }
 
     // Destroy the interpreter.
-    PyThreadState *tstate = PyThreadState_New(interp);
-    _PyThreadState_SetWhence(tstate, _PyThreadState_WHENCE_INTERP);
-    // XXX Possible GILState issues?
-    PyThreadState *save_tstate = PyThreadState_Swap(tstate);
-    Py_EndInterpreter(tstate);
-    PyThreadState_Swap(save_tstate);
+    _PyXI_EndInterpreter(interp, NULL, NULL);
 
     Py_RETURN_NONE;
 }

@@ -1682,3 +1682,82 @@ _PyXI_FiniTypes(PyInterpreterState *interp)
 {
     fini_exceptions(interp);
 }
+
+
+/*************/
+/* other API */
+/*************/
+
+PyInterpreterState *
+_PyXI_NewInterpreter(PyInterpreterConfig *config,
+                     PyThreadState **p_tstate, PyThreadState **p_save_tstate)
+{
+    PyThreadState *save_tstate = PyThreadState_Swap(NULL);
+    assert(save_tstate != NULL);
+
+    PyThreadState *tstate;
+    PyStatus status = Py_NewInterpreterFromConfig(&tstate, config);
+    if (PyStatus_Exception(status)) {
+        // Since no new thread state was created, there is no exception
+        // to propagate; raise a fresh one after swapping back in the
+        // old thread state.
+        PyThreadState_Swap(save_tstate);
+        _PyErr_SetFromPyStatus(status);
+        PyObject *exc = PyErr_GetRaisedException();
+        PyErr_SetString(PyExc_InterpreterError,
+                        "sub-interpreter creation failed");
+        _PyErr_ChainExceptions1(exc);
+        return NULL;
+    }
+    assert(tstate != NULL);
+    PyInterpreterState *interp = PyThreadState_GetInterpreter(tstate);
+
+    if (p_tstate != NULL) {
+        // We leave the new thread state as the current one.
+        *p_tstate = tstate;
+    }
+    else {
+        // Throw away the initial tstate.
+        PyThreadState_Clear(tstate);
+        PyThreadState_Swap(save_tstate);
+        PyThreadState_Delete(tstate);
+        save_tstate = NULL;
+    }
+    if (p_save_tstate != NULL) {
+        *p_save_tstate = save_tstate;
+    }
+    return interp;
+}
+
+void
+_PyXI_EndInterpreter(PyInterpreterState *interp,
+                     PyThreadState *tstate, PyThreadState **p_save_tstate)
+{
+    PyThreadState *cur_tstate = PyThreadState_GET();
+    PyThreadState *save_tstate = NULL;
+    if (tstate == NULL) {
+        if (PyThreadState_GetInterpreter(cur_tstate) == interp) {
+            tstate = cur_tstate;
+        }
+        else {
+            tstate = PyThreadState_New(interp);
+            _PyThreadState_SetWhence(tstate, _PyThreadState_WHENCE_INTERP);
+            assert(tstate != NULL);
+            save_tstate = PyThreadState_Swap(tstate);
+        }
+    }
+    else {
+        assert(PyThreadState_GetInterpreter(tstate) == interp);
+        if (tstate != cur_tstate) {
+            assert(PyThreadState_GetInterpreter(cur_tstate) != interp);
+            save_tstate = PyThreadState_Swap(tstate);
+        }
+    }
+
+    Py_EndInterpreter(tstate);
+
+    if (p_save_tstate != NULL) {
+        save_tstate = *p_save_tstate;
+    }
+    PyThreadState_Swap(save_tstate);
+}