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.

1 files changed, 35 insertions, 0 deletions

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) {