unsafe experiment: do not respect TLS access

python · neonene · Aug 17, 2024 · Aug 20, 2024 · Aug 20, 2024 · Aug 20, 2024
commit 6e62c38ce334b1a04778097ff232bad08ae4a5e9
diff --git a/Include/internal/pycore_pystate.h b/Include/internal/pycore_pystate.h
@@ -135,7 +135,7 @@ static inline PyThreadState*
 _PyThreadState_GET(void)
 {
 #if defined(HAVE_THREAD_LOCAL) && !defined(Py_BUILD_CORE_MODULE)
-    return _Py_tss_tstate;
+    return (PyThreadState*)_Py_atomic_load_ptr_relaxed(&_PyRuntime.tstate_current);
 #else
     return _PyThreadState_GetCurrent();
 #endif

diff --git a/Include/internal/pycore_runtime.h b/Include/internal/pycore_runtime.h
@@ -285,6 +285,7 @@ typedef struct pyruntimestate {
     struct _pythread_runtime_state threads;
     struct _signals_runtime_state signals;
 
+    PyThreadState *tstate_current;
     /* Used for the thread state bound to the current thread. */
     Py_tss_t autoTSSkey;
 

diff --git a/Python/pystate.c b/Python/pystate.c
@@ -79,23 +79,31 @@ current_fast_get(void)
 #endif
 }
 
+static inline PyThreadState *
+current_fast_get2(void)
+{
+    return (PyThreadState*)_Py_atomic_load_ptr_relaxed(&_PyRuntime.tstate_current);
+}
+
 static inline void
 current_fast_set(_PyRuntimeState *Py_UNUSED(runtime), PyThreadState *tstate)
 {
     assert(tstate != NULL);
 #ifdef HAVE_THREAD_LOCAL
     _Py_tss_tstate = tstate;
+    _Py_atomic_store_ptr_relaxed(&_PyRuntime.tstate_current, tstate);
 #else
     // XXX Fall back to the PyThread_tss_*() API.
 #  error "no supported thread-local variable storage classifier"
 #endif
 }
 
 static inline void
-current_fast_clear(_PyRuntimeState *Py_UNUSED(runtime))
+current_fast_clear(_PyRuntimeState *runtime)
 {
 #ifdef HAVE_THREAD_LOCAL
     _Py_tss_tstate = NULL;
+    _Py_atomic_store_ptr_relaxed(&runtime->tstate_current, NULL);
 #else
     // XXX Fall back to the PyThread_tss_*() API.
 #  error "no supported thread-local variable storage classifier"
@@ -110,7 +118,7 @@ current_fast_clear(_PyRuntimeState *Py_UNUSED(runtime))
 PyThreadState *
 _PyThreadState_GetCurrent(void)
 {
-    return current_fast_get();
+    return current_fast_get2();
 }
 
 
@@ -1331,7 +1339,7 @@ _PyInterpreterState_RequireIDRef(PyInterpreterState *interp, int required)
 PyInterpreterState*
 PyInterpreterState_Get(void)
 {
-    PyThreadState *tstate = current_fast_get();
+    PyThreadState *tstate = current_fast_get2();
     _Py_EnsureTstateNotNULL(tstate);
     PyInterpreterState *interp = tstate->interp;
     if (interp == NULL) {
@@ -2412,14 +2420,14 @@ PyThreadState_SetAsyncExc(unsigned long id, PyObject *exc)
 PyThreadState *
 PyThreadState_GetUnchecked(void)
 {
-    return current_fast_get();
+    return current_fast_get2();
 }
 
 
 PyThreadState *
 PyThreadState_Get(void)
 {
-    PyThreadState *tstate = current_fast_get();
+    PyThreadState *tstate = current_fast_get2();
     _Py_EnsureTstateNotNULL(tstate);
     return tstate;
 }