gh-105716: Update interp->threads.main After Fork (gh-117049)

I missed this in gh-109921.

We also update Py_Exit() to call _PyInterpreterState_SetNotRunningMain(), if necessary.

4 files changed, 46 insertions, 0 deletions

diff --git a/Include/internal/pycore_pystate.h b/Include/internal/pycore_pystate.h
index 6f9e6a3..9aa4392 100644
--- a/Include/internal/pycore_pystate.h
+++ b/Include/internal/pycore_pystate.h
@@ -83,6 +83,9 @@ PyAPI_FUNC(void) _PyInterpreterState_SetNotRunningMain(PyInterpreterState *);
 PyAPI_FUNC(int) _PyInterpreterState_IsRunningMain(PyInterpreterState *);
 PyAPI_FUNC(int) _PyInterpreterState_FailIfRunningMain(PyInterpreterState *);
 
+extern int _PyThreadState_IsRunningMain(PyThreadState *);
+extern void _PyInterpreterState_ReinitRunningMain(PyThreadState *);
+
 
 static inline const PyConfig *
 _Py_GetMainConfig(void)
diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c
index 644d4ba..8867916 100644
--- a/Modules/posixmodule.c
+++ b/Modules/posixmodule.c
@@ -646,6 +646,7 @@ PyOS_AfterFork_Child(void)
     PyThreadState *tstate = _PyThreadState_GET();
     _Py_EnsureTstateNotNULL(tstate);
 
+    assert(tstate->thread_id == PyThread_get_thread_ident());
 #ifdef PY_HAVE_THREAD_NATIVE_ID
     tstate->native_thread_id = PyThread_get_thread_native_id();
 #endif
@@ -655,6 +656,9 @@ PyOS_AfterFork_Child(void)
     _Py_qsbr_after_fork((_PyThreadStateImpl *)tstate);
 #endif
 
+    // Ideally we could guarantee tstate is running main.
+    _PyInterpreterState_ReinitRunningMain(tstate);
+
     status = _PyEval_ReInitThreads(tstate);
     if (_PyStatus_EXCEPTION(status)) {
         goto fatal_error;
diff --git a/Python/pylifecycle.c b/Python/pylifecycle.c
index bc76822..4bd01ad 100644
--- a/Python/pylifecycle.c
+++ b/Python/pylifecycle.c
@@ -3138,6 +3138,10 @@ call_ll_exitfuncs(_PyRuntimeState *runtime)
 void _Py_NO_RETURN
 Py_Exit(int sts)
 {
+    PyThreadState *tstate = _PyThreadState_GET();
+    if (tstate != NULL && _PyThreadState_IsRunningMain(tstate)) {
+        _PyInterpreterState_SetNotRunningMain(tstate->interp);
+    }
     if (Py_FinalizeEx() < 0) {
         sts = 120;
     }
diff --git a/Python/pystate.c b/Python/pystate.c
index 5332b8a..6d63eac 100644
--- a/Python/pystate.c
+++ b/Python/pystate.c
@@ -1050,6 +1050,30 @@ _PyInterpreterState_IsRunningMain(PyInterpreterState *interp)
     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;
+    }
+    return 0;
+}
+
 int
 _PyInterpreterState_FailIfRunningMain(PyInterpreterState *interp)
 {
@@ -1061,6 +1085,15 @@ _PyInterpreterState_FailIfRunningMain(PyInterpreterState *interp)
     return 0;
 }
 
+void
+_PyInterpreterState_ReinitRunningMain(PyThreadState *tstate)
+{
+    PyInterpreterState *interp = tstate->interp;
+    if (interp->threads.main != tstate) {
+        interp->threads.main = NULL;
+    }
+}
+
 
 //----------
 // accessors
@@ -1543,6 +1576,7 @@ PyThreadState_Clear(PyThreadState *tstate)
 {
     assert(tstate->_status.initialized && !tstate->_status.cleared);
     assert(current_fast_get()->interp == tstate->interp);
+    assert(!is_running_main(tstate));
     // XXX assert(!tstate->_status.bound || tstate->_status.unbound);
     tstate->_status.finalizing = 1;  // just in case
 
@@ -1641,6 +1675,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));
 
     PyInterpreterState *interp = tstate->interp;
     if (interp == NULL) {