Always try reusing the main tstate (and drop interp.reuse_init_tstate).

python · ericsnowcurrently · Jun 20, 2024 · Jun 20, 2024 · Jun 20, 2024 · Jun 20, 2024
commit e6f31ff2b5e133a4d854e1bc264092fe89ce15b3
diff --git a/Include/internal/pycore_interp.h b/Include/internal/pycore_interp.h
@@ -130,10 +130,6 @@ struct _is {
         PyThreadState *head;
         /* The thread currently executing in the __main__ module, if any. */
         PyThreadState *main;
-#ifdef Py_DEBUG
-        /* For testing (_testembed, etc.). */
-        int reuse_init_tstate;
-#endif
         /* Used in Modules/_threadmodule.c. */
         Py_ssize_t count;
         /* Support for runtime thread stack size tuning.

diff --git a/Lib/test/test_embed.py b/Lib/test/test_embed.py
@@ -227,17 +227,17 @@ def test_repeated_init_and_inittab(self):
         self.assertEqual(out, lines)
 
     def test_replace_main_tstate(self):
-        with self.subTest(exec=False):
-            self.run_embedded_interpreter(
-                'test_replace_main_tstate',
-            )
-
-        with self.subTest(exec=True):
-            out, _ = self.run_embedded_interpreter(
-                'test_replace_main_tstate',
-                'print("spam!")'
-            )
-            self.assertEqual(out.strip(), 'spam!')
+        for reuse in (True, False):
+            with self.subTest(reuse=reuse, exec=False):
+                self.run_embedded_interpreter(
+                    'test_replace_main_tstate',
+                )
+            with self.subTest(reuse=reuse, exec=True):
+                out, _ = self.run_embedded_interpreter(
+                    'test_replace_main_tstate',
+                    'print("spam!")'
+                )
+                self.assertEqual(out.strip(), 'spam!')
 
     def test_fini_in_subthread(self):
         self.run_embedded_interpreter('test_fini_in_subthread')

diff --git a/Programs/_testembed.c b/Programs/_testembed.c
@@ -216,38 +216,58 @@ static int test_replace_main_tstate(void)
 {
     int err = 0;
 
+    int reuse = 0;
     int hascode = 0;
     for (int i = 2; i < main_argc; i++) {
-        if (strncmp(main_argv[i], "--", 2) == 0) {
+        if (strcmp(main_argv[i], "--reuse") == 0) {
+            reuse = 1;
+        }
+        else if (strncmp(main_argv[i], "--", 2) == 0) {
             fprintf(stderr, "ERROR: unsupported arg %s\n", main_argv[i]);
             err = 1;
         }
         else {
             hascode = 1;
         }
         if (err) {
-            fprintf(stderr, "usage: %s test_replace_main_tstate [CODE ...]\n",
+            fprintf(stderr,
+                    "usage: %s test_replace_main_tstate --reuse [CODE ...]\n",
                     PROGRAM);
             return 1;
         }
     }
 
-#ifdef Py_DEBUG
     _testembed_Py_InitializeFromConfig();
     PyThreadState *main_tstate = PyThreadState_Get();
     PyInterpreterState *main_interp = main_tstate->interp;
     assert(_Py_IsMainInterpreter(main_interp));
 
+    // If the initial tstate is reused then there are slightly
+    // different possible failure paths through the code than if we got
+    // a completely new one.
+
+    PyThreadState *tstate = NULL;
+    if (!reuse) {
+        // The initial thread state is still alive,
+        // so this will create a completely new one,
+        // with its own distinct pointer.
+        tstate = PyThreadState_New(main_interp);
+        assert(tstate != NULL);
+        assert(tstate != main_tstate);
+    }
     PyThreadState_Clear(main_tstate);
     PyThreadState_DeleteCurrent();
-    assert(PyInterpreterState_ThreadHead(main_interp) == NULL);
-
-    main_interp->threads.reuse_init_tstate = 0;
+    assert(reuse == (PyInterpreterState_ThreadHead(main_interp) == NULL));
+    if (reuse) {
+        // The initial thread state has already been "destroyed",
+        // so this will re-use the statically allocated tstate
+        // (along with reinitializing it).
+        tstate = PyThreadState_New(main_interp);
+        assert(tstate != NULL);
+        assert(tstate == main_tstate);
+    }
+    assert(_PyThreadState_GET() == NULL);
 
-    PyThreadState *tstate = PyThreadState_New(main_interp);
-    // The runtime automatically re-uses the initial main tstate.
-    // (It is statically allocated as part of the global runtime state.)
-    assert(_PyThreadState_GET() != main_tstate);
     (void)PyThreadState_Swap(tstate);
 
     if (hascode) {
@@ -260,15 +280,10 @@ static int test_replace_main_tstate(void)
         }
     }
 
-    assert(PyThreadState_Get() != main_tstate);
+    assert(PyThreadState_Get() == tstate);
     Py_Finalize();
 
     return err;
-#else
-    // We want to always force a new "main" tstate, which requires Py_DEBUG.
-    fprintf(stderr, "ERROR: Py_DEBUG is required\n");
-    return 1;
-#endif
 }
 
 

diff --git a/Python/pystate.c b/Python/pystate.c
@@ -630,10 +630,6 @@ init_interpreter(PyInterpreterState *interp,
     assert(next != NULL || (interp == runtime->interpreters.main));
     interp->next = next;
 
-#ifdef Py_DEBUG
-    interp->threads.reuse_init_tstate = 1;
-#endif
-
     PyStatus status = _PyObject_InitState(interp);
     if (_PyStatus_EXCEPTION(status)) {
         return status;
@@ -1559,12 +1555,7 @@ new_threadstate(PyInterpreterState *interp, int whence)
 
     // Allocate the thread state and add it to the interpreter.
     PyThreadState *old_head = interp->threads.head;
-    if (old_head == NULL
-#ifdef Py_DEBUG
-        && interp->threads.reuse_init_tstate
-#endif
-        )
-    {
+    if (old_head == NULL) {
         // It's the interpreter's initial thread state.
         used_newtstate = 0;
         tstate = &interp->_initial_thread;
@@ -1573,7 +1564,7 @@ new_threadstate(PyInterpreterState *interp, int whence)
     else {
         // Every valid interpreter must have at least one thread.
         assert(id > 1);
-        assert(old_head == NULL || old_head->prev == NULL);
+        assert(old_head->prev == NULL);
         used_newtstate = 1;
         tstate = new_tstate;
         // Set to _PyThreadState_INIT.