python · ZeroIntensity · Apr 24, 2025 · Apr 24, 2025 · Apr 25, 2025 · Apr 25, 2025
@@ -206,6 +206,12 @@ struct _ts {
     */
     PyObject *threading_local_sentinel;
     _PyRemoteDebuggerSupport remote_debugger_support;
+
+    /* Number of nested PyThreadState_Ensure() calls on this thread state */
+    Py_ssize_t ensure_counter;
+
+    /* Thread state that was active before PyThreadState_Ensure() was called. */
+    PyThreadState *prior_ensure;
 };
 
 /* other API */
@@ -259,3 +265,45 @@ 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 *ptr);
+PyAPI_FUNC(PyInterpreterRef) PyInterpreterRef_Dup(PyInterpreterRef ref);
+PyAPI_FUNC(int) PyInterpreterRef_Main(PyInterpreterRef *strong_ptr);
+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)
+
+// Exports for '_testcapi' shared extension
+PyAPI_FUNC(Py_ssize_t) _PyInterpreterState_Refcount(PyInterpreterState *interp);
+PyAPI_FUNC(int) _PyInterpreterState_Incref(PyInterpreterState *interp);
+
+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
@@ -802,6 +802,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

diff --git a/Lib/test/test_embed.py b/Lib/test/test_embed.py
@@ -1914,10 +1914,12 @@ 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")
+
 
 class MiscTests(EmbeddingTestsMixin, unittest.TestCase):
     def test_unicode_id_init(self):

diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
@@ -2546,6 +2546,101 @@ 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 PyObject *
+test_interp_refcount(PyObject *self, PyObject *unused)
+{
+    PyInterpreterState *interp = PyInterpreterState_Get();
+    PyInterpreterRef ref1;
+    PyInterpreterRef ref2;
+
+    // Reference counts are technically 0 by default
+    assert(_PyInterpreterState_Refcount(interp) == 0);
+    ref1 = get_strong_ref();
+    assert(_PyInterpreterState_Refcount(interp) == 1);
+    ref2 = get_strong_ref();
+    assert(_PyInterpreterState_Refcount(interp) == 2);
+    PyInterpreterRef_Close(ref1);
+    assert(_PyInterpreterState_Refcount(interp) == 1);
+    PyInterpreterRef_Close(ref2);
+    assert(_PyInterpreterState_Refcount(interp) == 0);
+
+    ref1 = get_strong_ref();
+    ref2 = PyInterpreterRef_Dup(ref1);
+    assert(_PyInterpreterState_Refcount(interp) == 2);
+    assert(PyInterpreterRef_AsInterpreter(ref1) == interp);
+    assert(PyInterpreterRef_AsInterpreter(ref2) == interp);
+    PyInterpreterRef_Close(ref1);
+    PyInterpreterRef_Close(ref2);
+    assert(_PyInterpreterState_Refcount(interp) == 0);
+
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+test_interp_weak_ref(PyObject *self, PyObject *unused)
+{
+    PyInterpreterState *interp = PyInterpreterState_Get();
+    PyInterpreterWeakRef wref;
+    if (PyInterpreterWeakRef_Get(&wref) < 0) {
+        return NULL;
+    }
+    assert(_PyInterpreterState_Refcount(interp) == 0);
+
+    PyInterpreterRef ref;
+    int res = PyInterpreterWeakRef_AsStrong(wref, &ref);
+    assert(res == 0);
+    assert(PyInterpreterRef_AsInterpreter(ref) == interp);
+    assert(_PyInterpreterState_Refcount(interp) == 1);
+    PyInterpreterWeakRef_Close(wref);
+    PyInterpreterRef_Close(ref);
+
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+test_interp_ensure(PyObject *self, PyObject *unused)
+{
+    PyInterpreterState *interp = PyInterpreterState_Get();
+    PyInterpreterRef ref = get_strong_ref();
+    PyThreadState *save_tstate = PyThreadState_Swap(NULL);
+    PyThreadState *tstate = Py_NewInterpreter();
+    PyInterpreterRef sub_ref = get_strong_ref();
+    PyInterpreterState *subinterp = PyThreadState_GetInterpreter(tstate);
+
+    for (int i = 0; i < 10; ++i) {
+        int res = PyThreadState_Ensure(ref);
+        assert(res == 0);
+        assert(PyThreadState_GetUnchecked() != NULL);
+        assert(PyInterpreterState_Get() == interp);
+    }
+
+    for (int i = 0; i < 10; ++i) {
+        int res = PyThreadState_Ensure(sub_ref);
+        assert(res == 0);
+        assert(PyInterpreterState_Get() == subinterp);
+    }
+
+    for (int i = 0; i < 20; ++i) {
+        PyThreadState_Release();
+    }
+
+    PyInterpreterRef_Close(ref);
+    PyInterpreterRef_Close(sub_ref);
+
+    PyThreadState_Swap(save_tstate);
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef TestMethods[] = {
     {"set_errno",               set_errno,                       METH_VARARGS},
     {"test_config",             test_config,                     METH_NOARGS},
@@ -2640,6 +2735,9 @@ 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_interp_refcount", test_interp_refcount, METH_NOARGS},
+    {"test_interp_weak_ref", test_interp_weak_ref, METH_NOARGS},
+    {"test_interp_ensure", test_interp_ensure, METH_NOARGS},
     {NULL, NULL} /* sentinel */
 };
 

diff --git a/Programs/_testembed.c b/Programs/_testembed.c
@@ -2313,6 +2313,63 @@ test_get_incomplete_frame(void)
     return result;
 }
 
+const char *THREAD_CODE = \
+    "import time\n"
+    "time.sleep(0.2)\n"
+    "def fib(n):\n"
+    "  if n <= 1:\n"
+    "    return n\n"
+    "  else:\n"
+    "    return fib(n - 1) + fib(n - 2)\n"
+    "fib(10)";
+
+typedef struct {
+    PyInterpreterRef ref;
+    int done;
+} ThreadData;
+
+static void
+do_tstate_ensure(void *arg)
+{
+    ThreadData *data = (ThreadData *)arg;
+    int res = PyThreadState_Ensure(data->ref);
+    assert(res == 0);
+    PyThreadState_Ensure(data->ref);
+    PyThreadState_Ensure(data->ref);
+    PyGILState_STATE gstate = PyGILState_Ensure();
+    PyThreadState_Ensure(data->ref);
+    res = PyRun_SimpleString(THREAD_CODE);
+    PyThreadState_Release();
+    PyGILState_Release(gstate);
+    PyThreadState_Release();
+    PyThreadState_Release();
+    assert(res == 0);
+    PyThreadState_Release();
+    PyInterpreterRef_Close(data->ref);
+    data->done = 1;
+}
+
+static int
+test_thread_state_ensure(void)
+{
+    _testembed_initialize();
+    PyThread_handle_t handle;
+    PyThread_ident_t ident;
+    PyInterpreterRef ref;
+    if (PyInterpreterRef_Get(&ref) < 0) {
+        return -1;
+    };
+    ThreadData data = { ref };
+    if (PyThread_start_joinable_thread(do_tstate_ensure, &data,
+                                       &ident, &handle) < 0) {
+        PyInterpreterRef_Close(ref);
+        return -1;
+    }
+    Py_Finalize();
+    assert(data.done == 1);
+    return 0;
+}
+
 static void
 do_gilstate_ensure(void *event_ptr)
 {
@@ -2429,6 +2486,7 @@ static struct TestCase TestCases[] = {
     {"test_frozenmain", test_frozenmain},
 #endif
     {"test_get_incomplete_frame", test_get_incomplete_frame},
+    {"test_thread_state_ensure", test_thread_state_ensure},
     {"test_gilstate_after_finalization", test_gilstate_after_finalization},
     {NULL, NULL}
 };

@@ -13,6 +13,7 @@
 #include "pycore_freelist.h"      // _PyObject_ClearFreeLists()
 #include "pycore_global_objects_fini_generated.h"  // _PyStaticObjects_CheckRefcnt()
 #include "pycore_initconfig.h"    // _PyStatus_OK()
+#include "pycore_interp_structs.h"
 #include "pycore_interpolation.h" // _PyInterpolation_InitTypes()
 #include "pycore_long.h"          // _PyLong_InitTypes()
 #include "pycore_object.h"        // _PyDebug_PrintTotalRefs()
@@ -97,6 +98,7 @@ static PyStatus init_android_streams(PyThreadState *tstate);
 static PyStatus init_apple_streams(PyThreadState *tstate);
 #endif
 static void wait_for_thread_shutdown(PyThreadState *tstate);
+static void wait_for_interp_references(PyInterpreterState *interp);
 static void finalize_subinterpreters(void);
 static void call_ll_exitfuncs(_PyRuntimeState *runtime);
 
@@ -2022,6 +2024,9 @@ _Py_Finalize(_PyRuntimeState *runtime)
     // Wrap up existing "threading"-module-created, non-daemon threads.
     wait_for_thread_shutdown(tstate);
 
+    // Wait for the interpreter's reference count to reach zero
+    wait_for_interp_references(tstate->interp);
+
     // Make any remaining pending calls.
     _Py_FinishPendingCalls(tstate);
 
@@ -2438,6 +2443,9 @@ Py_EndInterpreter(PyThreadState *tstate)
     // Wrap up existing "threading"-module-created, non-daemon threads.
     wait_for_thread_shutdown(tstate);
 
+    // Wait for the interpreter's reference count to reach zero
+    wait_for_interp_references(tstate->interp);
+
     // Make any remaining pending calls.
     _Py_FinishPendingCalls(tstate);
 
@@ -3464,6 +3472,47 @@ wait_for_thread_shutdown(PyThreadState *tstate)
     Py_DECREF(threading);
 }
 
+/* Wait for the interpreter's reference count to reach zero.
+   See PEP 788. */
+static void
+wait_for_interp_references(PyInterpreterState *interp)
+{
+    assert(interp != NULL);
+    struct _Py_finalizing_threads *finalizing = &interp->threads.finalizing;
+    _Py_atomic_store_int_release(&finalizing->shutting_down, 1);
+    PyMutex_Lock(&finalizing->mutex);
+    if (_Py_atomic_load_ssize_relaxed(&finalizing->countdown) == 0) {
+        // Nothing to do.
+        PyMutex_Unlock(&finalizing->mutex);
+        return;
+    }
+    PyMutex_Unlock(&finalizing->mutex);
+
+    PyTime_t wait_max = 1000 * 1000 * 100; // 100 milliseconds
+    PyTime_t wait_ns = 1000; // 1 microsecond
+
+    while (true) {
+        if (PyEvent_WaitTimed(&finalizing->finished, wait_ns, 1)) {
+            // Event set
+            break;
+        }
+
+        wait_ns *= 2;
+        wait_ns = Py_MIN(wait_ns, wait_max);
+
+        if (PyErr_CheckSignals()) {
+            PyErr_Print();
+            /* The user CTRL+C'd us, bail out without waiting for a reference
+               count of zero.
+
+               This will probably cause threads to crash, but maybe that's
+               better than a deadlock. It might be worth intentionally
+               leaking subinterpreters to prevent some crashes here. */
+            break;
+        }
+    }
+}
+
 int Py_AtExit(void (*func)(void))
 {
     struct _atexit_runtime_state *state = &_PyRuntime.atexit;