From 15d48aea02099ffc5bdc5511cc53ced460cb31b9 Mon Sep 17 00:00:00 2001
From: Eric Snow <ericsnowcurrently@gmail.com>
Date: Mon, 8 Jul 2024 15:10:00 -0600
Subject: [PATCH] gh-121110: Fix Extension Module Tests Under Py_TRACE_REFS
 Builds (gh-121503)

The change in gh-118157 (b2cd54a) should have also updated clear_singlephase_extension() but didn't.  We fix that here.  Note that clear_singlephase_extension() (AKA _PyImport_ClearExtension()) is only used in tests.
---
 Lib/test/test_import/__init__.py |  7 ----
 Python/import.c                  | 62 ++++++++++++++++++++++----------
 2 files changed, 43 insertions(+), 26 deletions(-)

diff --git a/Lib/test/test_import/__init__.py b/Lib/test/test_import/__init__.py
index c10f689c4ea..e29097baaf5 100644
--- a/Lib/test/test_import/__init__.py
+++ b/Lib/test/test_import/__init__.py
@@ -3034,13 +3034,6 @@ class SinglephaseInitTests(unittest.TestCase):
     def test_basic_multiple_interpreters_reset_each(self):
         # resetting between each interpreter
 
-        if Py_TRACE_REFS:
-            # It's a Py_TRACE_REFS build.
-            # This test breaks interpreter isolation a little,
-            # which causes problems on Py_TRACE_REF builds.
-            # See gh-121110.
-            raise unittest.SkipTest('crashes on Py_TRACE_REFS builds')
-
         # At this point:
         #  * alive in 0 interpreters
         #  * module def may or may not be loaded already
diff --git a/Python/import.c b/Python/import.c
index 20ad1002004..40b7feac001 100644
--- a/Python/import.c
+++ b/Python/import.c
@@ -1532,6 +1532,35 @@ switch_to_main_interpreter(PyThreadState *tstate)
     return main_tstate;
 }
 
+static void
+switch_back_from_main_interpreter(PyThreadState *tstate,
+                                  PyThreadState *main_tstate,
+                                  PyObject *tempobj)
+{
+    assert(main_tstate == PyThreadState_GET());
+    assert(_Py_IsMainInterpreter(main_tstate->interp));
+    assert(tstate->interp != main_tstate->interp);
+
+    /* Handle any exceptions, which we cannot propagate directly
+     * to the subinterpreter. */
+    if (PyErr_Occurred()) {
+        if (PyErr_ExceptionMatches(PyExc_MemoryError)) {
+            /* We trust it will be caught again soon. */
+            PyErr_Clear();
+        }
+        else {
+            /* Printing the exception should be sufficient. */
+            PyErr_PrintEx(0);
+        }
+    }
+
+    Py_XDECREF(tempobj);
+
+    PyThreadState_Clear(main_tstate);
+    (void)PyThreadState_Swap(tstate);
+    PyThreadState_Delete(main_tstate);
+}
+
 static PyObject *
 get_core_module_dict(PyInterpreterState *interp,
                      PyObject *name, PyObject *path)
@@ -2027,27 +2056,10 @@ main_finally:
     /* Switch back to the subinterpreter. */
     if (switched) {
         assert(main_tstate != tstate);
-
-        /* Handle any exceptions, which we cannot propagate directly
-         * to the subinterpreter. */
-        if (PyErr_Occurred()) {
-            if (PyErr_ExceptionMatches(PyExc_MemoryError)) {
-                /* We trust it will be caught again soon. */
-                PyErr_Clear();
-            }
-            else {
-                /* Printing the exception should be sufficient. */
-                PyErr_PrintEx(0);
-            }
-        }
-
+        switch_back_from_main_interpreter(tstate, main_tstate, mod);
         /* Any module we got from the init function will have to be
          * reloaded in the subinterpreter. */
-        Py_CLEAR(mod);
-
-        PyThreadState_Clear(main_tstate);
-        (void)PyThreadState_Swap(tstate);
-        PyThreadState_Delete(main_tstate);
+        mod = NULL;
     }
 
     /*****************************************************************/
@@ -2141,9 +2153,21 @@ clear_singlephase_extension(PyInterpreterState *interp,
         }
     }
 
+    /* We must use the main interpreter to clean up the cache.
+     * See the note in import_run_extension(). */
+    PyThreadState *tstate = PyThreadState_GET();
+    PyThreadState *main_tstate = switch_to_main_interpreter(tstate);
+    if (main_tstate == NULL) {
+        return -1;
+    }
+
     /* Clear the cached module def. */
     _extensions_cache_delete(path, name);
 
+    if (main_tstate != tstate) {
+        switch_back_from_main_interpreter(tstate, main_tstate, NULL);
+    }
+
     return 0;
 }