gh-121621: Move asyncio running loop to thread state

Include/cpython/pystate.h

@@ -164,6 +164,9 @@ struct _ts {
     PyObject *async_gen_firstiter;
     PyObject *async_gen_finalizer;
 
+    PyObject *asyncio_cached_running_loop; // Borrowed reference
+    volatile uint64_t asyncio_cached_running_loop_tsid;
+
     PyObject *context;
     uint64_t context_ver;
 

Include/internal/pycore_freelist.h

@@ -124,6 +124,14 @@ struct _Py_object_stack_freelist {
     Py_ssize_t numfree;
 };
 
+
+struct _Py_asyncmodule_futureiter_freelist {
+#ifdef WITH_FREELISTS
+    struct futureiterobject *fi_freelist;
+    Py_ssize_t fi_freelist_len;
+#endif
+};
+
 struct _Py_object_freelists {
     struct _Py_float_freelist floats;
     struct _Py_tuple_freelist tuples;
@@ -135,6 +143,7 @@ struct _Py_object_freelists {
     struct _Py_async_gen_freelist async_gens;
     struct _Py_async_gen_asend_freelist async_gen_asends;
     struct _Py_object_stack_freelist object_stacks;
+    struct _Py_asyncmodule_futureiter_freelist futureiters;
 };
 
 extern void _PyObject_ClearFreeLists(struct _Py_object_freelists *freelists, int is_finalization);
@@ -147,6 +156,39 @@ extern void _PyAsyncGen_ClearFreeLists(struct _Py_object_freelists *freelists, i
 extern void _PyContext_ClearFreeList(struct _Py_object_freelists *freelists, int is_finalization);
 extern void _PyObjectStackChunk_ClearFreeList(struct _Py_object_freelists *freelists, int is_finalization);
 
+// Keep in sync with _asynciomodule.c !
+typedef struct futureiterobject_dummy {
+    PyObject_HEAD
+    void *future;
+} futureiterobject_dummy;
+
+static inline void
+_PyAsyncModule_ClearFreeLists(struct _Py_object_freelists *freelists, int is_finalization)
+{
+#ifdef WITH_FREELISTS
+    PyObject *next;
+    PyObject *current;
+
+    next = (PyObject*) freelists->futureiters.fi_freelist;
+    while (next != NULL) {
+        assert(freelists->futureiters.fi_freelist_len > 0);
+        freelists->futureiters.fi_freelist_len--;
+
+        current = next;
+        next = (PyObject*) ((futureiterobject_dummy*) current)->future;
+        PyObject_GC_Del(current);
+    }
+    assert(freelists->futureiters.fi_freelist_len == 0 || freelists->futureiters.fi_freelist_len == -1);
+    freelists->futureiters.fi_freelist = NULL;
+    if (is_finalization) {
+        freelists->futureiters.fi_freelist_len = -1;
+    }
+    else {
+        freelists->futureiters.fi_freelist_len = 0;
+    }
+#endif
+}
+
 #ifdef __cplusplus
 }
 #endif

Lib/test/libregrtest/utils.py

@@ -293,6 +293,13 @@ def clear_caches():
     else:
         importlib_metadata.FastPath.__new__.cache_clear()
 
+    try:
+        _asyncio = sys.modules['_asyncio']
+    except KeyError:
+        pass
+    else:
+        _asyncio._clear_freelist()
+
 
 def get_build_info():
     # Get most important configure and build options as a list of strings.

Modules/_asynciomodule.c

@@ -135,15 +135,9 @@ 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;
 
-    futureiterobject *fi_freelist;
-    Py_ssize_t fi_freelist_len;
-
     /* Linked-list of all tasks which are instances of asyncio.Task or subclasses
        of it. Third party tasks implementations which don't inherit from
        asyncio.Task are tracked separately using the 'non_asyncio_tasks' WeakSet.
@@ -220,6 +214,16 @@ get_asyncio_state_by_def(PyObject *self)
 #include "clinic/_asynciomodule.c.h"
 
 
+#ifdef WITH_FREELISTS
+static struct _Py_asyncmodule_futureiter_freelist *
+get_futureiter_freelist(void)
+{
+    struct _Py_object_freelists *freelists = _Py_object_freelists_GET();
+    assert(freelists != NULL);
+    return &freelists->futureiters;
+}
+#endif
+
 /*[clinic input]
 class _asyncio.Future "FutureObj *" "&Future_Type"
 [clinic start generated code]*/
@@ -329,11 +333,11 @@ get_running_loop(asyncio_state *state, PyObject **loop)
 
     PyThreadState *ts = _PyThreadState_GET();
     uint64_t ts_id = PyThreadState_GetID(ts);
-    if (state->cached_running_loop_tsid == ts_id &&
-        state->cached_running_loop != NULL)
+    if (ts->asyncio_cached_running_loop_tsid == ts_id &&
+        ts->asyncio_cached_running_loop != NULL)
     {
         // Fast path, check the cache.
-        rl = state->cached_running_loop;
+        rl = ts->asyncio_cached_running_loop;
     }
     else {
         PyObject *ts_dict = _PyThreadState_GetDict(ts);  // borrowed
@@ -352,10 +356,8 @@ get_running_loop(asyncio_state *state, PyObject **loop)
             }
         }
 
-        // 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;
+        ts->asyncio_cached_running_loop = rl;
+        ts->asyncio_cached_running_loop_tsid = ts_id;
     }
 
 
@@ -398,10 +400,8 @@ set_running_loop(asyncio_state *state, PyObject *loop)
     }
 
 
-    // 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);
+    tstate->asyncio_cached_running_loop = loop; // borrowed, kept alive by ts_dict
+    tstate->asyncio_cached_running_loop_tsid = PyThreadState_GetID(tstate);
 
     return 0;
 }
@@ -1668,25 +1668,19 @@ FutureIter_dealloc(futureiterobject *it)
 
     assert(_PyType_HasFeature(tp, Py_TPFLAGS_HEAPTYPE));
 
-    PyObject *module = ((PyHeapTypeObject*)tp)->ht_module;
-    asyncio_state *state = NULL;
-
     PyObject_GC_UnTrack(it);
     tp->tp_clear((PyObject *)it);
 
-    // GH-115874: We can't use PyType_GetModuleByDef here as the type might have
-    // already been cleared, which is also why we must check if ht_module != NULL.
-    if (module && _PyModule_GetDef(module) == &_asynciomodule) {
-        state = get_asyncio_state(module);
+#ifdef WITH_FREELISTS
+    struct _Py_asyncmodule_futureiter_freelist* freelist = get_futureiter_freelist();
+    if (freelist->fi_freelist_len < FI_FREELIST_MAXLEN) {
+        freelist->fi_freelist_len++;
+        it->future = (FutureObj*) freelist->fi_freelist;
+        freelist->fi_freelist = it;
     }
-
-    // TODO GH-121621: This should be moved to thread state as well.
-    if (state && state->fi_freelist_len < FI_FREELIST_MAXLEN) {
-        state->fi_freelist_len++;
-        it->future = (FutureObj*) state->fi_freelist;
-        state->fi_freelist = it;
-    }
-    else {
+    else
+#endif
+    {
         PyObject_GC_Del(it);
         Py_DECREF(tp);
     }
@@ -1890,14 +1884,18 @@ future_new_iter(PyObject *fut)
     asyncio_state *state = get_asyncio_state_by_def((PyObject *)fut);
     ENSURE_FUTURE_ALIVE(state, fut)
 
-    if (state->fi_freelist_len) {
-        state->fi_freelist_len--;
-        it = state->fi_freelist;
-        state->fi_freelist = (futureiterobject*) it->future;
+#ifdef WITH_FREELISTS
+    struct _Py_asyncmodule_futureiter_freelist* freelist = get_futureiter_freelist();
+    if (freelist->fi_freelist_len) {
+        freelist->fi_freelist_len--;
+        it = freelist->fi_freelist;
+        freelist->fi_freelist = (futureiterobject*) it->future;
         it->future = NULL;
         _Py_NewReference((PyObject*) it);
     }
-    else {
+    else
+#endif
+    {
         it = PyObject_GC_New(futureiterobject, state->FutureIterType);
         if (it == NULL) {
             return NULL;
@@ -3353,6 +3351,24 @@ task_wakeup(TaskObj *task, PyObject *o)
 /*********************** Functions **************************/
 
 
+/*[clinic input]
+_asyncio._clear_freelist
+
+Clears the asyncio freelist.
+
+Internal CPython implementation detail. Do not depend on this or use it!
+This function is thread-specific.
+
+[clinic start generated code]*/
+
+static PyObject *
+_asyncio__clear_freelist_impl(PyObject *module)
+/*[clinic end generated code: output=8d0e295bbbe2f8b6 input=f3ef7630d66cf63a]*/
+{
+    _PyAsyncModule_ClearFreeLists(_Py_object_freelists_GET(), 0);
+    Py_RETURN_NONE;
+}
+
 /*[clinic input]
 _asyncio._get_running_loop
 
@@ -3768,24 +3784,6 @@ _asyncio_all_tasks_impl(PyObject *module, PyObject *loop)
     return tasks;
 }
 
-static void
-module_free_freelists(asyncio_state *state)
-{
-    PyObject *next;
-    PyObject *current;
-
-    next = (PyObject*) state->fi_freelist;
-    while (next != NULL) {
-        assert(state->fi_freelist_len > 0);
-        state->fi_freelist_len--;
-
-        current = next;
-        next = (PyObject*) ((futureiterobject*) current)->future;
-        PyObject_GC_Del(current);
-    }
-    assert(state->fi_freelist_len == 0);
-    state->fi_freelist = NULL;
-}
 
 static int
 module_traverse(PyObject *mod, visitproc visit, void *arg)
@@ -3816,12 +3814,15 @@ module_traverse(PyObject *mod, visitproc visit, void *arg)
     Py_VISIT(state->context_kwname);
 
     // Visit freelist.
-    PyObject *next = (PyObject*) state->fi_freelist;
+#ifdef WITH_FREELISTS
+    struct _Py_asyncmodule_futureiter_freelist* freelist = get_futureiter_freelist();
+    PyObject *next = (PyObject*) freelist->fi_freelist;
     while (next != NULL) {
         PyObject *current = next;
         Py_VISIT(current);
         next = (PyObject*) ((futureiterobject*) current)->future;
     }
+#endif
     return 0;
 }
 
@@ -3853,7 +3854,8 @@ module_clear(PyObject *mod)
 
     Py_CLEAR(state->context_kwname);
 
-    module_free_freelists(state);
+    _PyAsyncModule_ClearFreeLists(_Py_object_freelists_GET(), 0);
+
 
     return 0;
 }
@@ -3965,6 +3967,7 @@ static PyMethodDef asyncio_methods[] = {
     _ASYNCIO__LEAVE_TASK_METHODDEF
     _ASYNCIO__SWAP_CURRENT_TASK_METHODDEF
     _ASYNCIO_ALL_TASKS_METHODDEF
+    _ASYNCIO__CLEAR_FREELIST_METHODDEF
     {NULL, NULL}
 };
 

Objects/object.c

@@ -815,6 +815,8 @@ _PyObject_ClearFreeLists(struct _Py_object_freelists *freelists, int is_finaliza
     _PyDict_ClearFreeList(freelists, is_finalization);
     _PyContext_ClearFreeList(freelists, is_finalization);
     _PyAsyncGen_ClearFreeLists(freelists, is_finalization);
+    _PyAsyncModule_ClearFreeLists(freelists, is_finalization);
+
     // Only be cleared if is_finalization is true.
     _PyObjectStackChunk_ClearFreeList(freelists, is_finalization);
     _PySlice_ClearFreeList(freelists, is_finalization);

Python/pystate.c

@@ -1700,6 +1700,8 @@ PyThreadState_Clear(PyThreadState *tstate)
 
     /* Don't clear tstate->pyframe: it is a borrowed reference */
 
+    /* Don't clear asyncio_cached_running_loop, it's borrowed. */
+
     Py_CLEAR(tstate->dict);
     Py_CLEAR(tstate->async_exc);