gh-109860: Use a New Thread State When Switching Interpreters, When Necessary (gh-110245)

ericsnowcurrently · web-flow · commit f5198b09e16b · 2023-10-03T09:20:48.000-06:00
In a few places we switch to another interpreter without knowing if it has a thread state associated with the current thread.  For the main interpreter there wasn't much of a problem, but for subinterpreters we were *mostly* okay re-using the tstate created with the interpreter (located via PyInterpreterState_ThreadHead()).  There was a good chance that tstate wasn't actually in use by another thread.

However, there are no guarantees of that.  Furthermore, re-using an already used tstate is currently fragile.  To address this, now we create a new thread state in each of those places and use it.

One consequence of this change is that PyInterpreterState_ThreadHead() may not return NULL (though that won't happen for the main interpreter).
diff --git a/Include/cpython/pystate.h b/Include/cpython/pystate.h
@@ -92,6 +92,15 @@ struct _ts {
         /* padding to align to 4 bytes */
         unsigned int :24;
     } _status;
+#ifdef Py_BUILD_CORE
+#  define _PyThreadState_WHENCE_NOTSET -1
+#  define _PyThreadState_WHENCE_UNKNOWN 0
+#  define _PyThreadState_WHENCE_INTERP 1
+#  define _PyThreadState_WHENCE_THREADING 2
+#  define _PyThreadState_WHENCE_GILSTATE 3
+#  define _PyThreadState_WHENCE_EXEC 4
+#endif
+    int _whence;
 
     int py_recursion_remaining;
     int py_recursion_limit;
diff --git a/Include/internal/pycore_pystate.h b/Include/internal/pycore_pystate.h
@@ -48,6 +48,7 @@ _Py_IsMainInterpreterFinalizing(PyInterpreterState *interp)
 PyAPI_FUNC(int) _PyInterpreterState_SetRunningMain(PyInterpreterState *);
 PyAPI_FUNC(void) _PyInterpreterState_SetNotRunningMain(PyInterpreterState *);
 PyAPI_FUNC(int) _PyInterpreterState_IsRunningMain(PyInterpreterState *);
+PyAPI_FUNC(int) _PyInterpreterState_FailIfRunningMain(PyInterpreterState *);
 
 
 static inline const PyConfig *
@@ -139,7 +140,9 @@ static inline PyInterpreterState* _PyInterpreterState_GET(void) {
 
 // PyThreadState functions
 
-extern PyThreadState * _PyThreadState_New(PyInterpreterState *interp);
+extern PyThreadState * _PyThreadState_New(
+    PyInterpreterState *interp,
+    int whence);
 extern void _PyThreadState_Bind(PyThreadState *tstate);
 extern void _PyThreadState_DeleteExcept(PyThreadState *tstate);
 
diff --git a/Include/internal/pycore_runtime_init.h b/Include/internal/pycore_runtime_init.h
@@ -185,6 +185,7 @@ extern PyTypeObject _PyExc_MemoryError;
 
 #define _PyThreadState_INIT \
     { \
+        ._whence = _PyThreadState_WHENCE_NOTSET, \
         .py_recursion_limit = Py_DEFAULT_RECURSION_LIMIT, \
         .context_ver = 1, \
     }
diff --git a/Lib/threading.py b/Lib/threading.py
@@ -1593,8 +1593,11 @@ def _shutdown():
         # The main thread isn't finished yet, so its thread state lock can't
         # have been released.
         assert tlock is not None
-        assert tlock.locked()
-        tlock.release()
+        if tlock.locked():
+            # It should have been released already by
+            # _PyInterpreterState_SetNotRunningMain(), but there may be
+            # embedders that aren't calling that yet.
+            tlock.release()
         _main_thread._stop()
     else:
         # bpo-1596321: _shutdown() must be called in the main thread.
diff --git a/Modules/_threadmodule.c b/Modules/_threadmodule.c
@@ -1205,7 +1205,7 @@ thread_PyThread_start_new_thread(PyObject *self, PyObject *fargs)
     if (boot == NULL) {
         return PyErr_NoMemory();
     }
-    boot->tstate = _PyThreadState_New(interp);
+    boot->tstate = _PyThreadState_New(interp, _PyThreadState_WHENCE_THREADING);
     if (boot->tstate == NULL) {
         PyMem_RawFree(boot);
         if (!PyErr_Occurred()) {
diff --git a/Modules/_xxsubinterpretersmodule.c b/Modules/_xxsubinterpretersmodule.c
@@ -242,6 +242,11 @@ _sharedns_apply(_sharedns *shared, PyObject *ns)
 // of the exception in the calling interpreter.
 
 typedef struct _sharedexception {
+    PyInterpreterState *interp;
+#define ERR_NOT_SET 0
+#define ERR_NO_MEMORY 1
+#define ERR_ALREADY_RUNNING 2
+    int code;
     const char *name;
     const char *msg;
 } _sharedexception;
@@ -263,14 +268,26 @@ _sharedexception_clear(_sharedexception *exc)
 }
 
 static const char *
-_sharedexception_bind(PyObject *exc, _sharedexception *sharedexc)
+_sharedexception_bind(PyObject *exc, int code, _sharedexception *sharedexc)
 {
+    if (sharedexc->interp == NULL) {
+        sharedexc->interp = PyInterpreterState_Get();
+    }
+
+    if (code != ERR_NOT_SET) {
+        assert(exc == NULL);
+        assert(code > 0);
+        sharedexc->code = code;
+        return NULL;
+    }
+
     assert(exc != NULL);
     const char *failure = NULL;
 
     PyObject *nameobj = PyUnicode_FromString(Py_TYPE(exc)->tp_name);
     if (nameobj == NULL) {
         failure = "unable to format exception type name";
+        code = ERR_NO_MEMORY;
         goto error;
     }
     sharedexc->name = _copy_raw_string(nameobj);
@@ -281,13 +298,15 @@ _sharedexception_bind(PyObject *exc, _sharedexception *sharedexc)
         } else {
             failure = "unable to encode and copy exception type name";
         }
+        code = ERR_NO_MEMORY;
         goto error;
     }
 
     if (exc != NULL) {
         PyObject *msgobj = PyUnicode_FromFormat("%S", exc);
         if (msgobj == NULL) {
             failure = "unable to format exception message";
+            code = ERR_NO_MEMORY;
             goto error;
         }
         sharedexc->msg = _copy_raw_string(msgobj);
@@ -298,6 +317,7 @@ _sharedexception_bind(PyObject *exc, _sharedexception *sharedexc)
             } else {
                 failure = "unable to encode and copy exception message";
             }
+            code = ERR_NO_MEMORY;
             goto error;
         }
     }
@@ -308,14 +328,18 @@ _sharedexception_bind(PyObject *exc, _sharedexception *sharedexc)
     assert(failure != NULL);
     PyErr_Clear();
     _sharedexception_clear(sharedexc);
-    *sharedexc = no_exception;
+    *sharedexc = (_sharedexception){
+        .interp = sharedexc->interp,
+        .code = code,
+    };
     return failure;
 }
 
 static void
 _sharedexception_apply(_sharedexception *exc, PyObject *wrapperclass)
 {
     if (exc->name != NULL) {
+        assert(exc->code == ERR_NOT_SET);
         if (exc->msg != NULL) {
             PyErr_Format(wrapperclass, "%s: %s",  exc->name, exc->msg);
         }
@@ -324,9 +348,19 @@ _sharedexception_apply(_sharedexception *exc, PyObject *wrapperclass)
         }
     }
     else if (exc->msg != NULL) {
+        assert(exc->code == ERR_NOT_SET);
         PyErr_SetString(wrapperclass, exc->msg);
     }
+    else if (exc->code == ERR_NO_MEMORY) {
+        PyErr_NoMemory();
+    }
+    else if (exc->code == ERR_ALREADY_RUNNING) {
+        assert(exc->interp != NULL);
+        assert(_PyInterpreterState_IsRunningMain(exc->interp));
+        _PyInterpreterState_FailIfRunningMain(exc->interp);
+    }
     else {
+        assert(exc->code == ERR_NOT_SET);
         PyErr_SetNone(wrapperclass);
     }
 }
@@ -362,9 +396,16 @@ static int
 _run_script(PyInterpreterState *interp, const char *codestr,
             _sharedns *shared, _sharedexception *sharedexc)
 {
+    int errcode = ERR_NOT_SET;
+
     if (_PyInterpreterState_SetRunningMain(interp) < 0) {
-        // We skip going through the shared exception.
-        return -1;
+        assert(PyErr_Occurred());
+        // In the case where we didn't switch interpreters, it would
+        // be more efficient to leave the exception in place and return
+        // immediately.  However, life is simpler if we don't.
+        PyErr_Clear();
+        errcode = ERR_ALREADY_RUNNING;
+        goto error;
     }
 
     PyObject *excval = NULL;
@@ -403,16 +444,17 @@ _run_script(PyInterpreterState *interp, const char *codestr,
 
 error:
     excval = PyErr_GetRaisedException();
-    const char *failure = _sharedexception_bind(excval, sharedexc);
+    const char *failure = _sharedexception_bind(excval, errcode, sharedexc);
     if (failure != NULL) {
         fprintf(stderr,
                 "RunFailedError: script raised an uncaught exception (%s)",
                 failure);
-        PyErr_Clear();
     }
     Py_XDECREF(excval);
+    if (errcode != ERR_ALREADY_RUNNING) {
+        _PyInterpreterState_SetNotRunningMain(interp);
+    }
     assert(!PyErr_Occurred());
-    _PyInterpreterState_SetNotRunningMain(interp);
     return -1;
 }
 
@@ -421,6 +463,7 @@ _run_script_in_interpreter(PyObject *mod, PyInterpreterState *interp,
                            const char *codestr, PyObject *shareables)
 {
     module_state *state = get_module_state(mod);
+    assert(state != NULL);
 
     _sharedns *shared = _get_shared_ns(shareables);
     if (shared == NULL && PyErr_Occurred()) {
@@ -429,50 +472,30 @@ _run_script_in_interpreter(PyObject *mod, PyInterpreterState *interp,
 
     // Switch to interpreter.
     PyThreadState *save_tstate = NULL;
+    PyThreadState *tstate = NULL;
     if (interp != PyInterpreterState_Get()) {
-        // XXX gh-109860: Using the "head" thread isn't strictly correct.
-        PyThreadState *tstate = PyInterpreterState_ThreadHead(interp);
-        assert(tstate != NULL);
-        // Hack (until gh-109860): The interpreter's initial thread state
-        // is least likely to break.
-        while(tstate->next != NULL) {
-            tstate = tstate->next;
-        }
-        // We must do this check before switching interpreters, so any
-        // exception gets raised in the right one.
-        // XXX gh-109860: Drop this redundant check once we stop
-        // re-using tstates that might already be in use.
-        if (_PyInterpreterState_IsRunningMain(interp)) {
-            PyErr_SetString(PyExc_RuntimeError,
-                            "interpreter already running");
-            if (shared != NULL) {
-                _sharedns_free(shared);
-            }
-            return -1;
-        }
+        tstate = PyThreadState_New(interp);
+        tstate->_whence = _PyThreadState_WHENCE_EXEC;
         // XXX Possible GILState issues?
         save_tstate = PyThreadState_Swap(tstate);
     }
 
     // Run the script.
-    _sharedexception exc = {NULL, NULL};
+    _sharedexception exc = (_sharedexception){ .interp = interp };
     int result = _run_script(interp, codestr, shared, &exc);
 
     // Switch back.
     if (save_tstate != NULL) {
+        PyThreadState_Clear(tstate);
         PyThreadState_Swap(save_tstate);
+        PyThreadState_Delete(tstate);
     }
 
     // Propagate any exception out to the caller.
-    if (exc.name != NULL) {
-        assert(state != NULL);
+    if (result < 0) {
+        assert(!PyErr_Occurred());
         _sharedexception_apply(&exc, state->RunFailedError);
-    }
-    else if (result != 0) {
-        if (!PyErr_Occurred()) {
-            // We were unable to allocate a shared exception.
-            PyErr_NoMemory();
-        }
+        assert(PyErr_Occurred());
     }
 
     if (shared != NULL) {
@@ -502,6 +525,7 @@ interp_create(PyObject *self, PyObject *args, PyObject *kwds)
     const PyInterpreterConfig config = isolated
         ? (PyInterpreterConfig)_PyInterpreterConfig_INIT
         : (PyInterpreterConfig)_PyInterpreterConfig_LEGACY_INIT;
+
     // XXX Possible GILState issues?
     PyThreadState *tstate = NULL;
     PyStatus status = Py_NewInterpreterFromConfig(&tstate, &config);
@@ -517,6 +541,7 @@ interp_create(PyObject *self, PyObject *args, PyObject *kwds)
         return NULL;
     }
     assert(tstate != NULL);
+
     PyInterpreterState *interp = PyThreadState_GetInterpreter(tstate);
     PyObject *idobj = PyInterpreterState_GetIDObject(interp);
     if (idobj == NULL) {
@@ -526,6 +551,10 @@ interp_create(PyObject *self, PyObject *args, PyObject *kwds)
         PyThreadState_Swap(save_tstate);
         return NULL;
     }
+
+    PyThreadState_Clear(tstate);
+    PyThreadState_Delete(tstate);
+
     _PyInterpreterState_RequireIDRef(interp, 1);
     return idobj;
 }
@@ -573,14 +602,8 @@ interp_destroy(PyObject *self, PyObject *args, PyObject *kwds)
     }
 
     // Destroy the interpreter.
-    // XXX gh-109860: Using the "head" thread isn't strictly correct.
-    PyThreadState *tstate = PyInterpreterState_ThreadHead(interp);
-    assert(tstate != NULL);
-    // Hack (until gh-109860): The interpreter's initial thread state
-    // is least likely to break.
-    while(tstate->next != NULL) {
-        tstate = tstate->next;
-    }
+    PyThreadState *tstate = PyThreadState_New(interp);
+    tstate->_whence = _PyThreadState_WHENCE_INTERP;
     // XXX Possible GILState issues?
     PyThreadState *save_tstate = PyThreadState_Swap(tstate);
     Py_EndInterpreter(tstate);
diff --git a/Python/pylifecycle.c b/Python/pylifecycle.c
@@ -655,7 +655,8 @@ pycore_create_interpreter(_PyRuntimeState *runtime,
         return status;
     }
 
-    PyThreadState *tstate = _PyThreadState_New(interp);
+    PyThreadState *tstate = _PyThreadState_New(interp,
+                                               _PyThreadState_WHENCE_INTERP);
     if (tstate == NULL) {
         return _PyStatus_ERR("can't make first thread");
     }
@@ -2050,7 +2051,8 @@ new_interpreter(PyThreadState **tstate_p, const PyInterpreterConfig *config)
         return _PyStatus_OK();
     }
 
-    PyThreadState *tstate = _PyThreadState_New(interp);
+    PyThreadState *tstate = _PyThreadState_New(interp,
+                                               _PyThreadState_WHENCE_INTERP);
     if (tstate == NULL) {
         PyInterpreterState_Delete(interp);
         *tstate_p = NULL;
diff --git a/Python/pystate.c b/Python/pystate.c

Original file line number	Diff line number	Diff line change
`@@ -185,6 +185,7 @@ extern PyTypeObject _PyExc_MemoryError;`
`185`	`185`
`186`	`186`	`#define _PyThreadState_INIT \`
`187`	`187`	`{ \`
	`188`	`+ ._whence = _PyThreadState_WHENCE_NOTSET, \`
`188`	`189`	`.py_recursion_limit = Py_DEFAULT_RECURSION_LIMIT, \`
`189`	`190`	`.context_ver = 1, \`
`190`	`191`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1205,7 +1205,7 @@ thread_PyThread_start_new_thread(PyObject self, PyObject fargs)`
`1205`	`1205`	`if (boot == NULL) {`
`1206`	`1206`	`return PyErr_NoMemory();`
`1207`	`1207`	`}`
`1208`		`- boot->tstate = _PyThreadState_New(interp);`
	`1208`	`+ boot->tstate = _PyThreadState_New(interp, _PyThreadState_WHENCE_THREADING);`
`1209`	`1209`	`if (boot->tstate == NULL) {`
`1210`	`1210`	`PyMem_RawFree(boot);`
`1211`	`1211`	`if (!PyErr_Occurred()) {`
Original file line number	Diff line number	Diff line change
`@@ -655,7 +655,8 @@ pycore_create_interpreter(_PyRuntimeState *runtime,`
`655`	`655`	`return status;`
`656`	`656`	`}`
`657`	`657`
`658`		`- PyThreadState *tstate = _PyThreadState_New(interp);`
	`658`	`+ PyThreadState *tstate = _PyThreadState_New(interp,`
	`659`	`+ _PyThreadState_WHENCE_INTERP);`
`659`	`660`	`if (tstate == NULL) {`
`660`	`661`	`return _PyStatus_ERR("can't make first thread");`
`661`	`662`	`}`
`@@ -2050,7 +2051,8 @@ new_interpreter(PyThreadState *tstate_p, const PyInterpreterConfig config)`
`2050`	`2051`	`return _PyStatus_OK();`
`2051`	`2052`	`}`
`2052`	`2053`
`2053`		`- PyThreadState *tstate = _PyThreadState_New(interp);`
	`2054`	`+ PyThreadState *tstate = _PyThreadState_New(interp,`
	`2055`	`+ _PyThreadState_WHENCE_INTERP);`
`2054`	`2056`	`if (tstate == NULL) {`
`2055`	`2057`	`PyInterpreterState_Delete(interp);`
`2056`	`2058`	`*tstate_p = NULL;`