python · ZeroIntensity · Apr 24, 2025 · Apr 24, 2025 · Apr 25, 2025 · Apr 25, 2025
@@ -208,6 +208,20 @@ struct _ts {
     */
     PyObject *threading_local_sentinel;
     _PyRemoteDebuggerSupport remote_debugger_support;
+
+    struct {
+        /* Number of nested PyThreadState_Ensure() calls on this thread state */
+        Py_ssize_t counter;
+
+        /* Thread state that was active before PyThreadState_Ensure() was called. */
+        PyThreadState *prior_tstate;
+
+        /* Should this thread state be deleted upon calling
+           PyThreadState_Release() (with the counter at 1)?
+
+           This is only true for thread states created by PyThreadState_Ensure() */
+        int delete_on_release;
+    } ensure;
 };
 
 /* other API */
@@ -261,3 +275,42 @@ PyAPI_FUNC(_PyFrameEvalFunction) _PyInterpreterState_GetEvalFrameFunc(
 PyAPI_FUNC(void) _PyInterpreterState_SetEvalFrameFunc(
     PyInterpreterState *interp,
     _PyFrameEvalFunction eval_frame);
+
+/* Strong interpreter references */
+
+typedef uintptr_t PyInterpreterRef;
+
+PyAPI_FUNC(int) PyInterpreterRef_Get(PyInterpreterRef *ref);
+PyAPI_FUNC(PyInterpreterRef) PyInterpreterRef_Dup(PyInterpreterRef ref);
+PyAPI_FUNC(int) PyInterpreterRef_Main(PyInterpreterRef *ref);
+PyAPI_FUNC(void) PyInterpreterRef_Close(PyInterpreterRef ref);
+PyAPI_FUNC(PyInterpreterState *) PyInterpreterRef_AsInterpreter(PyInterpreterRef ref);
+
+#define PyInterpreterRef_Close(ref) do {    \
+    PyInterpreterRef_Close(ref);            \
+    ref = 0;                                \
+} while (0)
+
+/* Weak interpreter references */
+
+typedef struct _PyInterpreterWeakRef {
+    int64_t id;
+    Py_ssize_t refcount;
+} _PyInterpreterWeakRef;
+
+typedef _PyInterpreterWeakRef *PyInterpreterWeakRef;
+
+PyAPI_FUNC(int) PyInterpreterWeakRef_Get(PyInterpreterWeakRef *ptr);
+PyAPI_FUNC(PyInterpreterWeakRef) PyInterpreterWeakRef_Dup(PyInterpreterWeakRef wref);
+PyAPI_FUNC(int) PyInterpreterWeakRef_AsStrong(PyInterpreterWeakRef wref, PyInterpreterRef *strong_ptr);
+PyAPI_FUNC(void) PyInterpreterWeakRef_Close(PyInterpreterWeakRef wref);
+
+#define PyInterpreterWeakRef_Close(ref) do {    \
+    PyInterpreterWeakRef_Close(ref);            \
+    ref = 0;                                    \
+} while (0)
+
+
+PyAPI_FUNC(int) PyThreadState_Ensure(PyInterpreterRef interp_ref);
+
+PyAPI_FUNC(void) PyThreadState_Release(void);
diff --git a/Include/internal/pycore_interp_structs.h b/Include/internal/pycore_interp_structs.h
@@ -810,6 +810,13 @@ struct _is {
            or the size specified by the THREAD_STACK_SIZE macro. */
         /* Used in Python/thread.c. */
         size_t stacksize;
+
+        struct _Py_finalizing_threads {
+            Py_ssize_t countdown;
+            PyEvent finished;
+            PyMutex mutex;
+            int shutting_down;
+        } finalizing;
     } threads;
 
     /* Reference to the _PyRuntime global variable. This field exists

@@ -328,6 +328,10 @@ _Py_RecursionLimit_GetMargin(PyThreadState *tstate)
     return Py_ARITHMETIC_RIGHT_SHIFT(intptr_t, here_addr - (intptr_t)_tstate->c_stack_soft_limit, PYOS_STACK_MARGIN_SHIFT);
 }
 
+// Exports for '_testinternalcapi' shared extension
+PyAPI_FUNC(Py_ssize_t) _PyInterpreterState_Refcount(PyInterpreterState *interp);
+PyAPI_FUNC(int) _PyInterpreterState_Incref(PyInterpreterState *interp);
+
 #ifdef __cplusplus
 }
 #endif

diff --git a/Lib/test/test_embed.py b/Lib/test/test_embed.py
@@ -1914,10 +1914,15 @@ def test_audit_run_stdin(self):
     def test_get_incomplete_frame(self):
         self.run_embedded_interpreter("test_get_incomplete_frame")
 
-
     def test_gilstate_after_finalization(self):
         self.run_embedded_interpreter("test_gilstate_after_finalization")
 
+    def test_thread_state_ensure(self):
+        self.run_embedded_interpreter("test_thread_state_ensure")
+
+    def test_main_interpreter_ref(self):
+        self.run_embedded_interpreter("test_main_interpreter_ref")
+
 
 class MiscTests(EmbeddingTestsMixin, unittest.TestCase):
     def test_unicode_id_init(self):

diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
@@ -2546,6 +2546,194 @@ toggle_reftrace_printer(PyObject *ob, PyObject *arg)
     Py_RETURN_NONE;
 }
 
+static PyInterpreterRef
+get_strong_ref(void)
+{
+    PyInterpreterRef ref;
+    if (PyInterpreterRef_Get(&ref) < 0) {
+        Py_FatalError("strong reference should not have failed");
+    }
+    return ref;
+}
+
+static void
+test_interp_ref_common(void)
+{
+    PyInterpreterState *interp = PyInterpreterState_Get();
+    PyInterpreterRef ref = get_strong_ref();
+    assert(PyInterpreterRef_AsInterpreter(ref) == interp);
+
+    PyInterpreterRef ref_2 = PyInterpreterRef_Dup(ref);
+    assert(PyInterpreterRef_AsInterpreter(ref_2) == interp);
+
+    // We can close the references in any order
+    PyInterpreterRef_Close(ref);
+    PyInterpreterRef_Close(ref_2);
+}
+
+static PyObject *
+test_interpreter_refs(PyObject *self, PyObject *unused)
+{
+    // Test the main interpreter
+    test_interp_ref_common();
+
+    // Test a (legacy) subinterpreter
+    PyThreadState *save_tstate = PyThreadState_Swap(NULL);
+    PyThreadState *interp_tstate = Py_NewInterpreter();
+    test_interp_ref_common();
+    Py_EndInterpreter(interp_tstate);
+
+    // Test an isolated subinterpreter
+    PyInterpreterConfig config = {
+        .gil = PyInterpreterConfig_OWN_GIL,
+        .check_multi_interp_extensions = 1
+    };
+
+    PyThreadState *isolated_interp_tstate;
+    PyStatus status = Py_NewInterpreterFromConfig(&isolated_interp_tstate, &config);
+    if (PyStatus_Exception(status)) {
+        PyErr_SetString(PyExc_RuntimeError, "interpreter creation failed");
+        return NULL;
+    }
+
+    test_interp_ref_common();
+    Py_EndInterpreter(isolated_interp_tstate);
+    PyThreadState_Swap(save_tstate);
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+test_thread_state_ensure_nested(PyObject *self, PyObject *unused)
+{
+    PyInterpreterRef ref = get_strong_ref();
+    PyThreadState *save_tstate = PyThreadState_Swap(NULL);
+    assert(PyGILState_GetThisThreadState() == save_tstate);
+
+    for (int i = 0; i < 10; ++i) {
+        // Test reactivation of the detached tstate.
+        if (PyThreadState_Ensure(ref) < 0) {
+            PyInterpreterRef_Close(ref);
+            return PyErr_NoMemory();
+        }
+
+        // No new thread state should've been created.
+        assert(PyThreadState_Get() == save_tstate);
+        PyThreadState_Release();
+    }
+
+    assert(PyThreadState_GetUnchecked() == NULL);
+
+    // Similarly, test ensuring with deep nesting and *then* releasing.
+    // If the (detached) gilstate matches the interpreter, then it shouldn't
+    // create a new thread state.
+    for (int i = 0; i < 10; ++i) {
+        if (PyThreadState_Ensure(ref) < 0) {
+            // This will technically leak other thread states, but it doesn't
+            // matter because this is a test.
+            PyInterpreterRef_Close(ref);
+            return PyErr_NoMemory();
+        }
+
+        assert(PyThreadState_Get() == save_tstate);
+    }
+
+    for (int i = 0; i < 10; ++i) {
+        assert(PyThreadState_Get() == save_tstate);
+        PyThreadState_Release();
+    }
+
+    assert(PyThreadState_GetUnchecked() == NULL);
+    PyInterpreterRef_Close(ref);
+    PyThreadState_Swap(save_tstate);
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+test_thread_state_ensure_crossinterp(PyObject *self, PyObject *unused)
+{
+    PyInterpreterRef ref = get_strong_ref();
+    PyThreadState *save_tstate = PyThreadState_Swap(NULL);
+    PyThreadState *interp_tstate = Py_NewInterpreter();
+    if (interp_tstate == NULL) {
+        PyInterpreterRef_Close(ref);
+        return PyErr_NoMemory();
+    }
+
+    /* This should create a new thread state for the calling interpreter, *not*
+       reactivate the old one. In a real-world scenario, this would arise in
+       something like this:
+
+       def some_func():
+           import something
+           # This re-enters the main interpreter, but we
+           # shouldn't have access to prior thread-locals.
+           something.call_something()
+
+       interp = interpreters.create()
+       interp.exec(some_func)
+       */
+    if (PyThreadState_Ensure(ref) < 0) {
+        PyInterpreterRef_Close(ref);
+        return PyErr_NoMemory();
+    }
+
+    PyThreadState *ensured_tstate = PyThreadState_Get();
+    assert(ensured_tstate != save_tstate);
+    assert(PyInterpreterState_Get() == PyInterpreterRef_AsInterpreter(ref));
+    assert(PyGILState_GetThisThreadState() == ensured_tstate);
+
+    // Now though, we should reactivate the thread state
+    if (PyThreadState_Ensure(ref) < 0) {
+        PyInterpreterRef_Close(ref);
+        return PyErr_NoMemory();
+    }
+
+    assert(PyThreadState_Get() == ensured_tstate);
+    PyThreadState_Release();
+
+    // Ensure that we're restoring the prior thread state
+    PyThreadState_Release();
+    assert(PyThreadState_Get() == interp_tstate);
+    assert(PyGILState_GetThisThreadState() == interp_tstate);
+
+    PyThreadState_Swap(interp_tstate);
+    Py_EndInterpreter(interp_tstate);
+
+    PyInterpreterRef_Close(ref);
+    PyThreadState_Swap(save_tstate);
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+test_weak_interpreter_ref_after_shutdown(PyObject *self, PyObject *unused)
+{
+    PyThreadState *save_tstate = PyThreadState_Swap(NULL);
+    PyInterpreterWeakRef wref;
+    PyThreadState *interp_tstate = Py_NewInterpreter();
+    if (interp_tstate == NULL) {
+        return PyErr_NoMemory();
+    }
+
+    int res = PyInterpreterWeakRef_Get(&wref);
+    (void)res;
+    assert(res == 0);
+
+    // As a sanity check, ensure that the weakref actually works
+    PyInterpreterRef ref;
+    res = PyInterpreterWeakRef_AsStrong(wref, &ref);
+    assert(res == 0);
+    PyInterpreterRef_Close(ref);
+
+    // Now, destroy the interpreter and try to acquire a weak reference.
+    // It should fail.
+    Py_EndInterpreter(interp_tstate);
+    res = PyInterpreterWeakRef_AsStrong(wref, &ref);
+    assert(res == -1);
+
+    PyThreadState_Swap(save_tstate);
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef TestMethods[] = {
     {"set_errno",               set_errno,                       METH_VARARGS},
     {"test_config",             test_config,                     METH_NOARGS},
@@ -2640,6 +2828,10 @@ static PyMethodDef TestMethods[] = {
     {"test_atexit", test_atexit, METH_NOARGS},
     {"code_offset_to_line", _PyCFunction_CAST(code_offset_to_line), METH_FASTCALL},
     {"toggle_reftrace_printer", toggle_reftrace_printer, METH_O},
+    {"test_interpreter_refs", test_interpreter_refs, METH_NOARGS},
+    {"test_thread_state_ensure_nested", test_thread_state_ensure_nested, METH_NOARGS},
+    {"test_thread_state_ensure_crossinterp", test_thread_state_ensure_crossinterp, METH_NOARGS},
+    {"test_weak_interpreter_ref_after_shutdown", test_weak_interpreter_ref_after_shutdown, METH_NOARGS},
     {NULL, NULL} /* sentinel */
 };
 

diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c
@@ -2345,6 +2345,60 @@ incref_decref_delayed(PyObject *self, PyObject *op)
     Py_RETURN_NONE;
 }
 
+#define NUM_REFS 100
+
+static PyObject *
+test_interp_refcount(PyObject *self, PyObject *unused)
+{
+    PyInterpreterState *interp = PyInterpreterState_Get();
+    assert(_PyInterpreterState_Refcount(interp) == 0);
+    PyInterpreterRef refs[NUM_REFS];
+    for (int i = 0; i < NUM_REFS; ++i) {
+        int res = PyInterpreterRef_Get(&refs[i]);
+        (void)res;
+        assert(res == 0);
+        assert(_PyInterpreterState_Refcount(interp) == i + 1);
+    }
+
+    for (int i = 0; i < NUM_REFS; ++i) {
+        PyInterpreterRef_Close(refs[i]);
+        assert(_PyInterpreterState_Refcount(interp) == (NUM_REFS - i - 1));
+    }
+
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+test_interp_weakref_incref(PyObject *self, PyObject *unused)
+{
+    PyInterpreterState *interp = PyInterpreterState_Get();
+    PyInterpreterWeakRef wref;
+    if (PyInterpreterWeakRef_Get(&wref) < 0) {
+        return NULL;
+    }
+    assert(_PyInterpreterState_Refcount(interp) == 0);
+
+    PyInterpreterRef refs[NUM_REFS];
+
+    for (int i = 0; i < NUM_REFS; ++i) {
+        int res = PyInterpreterWeakRef_AsStrong(wref, &refs[i]);
+        (void)res;
+        assert(res == 0);
+        assert(PyInterpreterRef_AsInterpreter(refs[i]) == interp);
+        assert(_PyInterpreterState_Refcount(interp) == i + 1);
+    }
+
+    for (int i = 0; i < NUM_REFS; ++i) {
+        PyInterpreterRef_Close(refs[i]);
+        assert(_PyInterpreterState_Refcount(interp) == (NUM_REFS - i - 1));
+    }
+
+    PyInterpreterWeakRef_Close(wref);
+    Py_RETURN_NONE;
+}
+
+#undef NUM_REFS
+
 static PyMethodDef module_functions[] = {
     {"get_configs", get_configs, METH_NOARGS},
     {"get_recursion_depth", get_recursion_depth, METH_NOARGS},
@@ -2447,6 +2501,8 @@ static PyMethodDef module_functions[] = {
     {"is_static_immortal", is_static_immortal, METH_O},
     {"incref_decref_delayed", incref_decref_delayed, METH_O},
     GET_NEXT_DICT_KEYS_VERSION_METHODDEF
+    {"test_interp_refcount", test_interp_refcount, METH_NOARGS},
+    {"test_interp_weakref_incref", test_interp_weakref_incref, METH_NOARGS},
     {NULL, NULL} /* sentinel */
 };