Merge branch 'pep-788-impl' of https://github.com/ZeroIntensity/cpython into pep-788-impl

Author: ZeroIntensity

Include/cpython/pystate.h

@@ -209,11 +209,19 @@ 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;
+    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)?
 
-    /* Thread state that was active before PyThreadState_Ensure() was called. */
-    PyThreadState *prior_ensure;
+           This is only true for thread states created by PyThreadState_Ensure() */
+        int delete_on_release;
+    } ensure;
 };
 
 /* other API */
@@ -272,32 +280,36 @@ PyAPI_FUNC(void) _PyInterpreterState_SetEvalFrameFunc(
 
 typedef uintptr_t PyInterpreterRef;
 
-PyAPI_FUNC(PyInterpreterRef) PyInterpreterRef_Get(void);
+PyAPI_FUNC(int) PyInterpreterRef_Get(PyInterpreterRef *ref);
 PyAPI_FUNC(PyInterpreterRef) PyInterpreterRef_Dup(PyInterpreterRef ref);
-PyAPI_FUNC(int) PyInterpreterState_AsStrong(PyInterpreterState *interp, PyInterpreterRef *strong_ptr);
+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);                                \
+} while (0)
 
 /* Weak interpreter references */
 
-typedef struct _interpreter_weakref {
+typedef struct _PyInterpreterWeakRef {
     int64_t id;
     Py_ssize_t refcount;
-} PyInterpreterWeakRef;
+} _PyInterpreterWeakRef;
 
-PyAPI_FUNC(PyInterpreterWeakRef) PyInterpreterWeakRef_Get(void);
+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);
 
-// Exports for '_testcapi' shared extension
-PyAPI_FUNC(Py_ssize_t) _PyInterpreterState_Refcount(PyInterpreterState *interp);
-PyAPI_FUNC(void) _PyInterpreterState_Incref(PyInterpreterState *interp);
+#define PyInterpreterWeakRef_Close(ref) do {    \
+    PyInterpreterWeakRef_Close(ref);            \
+    ref = 0;                                    \
+} while (0)
+
 
 PyAPI_FUNC(int) PyThreadState_Ensure(PyInterpreterRef interp_ref);
 

Include/internal/pycore_pystate.h

@@ -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

Lib/test/test_embed.py

@@ -1914,14 +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):

Modules/_testcapimodule.c

@@ -2546,82 +2546,189 @@ toggle_reftrace_printer(PyObject *ob, PyObject *arg)
     Py_RETURN_NONE;
 }
 
-static PyObject *
-test_interp_refcount(PyObject *self, PyObject *unused)
+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 ref1;
-    PyInterpreterRef ref2;
-
-    // Reference counts are technically 0 by default
-    assert(_PyInterpreterState_Refcount(interp) == 0);
-    ref1 = PyInterpreterRef_Get();
-    assert(_PyInterpreterState_Refcount(interp) == 1);
-    ref2 = PyInterpreterRef_Get();
-    assert(_PyInterpreterState_Refcount(interp) == 2);
-    PyInterpreterRef_Close(ref1);
-    assert(_PyInterpreterState_Refcount(interp) == 1);
-    PyInterpreterRef_Close(ref2);
-    assert(_PyInterpreterState_Refcount(interp) == 0);
-
-    ref1 = PyInterpreterRef_Get();
-    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);
+    PyInterpreterRef ref = get_strong_ref();
+    assert(PyInterpreterRef_AsInterpreter(ref) == interp);
 
-    Py_RETURN_NONE;
+    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_interp_weak_ref(PyObject *self, PyObject *unused)
+test_interpreter_refs(PyObject *self, PyObject *unused)
 {
-    PyInterpreterState *interp = PyInterpreterState_Get();
-    PyInterpreterWeakRef wref = PyInterpreterWeakRef_Get();
-    assert(_PyInterpreterState_Refcount(interp) == 0);
+    // Test the main interpreter
+    test_interp_ref_common();
 
-    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);
+    // 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_interp_ensure(PyObject *self, PyObject *unused)
+test_thread_state_ensure_nested(PyObject *self, PyObject *unused)
 {
-    PyInterpreterState *interp = PyInterpreterState_Get();
-    PyInterpreterRef ref = PyInterpreterRef_Get();
+    PyInterpreterRef ref = get_strong_ref();
     PyThreadState *save_tstate = PyThreadState_Swap(NULL);
-    PyThreadState *tstate = Py_NewInterpreter();
-    PyInterpreterRef sub_ref = PyInterpreterRef_Get();
-    PyInterpreterState *subinterp = PyThreadState_GetInterpreter(tstate);
+    assert(PyGILState_GetThisThreadState() == save_tstate);
 
     for (int i = 0; i < 10; ++i) {
-        int res = PyThreadState_Ensure(ref);
-        assert(res == 0);
-        assert(PyInterpreterState_Get() == interp);
+        // 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) {
-        int res = PyThreadState_Ensure(sub_ref);
-        assert(res == 0);
-        assert(PyInterpreterState_Get() == subinterp);
+        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 < 20; ++i) {
+    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);
-    PyInterpreterRef_Close(sub_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;
@@ -2721,9 +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_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},
+    {"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 */
 };