From 2cda9554b35c059f1766be897c2074d643839076 Mon Sep 17 00:00:00 2001
From: Eric Snow <ericsnowcurrently@gmail.com>
Date: Tue, 19 Mar 2024 16:31:43 -0600
Subject: [PATCH 1/2] Add _PyInterpreterState_ReinitRunningMain().

---
 Include/internal/pycore_pystate.h |  3 +++
 Modules/posixmodule.c             |  4 ++++
 Python/pylifecycle.c              |  4 ++++
 Python/pystate.c                  | 33 +++++++++++++++++++++++++++++++
 4 files changed, 44 insertions(+)

diff --git a/Include/internal/pycore_pystate.h b/Include/internal/pycore_pystate.h
index 6f9e6a332a7830..9aa439229cc8ea 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 7b2d3661ee5546..7cc7e5dcba158c 100644
--- a/Modules/posixmodule.c
+++ b/Modules/posixmodule.c
@@ -640,6 +640,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
@@ -649,6 +650,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 3a2c0a450ac9d9..fae074f5c8ec3d 100644
--- a/Python/pylifecycle.c
+++ b/Python/pylifecycle.c
@@ -3135,6 +3135,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 eedcb920cd1cf2..dba3d0db195fc6 100644
--- a/Python/pystate.c
+++ b/Python/pystate.c
@@ -1048,6 +1048,28 @@ _PyInterpreterState_IsRunningMain(PyInterpreterState *interp)
     return 0;
 }
 
+static int
+is_running_main(PyThreadState *tstate)
+{
+    if (tstate->interp->threads.main != NULL) {
+        return tstate == tstate->interp->threads.main;
+    }
+    return 0;
+}
+
+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)
 {
@@ -1059,6 +1081,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
@@ -1488,6 +1519,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
 
@@ -1586,6 +1618,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) {

From a7a1d1e1dcba5a878ef549675b57c06124ae8f82 Mon Sep 17 00:00:00 2001
From: Eric Snow <ericsnowcurrently@gmail.com>
Date: Wed, 20 Mar 2024 10:36:13 -0600
Subject: [PATCH 2/2] is_running_main() is only used in asserts.

---
 Python/pystate.c | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/Python/pystate.c b/Python/pystate.c
index dba3d0db195fc6..da17fcaa264e98 100644
--- a/Python/pystate.c
+++ b/Python/pystate.c
@@ -1048,6 +1048,7 @@ _PyInterpreterState_IsRunningMain(PyInterpreterState *interp)
     return 0;
 }
 
+#ifndef NDEBUG
 static int
 is_running_main(PyThreadState *tstate)
 {
@@ -1056,6 +1057,7 @@ is_running_main(PyThreadState *tstate)
     }
     return 0;
 }
+#endif
 
 int
 _PyThreadState_IsRunningMain(PyThreadState *tstate)