gh-105716: Fix _PyInterpreterState_IsRunningMain() For Embedders (gh-117140)

ericsnowcurrently · web-flow · commit b3d25df8d38b · 2024-03-21T18:20:20.000-06:00
When I added _PyInterpreterState_IsRunningMain() and friends last year, I tried to accommodate applications that embed Python but don't call _PyInterpreterState_SetRunningMain() (not that they're expected to). That mostly worked fine until my recent changes in gh-117049, where the subtleties with the fallback code led to failures; the change ended up breaking test_tools.test_freeze, which exercises a basic embedding situation. The simplest fix is to drop the fallback code I originally added to _PyInterpreterState_IsRunningMain() (and later to _PyThreadState_IsRunningMain()). I've kept the fallback in the _xxsubinterpreters module though. I've also updated Py_FrozenMain() to call _PyInterpreterState_SetRunningMain().
diff --git a/Modules/_xxsubinterpretersmodule.c b/Modules/_xxsubinterpretersmodule.c
@@ -56,6 +56,24 @@ _get_current_module(void)
 }
 
 
+static int
+is_running_main(PyInterpreterState *interp)
+{
+    if (_PyInterpreterState_IsRunningMain(interp)) {
+        return 1;
+    }
+    // Unlike with the general C-API, we can be confident that someone
+    // using this module for the main interpreter is doing so through
+    // the main program.  Thus we can make this extra check.  This benefits
+    // applications that embed Python but haven't been updated yet
+    // to call_PyInterpreterState_SetRunningMain().
+    if (_Py_IsMainInterpreter(interp)) {
+        return 1;
+    }
+    return 0;
+}
+
+
 /* Cross-interpreter Buffer Views *******************************************/
 
 // XXX Release when the original interpreter is destroyed.
@@ -509,7 +527,7 @@ 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 (_PyInterpreterState_IsRunningMain(interp)) {
+    if (is_running_main(interp)) {
         PyErr_Format(PyExc_RuntimeError, "interpreter running");
         return NULL;
     }
@@ -977,7 +995,7 @@ interp_is_running(PyObject *self, PyObject *args, PyObject *kwds)
     if (interp == NULL) {
         return NULL;
     }
-    if (_PyInterpreterState_IsRunningMain(interp)) {
+    if (is_running_main(interp)) {
         Py_RETURN_TRUE;
     }
     Py_RETURN_FALSE;
diff --git a/Python/frozenmain.c b/Python/frozenmain.c
@@ -54,6 +54,12 @@ Py_FrozenMain(int argc, char **argv)
         Py_ExitStatusException(status);
     }
 
+    PyInterpreterState *interp = PyInterpreterState_Get();
+    if (_PyInterpreterState_SetRunningMain(interp) < 0) {
+        PyErr_Print();
+        exit(1);
+    }
+
 #ifdef MS_WINDOWS
     PyWinFreeze_ExeInit();
 #endif
@@ -83,6 +89,9 @@ Py_FrozenMain(int argc, char **argv)
 #ifdef MS_WINDOWS
     PyWinFreeze_ExeTerm();
 #endif
+
+    _PyInterpreterState_SetNotRunningMain(interp);
+
     if (Py_FinalizeEx() < 0) {
         sts = 120;
     }
diff --git a/Python/pystate.c b/Python/pystate.c
@@ -1042,34 +1042,24 @@ _PyInterpreterState_IsRunningMain(PyInterpreterState *interp)
     if (interp->threads.main != NULL) {
         return 1;
     }
-    // For now, we assume the main interpreter is always running.
-    if (_Py_IsMainInterpreter(interp)) {
-        return 1;
-    }
+    // Embedders might not know to call _PyInterpreterState_SetRunningMain(),
+    // so their main thread wouldn't show it is running the main interpreter's
+    // program.  (Py_Main() doesn't have this problem.)  For now this isn't
+    // critical.  If it were, we would need to infer "running main" from other
+    // information, like if it's the main interpreter.  We used to do that
+    // but the naive approach led to some inconsistencies that caused problems.
     return 0;
 }
 
-#ifndef NDEBUG
-static int
-is_running_main(PyThreadState *tstate)
-{
-    if (tstate->interp->threads.main != NULL) {
-        return tstate == tstate->interp->threads.main;
-    }
-    return 0;
-}
-#endif
-
 int
 _PyThreadState_IsRunningMain(PyThreadState *tstate)
 {
     PyInterpreterState *interp = tstate->interp;
     if (interp->threads.main != NULL) {
         return tstate == interp->threads.main;
     }
-    if (_Py_IsMainInterpreter(interp)) {
-        return tstate->thread_id == interp->runtime->main_thread;
-    }
+    // See the note in _PyInterpreterState_IsRunningMain() about
+    // possible false negatives here for embedders.
     return 0;
 }
 
@@ -1571,7 +1561,7 @@ PyThreadState_Clear(PyThreadState *tstate)
 {
     assert(tstate->_status.initialized && !tstate->_status.cleared);
     assert(current_fast_get()->interp == tstate->interp);
-    assert(!is_running_main(tstate));
+    assert(!_PyThreadState_IsRunningMain(tstate));
     // XXX assert(!tstate->_status.bound || tstate->_status.unbound);
     tstate->_status.finalizing = 1;  // just in case
 
@@ -1670,7 +1660,7 @@ tstate_delete_common(PyThreadState *tstate)
     assert(tstate->_status.cleared && !tstate->_status.finalized);
     assert(tstate->state != _Py_THREAD_ATTACHED);
     tstate_verify_not_active(tstate);
-    assert(!is_running_main(tstate));
+    assert(!_PyThreadState_IsRunningMain(tstate));
 
     PyInterpreterState *interp = tstate->interp;
     if (interp == NULL) {
