gh-120838: Disallow Py_Finalize() if Py_RunMain() is Running

Include/internal/pycore_pylifecycle.h

@@ -114,6 +114,14 @@ PyAPI_FUNC(char*) _Py_SetLocaleFromEnv(int category);
 // Export for special main.c string compiling with source tracebacks
 int _PyRun_SimpleStringFlagsWithName(const char *command, const char* name, PyCompilerFlags *flags);
 
+struct pyfinalize_args {
+    const char *caller;
+    int check_pymain;
+};
+
+// Export for _testembed
+extern int _Py_Finalize(_PyRuntimeState *, struct pyfinalize_args *);
+
 
 /* interpreter config */
 

Include/internal/pycore_runtime.h

@@ -209,6 +209,8 @@ typedef struct pyruntimestate {
     unsigned long main_thread;
     PyThreadState *main_tstate;
 
+    int is_pymain;
+
     /* ---------- IMPORTANT ---------------------------
      The fields above this line are declared as early as
      possible to facilitate out-of-process observability

Misc/NEWS.d/next/C API/2024-06-21-09-55-56.gh-issue-120838.-xuGRW.rst

@@ -0,0 +1,3 @@
+Calling :c:func:`Py_Finalize` while :c:func:`Py_RunMain` is running is now a
+failure.  When used, ``Py_RunMain()`` is solely responsible for finalizing
+the runtime.

Modules/main.c

@@ -711,19 +711,25 @@ pymain_exit_error(PyStatus status)
 }
 
 
+extern int _Py_FinalizeMain(void);
+
 int
 Py_RunMain(void)
 {
     int exitcode = 0;
 
+    _PyRuntime.is_pymain = 1;
+
     pymain_run_python(&exitcode);
 
-    if (Py_FinalizeEx() < 0) {
+    if (_Py_FinalizeMain() < 0) {
         /* Value unlikely to be confused with a non-error exit status or
            other special meaning */
         exitcode = 120;
     }
 
+    _PyRuntime.is_pymain = 0;
+
     pymain_free();
 
     if (_PyRuntime.signals.unhandled_keyboard_interrupt) {

Python/pylifecycle.c

@@ -1910,17 +1910,27 @@ finalize_interp_delete(PyInterpreterState *interp)
 
 
 int
-Py_FinalizeEx(void)
+_Py_Finalize(_PyRuntimeState *runtime, struct pyfinalize_args *args)
 {
     int status = 0;
 
-    _PyRuntimeState *runtime = &_PyRuntime;
+    /* Bail out early if already finalized. */
     if (!runtime->initialized) {
+        assert(!runtime->is_pymain);
         return status;
     }
 
-    /* Get current thread state and interpreter pointer */
+    /* Make sure Py_RunMain() users aren't calling Py_Finalize(). */
+    if (args->check_pymain && runtime->is_pymain) {
+        fprintf(stderr,
+                "%s() should not be called while Py_RunMain() is running\n",
+                args->caller);
+        return -1;
+    }
+
+    /* Get final thread state pointer. */
     PyThreadState *tstate = _PyThreadState_GET();
+    assert(tstate->interp->runtime == runtime);
     // XXX assert(_Py_IsMainInterpreter(tstate->interp));
     // XXX assert(_Py_IsMainThread());
 
@@ -2142,10 +2152,39 @@ Py_FinalizeEx(void)
     return status;
 }
 
+/* _Py_FinalizeMain() is used exclusively by Py_RunMain(). */
+int
+_Py_FinalizeMain(void)
+{
+    _PyRuntimeState *runtime = &_PyRuntime;
+    assert(runtime->is_pymain);
+    assert(_Py_IsMainInterpreter(_PyInterpreterState_GET()));
+    assert(_Py_IsMainThread());
+    assert(_PyThreadState_GET() == runtime->main_tstate);
+    struct pyfinalize_args args = {
+        .caller = "Py_RunMain",
+    };
+    return _Py_Finalize(runtime, &args);
+}
+
+int
+Py_FinalizeEx(void)
+{
+    struct pyfinalize_args args = {
+        .caller = "Py_FinalizeEx",
+        .check_pymain = 1,
+    };
+    return _Py_Finalize(&_PyRuntime, &args);
+}
+
 void
 Py_Finalize(void)
 {
-    Py_FinalizeEx();
+    struct pyfinalize_args args = {
+        .caller = "Py_Finalize",
+        .check_pymain = 1,
+    };
+    (void)_Py_Finalize(&_PyRuntime, &args);
 }
 
 
@@ -3218,7 +3257,11 @@ Py_Exit(int sts)
     if (tstate != NULL && _PyThreadState_IsRunningMain(tstate)) {
         _PyInterpreterState_SetNotRunningMain(tstate->interp);
     }
-    if (Py_FinalizeEx() < 0) {
+    struct pyfinalize_args args = {
+        .caller = "Py_Exit",
+        /* We don't worry about checking if Py_RunMain() is running. */
+    };
+    if (_Py_Finalize(&_PyRuntime, &args) < 0) {
         sts = 120;
     }