Use a thread state cache instead of creating a new one for each access.

Author: ZeroIntensity

Include/internal/pycore_pystate.h

@@ -231,6 +231,9 @@ extern PyThreadState * _PyThreadState_RemoveExcept(PyThreadState *tstate);
 extern void _PyThreadState_DeleteList(PyThreadState *list, int is_after_fork);
 extern void _PyThreadState_ClearMimallocHeaps(PyThreadState *tstate);
 
+// Export for '_interpreters' shared extension
+PyAPI_FUNC(PyThreadState *) _PyThreadState_NewForExec(PyInterpreterState *interp);
+
 // Export for '_testinternalcapi' shared extension
 PyAPI_FUNC(PyObject*) _PyThreadState_GetDict(PyThreadState *tstate);
 

Lib/test/test_interpreters/test_object_proxy.py

@@ -7,7 +7,7 @@
 import_helper.import_module("_interpreters")
 from concurrent.interpreters import share, SharedObjectProxy
 from test.test_interpreters.utils import TestBase
-from threading import Barrier, Thread, Lock
+from threading import Barrier, Thread, Lock, local
 from concurrent import interpreters
 from contextlib import contextmanager
 
@@ -242,6 +242,20 @@ def thread(interp):
 
         self.run_concurrently(thread, proxy=proxy)
 
+    def test_retain_thread_local_variables(self):
+        thread_local = local()
+        thread_local.value = 42
+
+        def test():
+            old = thread_local.value
+            thread_local.value = 24
+            return old
+
+        proxy = share(test)
+        with self.create_interp(proxy=proxy) as interp:
+            interp.exec("assert proxy() == 42")
+            self.assertEqual(thread_local.value, 24)
+
 
 if __name__ == "__main__":
     unittest.main()

Modules/_interpretersmodule.c

@@ -414,8 +414,7 @@ _sharedobjectproxy_enter(SharedObjectProxy *self, _PyXI_proxy_state *state)
         return 0;
     }
     state->to_restore = tstate;
-    PyThreadState *for_call = _PyThreadState_NewBound(self->interp,
-                                                      _PyThreadState_WHENCE_EXEC);
+    PyThreadState *for_call = _PyThreadState_NewForExec(self->interp);
     state->for_call = for_call;
     if (for_call == NULL) {
         PyErr_NoMemory();
@@ -446,9 +445,10 @@ _sharedobjectproxy_exit(SharedObjectProxy *self, _PyXI_proxy_state *state)
     }
 
     assert(state->for_call == _PyThreadState_GET());
-    PyThreadState_Clear(state->for_call);
     PyThreadState_Swap(state->to_restore);
-    PyThreadState_Delete(state->for_call);
+    // If we created a new thread state, we don't want to delete it.
+    // It's likely that it will be used again, but if not, the interpreter
+    // will clean it up at the end anyway.
 
     if (should_throw) {
         _PyErr_SetString(state->to_restore, PyExc_RuntimeError, "exception in interpreter");

Python/pystate.c

@@ -80,6 +80,56 @@ _Py_thread_local PyThreadState *_Py_tss_gilstate = NULL;
    and is same as tstate->interp. */
 _Py_thread_local PyInterpreterState *_Py_tss_interp = NULL;
 
+/* The last thread state used for each interpreter by this thread. */
+_Py_thread_local _Py_hashtable_t *_Py_tss_tstate_map = NULL;
+
+// TODO: Let's add a way to use _Py_hashtable_t statically to avoid the
+// extra heap allocation.
+
+static void
+mark_thread_state_used(PyThreadState *tstate)
+{
+    assert(tstate != NULL);
+    if (_Py_tss_tstate_map == NULL) {
+        _Py_hashtable_allocator_t alloc = {
+            .malloc = PyMem_RawMalloc,
+            .free = PyMem_RawFree
+        };
+        _Py_tss_tstate_map = _Py_hashtable_new_full(_Py_hashtable_hash_ptr,
+                                                    _Py_hashtable_compare_direct,
+                                                    NULL, NULL, &alloc);
+        if (_Py_tss_tstate_map == NULL) {
+            return;
+        }
+    }
+
+    (void)_Py_hashtable_steal(_Py_tss_tstate_map, tstate->interp);
+    (void)_Py_hashtable_set(_Py_tss_tstate_map, tstate->interp, tstate);
+}
+
+static PyThreadState *
+last_thread_state_for_interp(PyInterpreterState *interp)
+{
+    assert(interp != NULL);
+    if (_Py_tss_tstate_map == NULL) {
+        return NULL;
+    }
+
+    return _Py_hashtable_get(_Py_tss_tstate_map, interp);
+}
+
+static void
+mark_thread_state_dead(PyThreadState *tstate)
+{
+    if (_Py_tss_tstate_map == NULL) {
+        return;
+    }
+
+    if (tstate == _Py_hashtable_get(_Py_tss_tstate_map, tstate->interp)) {
+        (void)_Py_hashtable_steal(_Py_tss_tstate_map, tstate->interp);
+    }
+}
+
 static inline PyThreadState *
 current_fast_get(void)
 {
@@ -1603,6 +1653,21 @@ _PyThreadState_NewBound(PyInterpreterState *interp, int whence)
     return tstate;
 }
 
+/* Get the last thread state used for this interpreter, or create a new
+ * one if none exists.
+ * The thread state returned by this may or may not be attached. */
+PyThreadState *
+_PyThreadState_NewForExec(PyInterpreterState *interp)
+{
+    assert(interp != NULL);
+    PyThreadState *cached = last_thread_state_for_interp(interp);
+    if (cached != NULL) {
+        return cached;
+    }
+
+    return _PyThreadState_NewBound(interp, _PyThreadState_WHENCE_EXEC);
+}
+
 // This must be followed by a call to _PyThreadState_Bind();
 PyThreadState *
 _PyThreadState_New(PyInterpreterState *interp, int whence)
@@ -1649,6 +1714,7 @@ PyThreadState_Clear(PyThreadState *tstate)
     // disabled.
     // XXX assert(!_PyThreadState_IsRunningMain(tstate));
     // XXX assert(!tstate->_status.bound || tstate->_status.unbound);
+    mark_thread_state_dead(tstate);
     tstate->_status.finalizing = 1;  // just in case
 
     /* XXX Conditions we need to enforce:
@@ -1961,7 +2027,6 @@ _PyThreadState_DeleteList(PyThreadState *list, int is_after_fork)
     }
 }
 
-
 //----------
 // accessors
 //----------
@@ -2168,6 +2233,8 @@ _PyThreadState_Attach(PyThreadState *tstate)
 #if defined(Py_DEBUG)
     errno = err;
 #endif
+
+    mark_thread_state_used(tstate);
 }
 
 static void