python · vstinner · Jun 9, 2020
diff --git a/Include/internal/pycore_pystate.h b/Include/internal/pycore_pystate.h
@@ -11,6 +11,21 @@ extern "C" {
 #include "pycore_runtime.h"   /* PyRuntimeState */
 
 
+PyAPI_FUNC(void) _Py_NO_RETURN _Py_FatalError_TstateNULL(const char *func);
+
+static inline void
+_Py_EnsureFuncTstateNotNULL(const char *func, PyThreadState *tstate)
+{
+    if (tstate == NULL) {
+        _Py_FatalError_TstateNULL(func);
+    }
+}
+
+// Call Py_FatalError() if tstate is NULL
+#define _Py_EnsureTstateNotNULL(tstate) \
+    _Py_EnsureFuncTstateNotNULL(__func__, tstate)
+
+
 /* Check if the current thread is the main thread.
    Use _Py_IsMainInterpreter() to check if it's the main interpreter. */
 static inline int
@@ -53,6 +68,7 @@ _Py_ThreadCanHandlePendingCalls(void)
 PyAPI_FUNC(PyThreadState*) _PyThreadState_GetTSS(void);
 #endif
 
+
 static inline PyThreadState*
 _PyRuntimeState_GetThreadState(_PyRuntimeState *runtime)
 {
@@ -63,6 +79,32 @@ _PyRuntimeState_GetThreadState(_PyRuntimeState *runtime)
 #endif
 }
 
+
+static inline PyInterpreterState*
+_PyRuntimeState_GetInterpreter(_PyRuntimeState *runtime)
+{
+#ifdef EXPERIMENTAL_ISOLATED_SUBINTERPRETERS
+    PyThreadState *tstate = _PyRuntimeState_GetThreadState(runtime);
+#ifdef Py_DEBUG
+    _Py_EnsureTstateNotNULL(tstate);
+#endif
+    return tstate->interp;
+#else
+    PyInterpreterState* interp;
+    interp = (PyInterpreterState*)_Py_atomic_load_relaxed(&runtime->gilstate.interp_current);
+#ifdef Py_DEBUG
+    if (interp == NULL) {
+        _Py_FatalErrorFunc(__func__,
+                           "the function must be called with the GIL held, "
+                           "but the GIL is released "
+                           "(the current Python interpreter is NULL)");
+    }
+#endif
+    return interp;
+#endif
+}
+
+
 /* Get the current Python thread state.
 
    Efficient macro reading directly the 'gilstate.tstate_current' atomic
@@ -86,20 +128,6 @@ _PyThreadState_GET(void)
 #undef PyThreadState_GET
 #define PyThreadState_GET() _PyThreadState_GET()
 
-PyAPI_FUNC(void) _Py_NO_RETURN _Py_FatalError_TstateNULL(const char *func);
-
-static inline void
-_Py_EnsureFuncTstateNotNULL(const char *func, PyThreadState *tstate)
-{
-    if (tstate == NULL) {
-        _Py_FatalError_TstateNULL(func);
-    }
-}
-
-// Call Py_FatalError() if tstate is NULL
-#define _Py_EnsureTstateNotNULL(tstate) \
-    _Py_EnsureFuncTstateNotNULL(__func__, tstate)
-
 
 /* Get the current interpreter state.
 
@@ -110,11 +138,7 @@ _Py_EnsureFuncTstateNotNULL(const char *func, PyThreadState *tstate)
    See also _PyInterpreterState_Get()
    and _PyGILState_GetInterpreterStateUnsafe(). */
 static inline PyInterpreterState* _PyInterpreterState_GET(void) {
-    PyThreadState *tstate = _PyThreadState_GET();
-#ifdef Py_DEBUG
-    _Py_EnsureTstateNotNULL(tstate);
-#endif
-    return tstate->interp;
+    return _PyRuntimeState_GetInterpreter(&_PyRuntime);
 }
 
 

diff --git a/Include/internal/pycore_runtime.h b/Include/internal/pycore_runtime.h
@@ -33,6 +33,9 @@ struct _gilstate_runtime_state {
     /* Assuming the current thread holds the GIL, this is the
        PyThreadState for the current thread. */
     _Py_atomic_address tstate_current;
+    /* Assuming the current thread holds the GIL, this is the
+       PyInterpreterState of tstate_current. */
+    _Py_atomic_address interp_current;
     /* The single PyInterpreterState used by this process'
        GILState implementation
     */

diff --git a/Python/pystate.c b/Python/pystate.c
@@ -38,6 +38,9 @@ extern "C" {
 #define _PyRuntimeGILState_SetThreadState(gilstate, value) \
     _Py_atomic_store_relaxed(&(gilstate)->tstate_current, \
                              (uintptr_t)(value))
+#define _PyRuntimeGILState_SetInterpreter(gilstate, value) \
+    _Py_atomic_store_relaxed(&(gilstate)->interp_current, \
+                             (uintptr_t)(value))
 
 /* Forward declarations */
 static PyThreadState *_PyGILState_GetThisThreadState(struct _gilstate_runtime_state *gilstate);
@@ -889,6 +892,7 @@ _PyThreadState_DeleteCurrent(PyThreadState *tstate)
     struct _gilstate_runtime_state *gilstate = &tstate->interp->runtime->gilstate;
     tstate_delete_common(tstate, gilstate);
     _PyRuntimeGILState_SetThreadState(gilstate, NULL);
+    _PyRuntimeGILState_SetInterpreter(gilstate, NULL);
     _PyEval_ReleaseLock(tstate);
     PyMem_RawFree(tstate);
 }
@@ -978,6 +982,9 @@ _PyThreadState_Swap(struct _gilstate_runtime_state *gilstate, PyThreadState *new
 #endif
 
     _PyRuntimeGILState_SetThreadState(gilstate, newts);
+    PyInterpreterState* interp = (newts != NULL) ? newts->interp : NULL;
+    _PyRuntimeGILState_SetInterpreter(gilstate, interp);
+
     /* It should not be possible for more than one thread state
        to be used for a thread.  Check this the best we can in debug
        builds.