From de2be447b3083975cf84095bef5d6505b396887f Mon Sep 17 00:00:00 2001
From: Sam Gross <colesbury@gmail.com>
Date: Mon, 30 Aug 2021 14:31:18 -0700
Subject: [PATCH] Make PyThreadState_GET thread-local

---
 Include/internal/pycore_pystate.h | 28 +++++++++++++++++++++-------
 Include/internal/pycore_runtime.h |  3 ---
 Python/pystate.c                  | 18 ++++++++----------
 3 files changed, 29 insertions(+), 20 deletions(-)

diff --git a/Include/internal/pycore_pystate.h b/Include/internal/pycore_pystate.h
index 3d6d400f74d..c13a777a80d 100644
--- a/Include/internal/pycore_pystate.h
+++ b/Include/internal/pycore_pystate.h
@@ -63,12 +63,10 @@ _Py_ThreadCanHandlePendingCalls(void)
 
 /* Variable and macro for in-line access to current thread
    and interpreter state */
-
-static inline PyThreadState*
-_PyRuntimeState_GetThreadState(_PyRuntimeState *runtime)
-{
-    return (PyThreadState*)_Py_atomic_load_relaxed(&runtime->gilstate.tstate_current);
-}
+#if defined(__GNUC__) && !defined(Py_ENABLE_SHARED)
+__attribute__((tls_model("local-exec")))
+#endif
+extern Py_DECL_THREAD PyThreadState *_Py_current_tstate;
 
 /* Get the current Python thread state.
 
@@ -82,7 +80,23 @@ _PyRuntimeState_GetThreadState(_PyRuntimeState *runtime)
 static inline PyThreadState*
 _PyThreadState_GET(void)
 {
-    return _PyRuntimeState_GetThreadState(&_PyRuntime);
+#if defined(Py_BUILD_CORE_MODULE)
+    return _PyThreadState_UncheckedGet();
+#else
+    return _Py_current_tstate;
+#endif
+}
+
+static inline void
+_PyThreadState_SET(PyThreadState *tstate)
+{
+    _Py_current_tstate = tstate;
+}
+
+static inline PyThreadState*
+_PyRuntimeState_GetThreadState(_PyRuntimeState *runtime)
+{
+    return _PyThreadState_GET();
 }
 
 static inline void
diff --git a/Include/internal/pycore_runtime.h b/Include/internal/pycore_runtime.h
index d100e836c7d..442d64b7c8c 100644
--- a/Include/internal/pycore_runtime.h
+++ b/Include/internal/pycore_runtime.h
@@ -41,9 +41,6 @@ struct _gilstate_runtime_state {
     /* bpo-26558: Flag to disable PyGILState_Check().
        If set to non-zero, PyGILState_Check() always return 1. */
     int check_enabled;
-    /* Assuming the current thread holds the GIL, this is the
-       PyThreadState for the current thread. */
-    _Py_atomic_address tstate_current;
     /* The single PyInterpreterState used by this process'
        GILState implementation
     */
diff --git a/Python/pystate.c b/Python/pystate.c
index f2f571faf40..f0339b9d9e3 100644
--- a/Python/pystate.c
+++ b/Python/pystate.c
@@ -37,11 +37,8 @@ to avoid the expense of doing their own locking).
 extern "C" {
 #endif
 
-#define _PyRuntimeGILState_GetThreadState(gilstate) \
-    ((PyThreadState*)_Py_atomic_load_relaxed(&(gilstate)->tstate_current))
-#define _PyRuntimeGILState_SetThreadState(gilstate, value) \
-    _Py_atomic_store_relaxed(&(gilstate)->tstate_current, \
-                             (uintptr_t)(value))
+#define _PyRuntimeGILState_GetThreadState(gilstate) _PyThreadState_GET()
+#define _PyRuntimeGILState_SetThreadState(gilstate, value) _PyThreadState_SET(value)
 
 /* Forward declarations */
 static PyThreadState *_PyGILState_GetThisThreadState(struct _gilstate_runtime_state *gilstate);
@@ -58,6 +55,9 @@ _Py_COMP_DIAG_IGNORE_DEPR_DECLS
 static const _PyRuntimeState initial = _PyRuntimeState_INIT(_PyRuntime);
 _Py_COMP_DIAG_POP
 
+Py_DECL_THREAD PyThreadState *_Py_current_tstate;
+
+
 static int
 alloc_for_runtime(PyThread_type_lock *plock1, PyThread_type_lock *plock2,
                   PyThread_type_lock *plock3, PyThread_type_lock *plock4)
@@ -1164,8 +1164,7 @@ _PyThreadState_DeleteCurrent(PyThreadState *tstate)
 void
 PyThreadState_DeleteCurrent(void)
 {
-    struct _gilstate_runtime_state *gilstate = &_PyRuntime.gilstate;
-    PyThreadState *tstate = _PyRuntimeGILState_GetThreadState(gilstate);
+    PyThreadState *tstate = _PyThreadState_GET();
     _PyThreadState_DeleteCurrent(tstate);
 }
 
@@ -1531,9 +1530,8 @@ static int
 PyThreadState_IsCurrent(PyThreadState *tstate)
 {
     /* Must be the tstate for this thread */
-    struct _gilstate_runtime_state *gilstate = &_PyRuntime.gilstate;
-    assert(_PyGILState_GetThisThreadState(gilstate) == tstate);
-    return tstate == _PyRuntimeGILState_GetThreadState(gilstate);
+    assert(_PyGILState_GetThisThreadState(&_PyRuntime.gilstate) == tstate);
+    return tstate == _PyThreadState_GET();
 }
 
 /* Internal initialization/finalization functions called by