Add PyInterpreterState_*RunningMain().

python · ericsnowcurrently · Oct 2, 2023 · Sep 25, 2023 · Sep 26, 2023 · Sep 26, 2023
commit d08ac68ded822c4043690523c72ea82afc6fd202
@@ -211,6 +211,12 @@ struct _ts {
  * if it is NULL. */
 PyAPI_FUNC(PyThreadState *) _PyThreadState_UncheckedGet(void);
 
+PyAPI_FUNC(int) PyInterpreterState_SetRunningMain(PyInterpreterState *);
+PyAPI_FUNC(void) PyInterpreterState_SetNotRunningMain(PyInterpreterState *);
+PyAPI_FUNC(int) PyInterpreterState_IsRunningMain(PyInterpreterState *);
+PyAPI_FUNC(int) PyThreadState_IsRunning(PyThreadState *);
+
+
 // Disable tracing and profiling.
 PyAPI_FUNC(void) PyThreadState_EnterTracing(PyThreadState *tstate);
 

@@ -73,6 +73,8 @@ struct _is {
         uint64_t next_unique_id;
         /* The linked list of threads, newest first. */
         PyThreadState *head;
+        /* The thread currently executing in the __main__ module, if any. */
+        PyThreadState *main;
         /* Used in Modules/_threadmodule.c. */
         long count;
         /* Support for runtime thread stack size tuning.

diff --git a/Modules/_xxsubinterpretersmodule.c b/Modules/_xxsubinterpretersmodule.c
@@ -312,7 +312,8 @@ _sharedexception_bind(PyObject *exc, _sharedexception *sharedexc)
 }
 
 static void
-_sharedexception_apply(_sharedexception *exc, PyObject *wrapperclass)
+_sharedexception_apply(_sharedexception *exc, PyObject *wrapperclass,
+                       PyInterpreterState *interp)
 {
     if (exc->name != NULL) {
         if (exc->msg != NULL) {
@@ -358,41 +359,14 @@ exceptions_init(PyObject *mod)
 }
 
 static int
-_is_running(PyInterpreterState *interp)
-{
-    PyThreadState *tstate = PyInterpreterState_ThreadHead(interp);
-    if (PyThreadState_Next(tstate) != NULL) {
-        PyErr_SetString(PyExc_RuntimeError,
-                        "interpreter has more than one thread");
-        return -1;
-    }
-
-    assert(!PyErr_Occurred());
-    struct _PyInterpreterFrame *frame = tstate->current_frame;
-    if (frame == NULL) {
-        return 0;
-    }
-    return 1;
-}
-
-static int
-_ensure_not_running(PyInterpreterState *interp)
+_run_script(PyInterpreterState *interp, const char *codestr,
+            _sharedns *shared, _sharedexception *sharedexc)
 {
-    int is_running = _is_running(interp);
-    if (is_running < 0) {
-        return -1;
-    }
-    if (is_running) {
-        PyErr_Format(PyExc_RuntimeError, "interpreter already running");
+    if (PyInterpreterState_SetRunningMain(interp) < 0) {
+        // We skip going through the shared exception.
         return -1;
     }
-    return 0;
-}
 
-static int
-_run_script(PyInterpreterState *interp, const char *codestr,
-            _sharedns *shared, _sharedexception *sharedexc)
-{
     PyObject *excval = NULL;
     PyObject *main_mod = PyUnstable_InterpreterState_GetMainModule(interp);
     if (main_mod == NULL) {
@@ -415,6 +389,7 @@ _run_script(PyInterpreterState *interp, const char *codestr,
 
     // Run the string (see PyRun_SimpleStringFlags).
     PyObject *result = PyRun_StringFlags(codestr, Py_file_input, ns, ns, NULL);
+    PyInterpreterState_SetNotRunningMain(interp);
     Py_DECREF(ns);
     if (result == NULL) {
         goto error;
@@ -437,16 +412,14 @@ _run_script(PyInterpreterState *interp, const char *codestr,
     }
     Py_XDECREF(excval);
     assert(!PyErr_Occurred());
+    PyInterpreterState_SetNotRunningMain(interp);
     return -1;
 }
 
 static int
 _run_script_in_interpreter(PyObject *mod, PyInterpreterState *interp,
                            const char *codestr, PyObject *shareables)
 {
-    if (_ensure_not_running(interp) < 0) {
-        return -1;
-    }
     module_state *state = get_module_state(mod);
 
     _sharedns *shared = _get_shared_ns(shareables);
@@ -456,9 +429,11 @@ _run_script_in_interpreter(PyObject *mod, PyInterpreterState *interp,
 
     // Switch to interpreter.
     PyThreadState *save_tstate = NULL;
-    if (interp != PyInterpreterState_Get()) {
+    PyThreadState *tcur = PyThreadState_Get();
+    PyThreadState *tstate = tcur;
+    if (interp != tstate->interp) {
         // XXX Using the "head" thread isn't strictly correct.
-        PyThreadState *tstate = PyInterpreterState_ThreadHead(interp);
+        tstate = PyInterpreterState_ThreadHead(interp);
         // XXX Possible GILState issues?
         save_tstate = PyThreadState_Swap(tstate);
     }
@@ -475,11 +450,13 @@ _run_script_in_interpreter(PyObject *mod, PyInterpreterState *interp,
     // Propagate any exception out to the caller.
     if (exc.name != NULL) {
         assert(state != NULL);
-        _sharedexception_apply(&exc, state->RunFailedError);
+        _sharedexception_apply(&exc, state->RunFailedError, interp);
     }
     else if (result != 0) {
-        // We were unable to allocate a shared exception.
-        PyErr_NoMemory();
+        if (!PyErr_Occurred()) {
+            // We were unable to allocate a shared exception.
+            PyErr_NoMemory();
+        }
     }
 
     if (shared != NULL) {
@@ -574,7 +551,8 @@ interp_destroy(PyObject *self, PyObject *args, PyObject *kwds)
     // Ensure the interpreter isn't running.
     /* XXX We *could* support destroying a running interpreter but
        aren't going to worry about it for now. */
-    if (_ensure_not_running(interp) < 0) {
+    if (PyInterpreterState_IsRunningMain(interp)) {
+        PyErr_Format(PyExc_RuntimeError, "interpreter running");
         return NULL;
     }
 
@@ -748,11 +726,7 @@ interp_is_running(PyObject *self, PyObject *args, PyObject *kwds)
     if (interp == NULL) {
         return NULL;
     }
-    int is_running = _is_running(interp);
-    if (is_running < 0) {
-        return NULL;
-    }
-    if (is_running) {
+    if (PyInterpreterState_IsRunningMain(interp)) {
         Py_RETURN_TRUE;
     }
     Py_RETURN_FALSE;
@@ -763,6 +737,7 @@ PyDoc_STRVAR(is_running_doc,
 \n\
 Return whether or not the identified interpreter is running.");
 
+
 static PyMethodDef module_functions[] = {
     {"create",                    _PyCFunction_CAST(interp_create),
      METH_VARARGS | METH_KEYWORDS, create_doc},

@@ -1095,6 +1095,50 @@ _PyInterpreterState_DeleteExceptMain(_PyRuntimeState *runtime)
 #endif
 
 
+int
+PyInterpreterState_SetRunningMain(PyInterpreterState *interp)
+{
+    if (interp->threads.main != NULL) {
+        PyErr_SetString(PyExc_RuntimeError,
+                        "interpreter already running");
+        return -1;
+    }
+    PyThreadState *tstate = current_fast_get(&_PyRuntime);
+    _Py_EnsureTstateNotNULL(tstate);
+    if (tstate->interp != interp) {
+        PyErr_SetString(PyExc_RuntimeError,
+                        "current tstate has wrong interpreter");
+        return -1;
+    }
+    interp->threads.main = tstate;
+    return 0;
+}
+
+void
+PyInterpreterState_SetNotRunningMain(PyInterpreterState *interp)
+{
+    assert(interp->threads.main == current_fast_get(&_PyRuntime));
+    interp->threads.main = NULL;
+}
+
+static int interpreter_exists(PyInterpreterState *);
+
+int
+PyInterpreterState_IsRunningMain(PyInterpreterState *interp)
+{
+    HEAD_LOCK(interp->runtime);
+    int exists = interpreter_exists(interp);
+    HEAD_UNLOCK(interp->runtime);
+    if (!exists) {
+        return 0;
+    }
+    if (interp->threads.main == NULL) {
+        return 0;
+    }
+    return PyThreadState_IsRunning(interp->threads.main);
+}
+
+
 //----------
 // accessors
 //----------
@@ -1204,6 +1248,22 @@ PyInterpreterState_GetDict(PyInterpreterState *interp)
 // look up an interpreter state
 //-----------------------------
 
+/* This must be called with the "head" lock held. */
+static int
+interpreter_exists(PyInterpreterState *interp)
+{
+    PyInterpreterState *actual = interp->runtime->interpreters.head;
+    while (actual != NULL) {
+        if (actual == interp) {
+            return 1;
+        }
+        actual = actual->next;
+    }
+    // It's a dangling pointer.
+    return 0;
+}
+
+
 /* Return the interpreter associated with the current OS thread.
 
    The GIL must be held.
@@ -1683,6 +1743,13 @@ _PyThreadState_DeleteExcept(PyThreadState *tstate)
 }
 
 
+int
+PyThreadState_IsRunning(PyThreadState *tstate)
+{
+    return tstate->current_frame != NULL;
+}
+
+
 //----------
 // accessors
 //----------
@@ -2805,6 +2872,10 @@ _register_builtins_for_crossinterpreter_data(struct _xidregistry *xidregistry)
 }
 
 
+/*************/
+/* Other API */
+/*************/
+
 _PyFrameEvalFunction
 _PyInterpreterState_GetEvalFrameFunc(PyInterpreterState *interp)
 {