gh-121621: Move asyncio running loop to thread state (GH-121695)

Author: Fidget-Spinner

Include/cpython/pystate.h

@@ -68,6 +68,8 @@ struct _ts {
        pycore_ceval.h. */
     uintptr_t eval_breaker;
 
+    PyObject *asyncio_running_loop; // Strong reference
+
     struct {
         /* Has been initialized to a safe state.
 

Include/internal/pycore_global_objects_fini_generated.h

@@ -77,7 +77,6 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(__annotate__)
         STRUCT_FOR_ID(__annotations__)
         STRUCT_FOR_ID(__args__)
-        STRUCT_FOR_ID(__asyncio_running_event_loop__)
         STRUCT_FOR_ID(__await__)
         STRUCT_FOR_ID(__bases__)
         STRUCT_FOR_ID(__bool__)

Include/internal/pycore_global_strings.h

@@ -135,9 +135,6 @@ typedef struct {
     /* Imports from traceback. */
     PyObject *traceback_extract_stack;
 
-    PyObject *cached_running_loop; // Borrowed reference
-    volatile uint64_t cached_running_loop_tsid;
-
     /* Counter for autogenerated Task names */
     uint64_t task_name_counter;
 
@@ -321,101 +318,15 @@ get_future_loop(asyncio_state *state, PyObject *fut)
     return PyObject_GetAttr(fut, &_Py_ID(_loop));
 }
 
-
-static int
-get_running_loop(asyncio_state *state, PyObject **loop)
-{
-    PyObject *rl;
-
-    PyThreadState *ts = _PyThreadState_GET();
-    uint64_t ts_id = PyThreadState_GetID(ts);
-    if (state->cached_running_loop_tsid == ts_id &&
-        state->cached_running_loop != NULL)
-    {
-        // Fast path, check the cache.
-        rl = state->cached_running_loop;
-    }
-    else {
-        PyObject *ts_dict = _PyThreadState_GetDict(ts);  // borrowed
-        if (ts_dict == NULL) {
-            goto not_found;
-        }
-
-        rl = PyDict_GetItemWithError(
-            ts_dict, &_Py_ID(__asyncio_running_event_loop__));  // borrowed
-        if (rl == NULL) {
-            if (PyErr_Occurred()) {
-                goto error;
-            }
-            else {
-                goto not_found;
-            }
-        }
-
-        // TODO GH-121621: This should be moved to PyThreadState
-        // for easier and quicker access.
-        state->cached_running_loop = rl;
-        state->cached_running_loop_tsid = ts_id;
-    }
-
-
-    if (rl == Py_None) {
-        goto not_found;
-    }
-
-    *loop = Py_NewRef(rl);
-    return 0;
-
-not_found:
-    *loop = NULL;
-    return 0;
-
-error:
-    *loop = NULL;
-    return -1;
-}
-
-
-static int
-set_running_loop(asyncio_state *state, PyObject *loop)
-{
-    PyObject *ts_dict = NULL;
-
-    PyThreadState *tstate = _PyThreadState_GET();
-    if (tstate != NULL) {
-        ts_dict = _PyThreadState_GetDict(tstate);  // borrowed
-    }
-
-    if (ts_dict == NULL) {
-        PyErr_SetString(
-            PyExc_RuntimeError, "thread-local storage is not available");
-        return -1;
-    }
-    if (PyDict_SetItem(
-            ts_dict, &_Py_ID(__asyncio_running_event_loop__), loop) < 0)
-    {
-        return -1;
-    }
-
-
-    // TODO GH-121621: This should be moved to PyThreadState
-    // for easier and quicker access.
-    state->cached_running_loop = loop; // borrowed, kept alive by ts_dict
-    state->cached_running_loop_tsid = PyThreadState_GetID(tstate);
-
-    return 0;
-}
-
-
 static PyObject *
 get_event_loop(asyncio_state *state)
 {
     PyObject *loop;
     PyObject *policy;
 
-    if (get_running_loop(state, &loop)) {
-        return NULL;
-    }
+    PyThreadState *ts = _PyThreadState_GET();
+    loop = Py_XNewRef(ts->asyncio_running_loop);
+
     if (loop != NULL) {
         return loop;
     }
@@ -3367,11 +3278,8 @@ static PyObject *
 _asyncio__get_running_loop_impl(PyObject *module)
 /*[clinic end generated code: output=b4390af721411a0a input=0a21627e25a4bd43]*/
 {
-    PyObject *loop;
-    asyncio_state *state = get_asyncio_state(module);
-    if (get_running_loop(state, &loop)) {
-        return NULL;
-    }
+    PyThreadState *ts = _PyThreadState_GET();
+    PyObject *loop = Py_XNewRef(ts->asyncio_running_loop);
     if (loop == NULL) {
         /* There's no currently running event loop */
         Py_RETURN_NONE;
@@ -3394,10 +3302,11 @@ static PyObject *
 _asyncio__set_running_loop(PyObject *module, PyObject *loop)
 /*[clinic end generated code: output=ae56bf7a28ca189a input=4c9720233d606604]*/
 {
-    asyncio_state *state = get_asyncio_state(module);
-    if (set_running_loop(state, loop)) {
-        return NULL;
+    PyThreadState *ts = _PyThreadState_GET();
+    if (loop == Py_None) {
+        loop = NULL;
     }
+    Py_XSETREF(ts->asyncio_running_loop, Py_XNewRef(loop));
     Py_RETURN_NONE;
 }
 
@@ -3435,14 +3344,13 @@ _asyncio_get_running_loop_impl(PyObject *module)
 /*[clinic end generated code: output=c247b5f9e529530e input=2a3bf02ba39f173d]*/
 {
     PyObject *loop;
-    asyncio_state *state = get_asyncio_state(module);
-    if (get_running_loop(state, &loop)) {
-        return NULL;
-    }
+    PyThreadState *ts = _PyThreadState_GET();
+    loop = Py_XNewRef(ts->asyncio_running_loop);
     if (loop == NULL) {
         /* There's no currently running event loop */
         PyErr_SetString(
             PyExc_RuntimeError, "no running event loop");
+        return NULL;
     }
     return loop;
 }

Include/internal/pycore_runtime_init_generated.h

@@ -1499,6 +1499,8 @@ init_threadstate(_PyThreadStateImpl *_tstate,
     tstate->previous_executor = NULL;
     tstate->dict_global_version = 0;
 
+    tstate->asyncio_running_loop = NULL;
+
     tstate->delete_later = NULL;
 
     llist_init(&_tstate->mem_free_queue);
@@ -1700,6 +1702,8 @@ PyThreadState_Clear(PyThreadState *tstate)
 
     /* Don't clear tstate->pyframe: it is a borrowed reference */
 
+    Py_CLEAR(tstate->asyncio_running_loop);
+
     Py_CLEAR(tstate->dict);
     Py_CLEAR(tstate->async_exc);