Implement new version of PyThreadState_Ensure() and PyThreadState_Release()

Author: ZeroIntensity

Include/cpython/pystate.h

@@ -61,12 +61,6 @@ 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 */
 
@@ -213,10 +207,11 @@ struct _ts {
     PyObject *threading_local_sentinel;
     _PyRemoteDebuggerSupport remote_debugger_support;
 
-    /* Whether this thread hangs when the interpreter is finalizing. */
-    uint8_t daemon;
+    /* Number of nested PyThreadState_Ensure() calls on this thread state */
+    Py_ssize_t ensure_counter;
 
-    _Py_ensured_tstate *ensured;
+    /* Thread state that was active before PyThreadState_Ensure() was called. */
+    PyThreadState *prior_ensure;
 };
 
 /* other API */
@@ -279,6 +274,7 @@ PyAPI_FUNC(PyInterpreterRef) PyInterpreterRef_Get(void);
 PyAPI_FUNC(PyInterpreterRef) PyInterpreterRef_Dup(PyInterpreterRef ref);
 PyAPI_FUNC(int) PyInterpreterState_AsStrong(PyInterpreterState *interp, 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);            \
@@ -301,6 +297,6 @@ PyAPI_FUNC(void) PyInterpreterWeakRef_Close(PyInterpreterWeakRef wref);
 PyAPI_FUNC(Py_ssize_t) _PyInterpreterState_Refcount(PyInterpreterState *interp);
 PyAPI_FUNC(void) _PyInterpreterState_Incref(PyInterpreterState *interp);
 
-PyAPI_FUNC(int) PyThreadState_Ensure(PyInterpreterState *interp);
+PyAPI_FUNC(int) PyThreadState_Ensure(PyInterpreterRef interp_ref);
 
 PyAPI_FUNC(void) PyThreadState_Release(void);

Python/pystate.c

@@ -3188,6 +3188,13 @@ PyInterpreterRef_Close(PyInterpreterRef ref)
     decref_interpreter(interp);
 }
 
+PyInterpreterState *
+PyInterpreterRef_AsInterpreter(PyInterpreterRef ref)
+{
+    PyInterpreterState *interp = ref_as_interp(ref);
+    return interp;
+}
+
 PyInterpreterWeakRef
 PyInterpreterWeakRef_Get(void)
 {
@@ -3260,13 +3267,56 @@ PyInterpreterState_AsStrong(PyInterpreterState *interp, PyInterpreterRef *strong
 }
 
 int
-PyThreadState_Ensure(PyInterpreterState *interp)
+PyThreadState_Ensure(PyInterpreterRef interp_ref)
 {
+    PyInterpreterState *interp = ref_as_interp(interp_ref);
+    PyThreadState *attached_tstate = current_fast_get();
+    if (attached_tstate != NULL && attached_tstate->interp == interp) {
+        /* Yay! We already have an attached thread state that matches. */
+        ++attached_tstate->ensure_counter;
+        return 0;
+    }
+
+    PyThreadState *detached_gilstate = gilstate_get();
+    if (detached_gilstate != NULL && detached_gilstate->interp == interp) {
+        /* There's a detached thread state that works. */
+        assert(attached_tstate == NULL);
+        ++detached_gilstate->ensure_counter;
+        _PyThreadState_Attach(detached_gilstate);
+        return 0;
+    }
+
+    PyThreadState *fresh_tstate = _PyThreadState_NewBound(interp,
+                                                          _PyThreadState_WHENCE_GILSTATE);
+    if (fresh_tstate == NULL) {
+        return -1;
+    }
+
+    if (attached_tstate != NULL) {
+        fresh_tstate->ensure_counter = 1;
+        fresh_tstate->prior_ensure = PyThreadState_Swap(fresh_tstate);
+    } else {
+        _PyThreadState_Attach(fresh_tstate);
+    }
+
     return 0;
 }
 
 void
 PyThreadState_Release(void)
 {
+    PyThreadState *tstate = current_fast_get();
+    _Py_EnsureTstateNotNULL(tstate);
+    Py_ssize_t remaining = --tstate->ensure_counter;
+    if (remaining < 0) {
+        Py_FatalError("PyThreadState_Release() called more times than PyThreadState_Ensure()");
+    }
+    if (remaining == 0) {
+        PyThreadState *to_restore = tstate->prior_ensure;
+        PyThreadState_Clear(tstate);
+        PyThreadState_Swap(to_restore);
+        PyThreadState_Delete(tstate);
+    }
+
     return;
 }