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..da17fcaa264e98 100644 --- a/Python/pystate.c +++ b/Python/pystate.c @@ -1048,6 +1048,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) { @@ -1059,6 +1083,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 +1521,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 +1620,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) {