Implement PEP 788

Include/cpython/pystate.h

@@ -63,6 +63,12 @@ typedef struct _stack_chunk {
     PyObject * data[1]; /* Variable sized */
 } _PyStackChunk;
 
+typedef struct _ensured_tstate {
+    struct _ensured_tstate *next;
+    PyThreadState *prior_tstate;
+    uint8_t was_daemon;
+} _Py_ensured_tstate;
+
 struct _ts {
     /* See Python/ceval.c for comments explaining most fields */
 
@@ -210,6 +216,11 @@ struct _ts {
     */
     PyObject *threading_local_sentinel;
     _PyRemoteDebuggerSupport remote_debugger_support;
+
+    /* Whether this thread hangs when the interpreter is finalizing. */
+    uint8_t daemon;
+
+    _Py_ensured_tstate *ensured;
 };
 
 # define Py_C_RECURSION_LIMIT 5000
@@ -265,3 +276,24 @@ PyAPI_FUNC(_PyFrameEvalFunction) _PyInterpreterState_GetEvalFrameFunc(
 PyAPI_FUNC(void) _PyInterpreterState_SetEvalFrameFunc(
     PyInterpreterState *interp,
     _PyFrameEvalFunction eval_frame);
+
+/* Similar to PyInterpreterState_Get(), but returns the interpreter with an
+ * incremented reference count. PyInterpreterState_Delete() won't delete the
+ * full interpreter structure until the reference is released by
+ * PyThreadState_Ensure() or PyInterpreterState_Release(). */
+PyAPI_FUNC(PyInterpreterState *) PyInterpreterState_Hold(void);
+
+PyAPI_FUNC(PyInterpreterState *) PyInterpreterState_Lookup(int64_t interp_id);
+
+/* Release a reference to an interpreter incremented by PyInterpreterState_Hold() */
+PyAPI_FUNC(void) PyInterpreterState_Release(PyInterpreterState *interp);
+
+// Exports for '_testcapi' shared extension
+PyAPI_FUNC(Py_ssize_t) _PyInterpreterState_Refcount(PyInterpreterState *interp);
+PyAPI_FUNC(void) _PyInterpreterState_Incref(PyInterpreterState *interp);
+
+PyAPI_FUNC(int) PyThreadState_SetDaemon(int daemon);
+
+PyAPI_FUNC(int) PyThreadState_Ensure(PyInterpreterState *interp);
+
+PyAPI_FUNC(void) PyThreadState_Release(void);

Include/internal/pycore_interp_structs.h

@@ -792,6 +792,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

Modules/_testcapimodule.c

@@ -2546,6 +2546,81 @@ toggle_reftrace_printer(PyObject *ob, PyObject *arg)
     Py_RETURN_NONE;
 }
 
+static PyObject *
+test_interp_refcount(PyObject *self, PyObject *unused)
+{
+    PyThreadState *save = PyThreadState_Get();
+    PyThreadState *tstate = Py_NewInterpreter();
+    assert(tstate == PyThreadState_Get());
+    PyInterpreterState *interp = PyThreadState_GetInterpreter(tstate);
+    assert(interp != NULL);
+
+    assert(_PyInterpreterState_Refcount(interp) == 0);
+    PyInterpreterState *held = PyInterpreterState_Hold();
+    assert(_PyInterpreterState_Refcount(interp) == 1);
+    PyInterpreterState_Release(held);
+    assert(_PyInterpreterState_Refcount(interp) == 0);
+
+    held = PyInterpreterState_Hold();
+    Py_EndInterpreter(tstate);
+    PyThreadState_Swap(save);
+    assert(_PyInterpreterState_Refcount(interp) == 1);
+    PyInterpreterState_Release(held);
+    assert(_PyInterpreterState_Refcount(interp) == 0);
+
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+test_interp_lookup(PyObject *self, PyObject *unused)
+{
+    PyInterpreterState *interp = PyInterpreterState_Get();
+    assert(_PyInterpreterState_Refcount(interp) == 0);
+    int64_t interp_id = PyInterpreterState_GetID(interp);
+    PyInterpreterState *ref = PyInterpreterState_Lookup(interp_id);
+    assert(ref == interp);
+    assert(_PyInterpreterState_Refcount(interp) == 1);
+    PyInterpreterState_Release(ref);
+    assert(PyInterpreterState_Lookup(10000) == NULL);
+    Py_BEGIN_ALLOW_THREADS;
+    ref = PyInterpreterState_Lookup(interp_id);
+    assert(ref == interp);
+    PyInterpreterState_Release(ref);
+    Py_END_ALLOW_THREADS;
+
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+test_interp_ensure(PyObject *self, PyObject *unused)
+{
+    PyInterpreterState *interp = PyInterpreterState_Get();
+    PyThreadState *save_tstate = PyThreadState_Swap(NULL);
+    PyThreadState *tstate = Py_NewInterpreter();
+    PyInterpreterState *subinterp = PyThreadState_GetInterpreter(tstate);
+
+    for (int i = 0; i < 10; ++i) {
+        _PyInterpreterState_Incref(interp);
+        int res = PyThreadState_Ensure(interp);
+        assert(res == 0);
+        assert(PyInterpreterState_Get() == interp);
+    }
+
+    for (int i = 0; i < 10; ++i) {
+        _PyInterpreterState_Incref(subinterp);
+        int res = PyThreadState_Ensure(subinterp);
+        assert(res == 0);
+        assert(PyInterpreterState_Get() == subinterp);
+    }
+
+    for (int i = 0; i < 20; ++i) {
+        PyThreadState_Release();
+    }
+
+    PyThreadState_Swap(save_tstate);
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef TestMethods[] = {
     {"set_errno",               set_errno,                       METH_VARARGS},
     {"test_config",             test_config,                     METH_NOARGS},
@@ -2640,6 +2715,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_lookup", test_interp_lookup, METH_NOARGS},
+    {"test_interp_ensure", test_interp_ensure, METH_NOARGS},
     {NULL, NULL} /* sentinel */
 };
 

Programs/_testembed.c

@@ -2341,6 +2341,56 @@ 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 {
+    PyInterpreterState *interp;
+    int done;
+} ThreadData;
+
+static void
+do_tstate_ensure(void *arg)
+{
+    ThreadData *data = (ThreadData *)arg;
+    int res = PyThreadState_Ensure(data->interp);
+    assert(res == 0);
+    PyThreadState_Ensure(PyInterpreterState_Hold());
+    PyThreadState_Ensure(PyInterpreterState_Hold());
+    PyGILState_STATE gstate = PyGILState_Ensure();
+    PyThreadState_Ensure(PyInterpreterState_Hold());
+    res = PyRun_SimpleString(THREAD_CODE);
+    PyThreadState_Release();
+    PyGILState_Release(gstate);
+    PyThreadState_Release();
+    PyThreadState_Release();
+    assert(res == 0);
+    PyThreadState_Release();
+    data->done = 1;
+}
+
+static int
+test_thread_state_ensure(void)
+{
+    _testembed_Py_InitializeFromConfig();
+    PyThread_handle_t handle;
+    PyThread_ident_t ident;
+    ThreadData data = { PyInterpreterState_Hold() };
+    if (PyThread_start_joinable_thread(do_tstate_ensure, &data,
+                                       &ident, &handle) < 0) {
+        PyInterpreterState_Release(data.interp);
+        return -1;
+    }
+    Py_Finalize();
+    assert(data.done == 1);
+    return 0;
+}
 
 /* *********************************************************
  * List of test cases and the function that implements it.
@@ -2431,6 +2481,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},
 
     {NULL, NULL}
 };

Python/pylifecycle.c

@@ -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_long.h"          // _PyLong_InitTypes()
 #include "pycore_object.h"        // _PyDebug_PrintTotalRefs()
 #include "pycore_obmalloc.h"      // _PyMem_init_obmalloc()
@@ -96,6 +97,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_native_shutdown(PyInterpreterState *interp);
 static void finalize_subinterpreters(void);
 static void call_ll_exitfuncs(_PyRuntimeState *runtime);
 
@@ -2012,6 +2014,9 @@ _Py_Finalize(_PyRuntimeState *runtime)
     // Wrap up existing "threading"-module-created, non-daemon threads.
     wait_for_thread_shutdown(tstate);
 
+    // Wrap up non-daemon native threads
+    wait_for_native_shutdown(tstate->interp);
+
     // Make any remaining pending calls.
     _Py_FinishPendingCalls(tstate);
 
@@ -2428,6 +2433,9 @@ Py_EndInterpreter(PyThreadState *tstate)
     // Wrap up existing "threading"-module-created, non-daemon threads.
     wait_for_thread_shutdown(tstate);
 
+    // Wrap up non-daemon native threads
+    wait_for_native_shutdown(tstate->interp);
+
     // Make any remaining pending calls.
     _Py_FinishPendingCalls(tstate);
 
@@ -3454,6 +3462,25 @@ wait_for_thread_shutdown(PyThreadState *tstate)
     Py_DECREF(threading);
 }
 
+/* Wait for all non-daemon native threads to finish.
+   See PEP 788. */
+static void
+wait_for_native_shutdown(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);
+
+    PyEvent_Wait(&finalizing->finished);
+}
+
 int Py_AtExit(void (*func)(void))
 {
     struct _atexit_runtime_state *state = &_PyRuntime.atexit;