gh-121621: Move asyncio running loop to thread state

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

@@ -77,10 +77,21 @@ typedef struct {
 
 #define FI_FREELIST_MAXLEN 255
 
+#ifdef Py_GIL_DISABLED
+#   define ASYNCIO_STATE_LOCK(state) PyMutex_Lock(&state->mutex)
+#   define ASYNCIO_STATE_UNLOCK(state) PyMutex_Unlock(&state->mutex)
+#else
+#   define ASYNCIO_STATE_LOCK(state) ((void)state)
+#   define ASYNCIO_STATE_UNLOCK(state) ((void)state)
+#endif
+
 typedef struct futureiterobject futureiterobject;
 
 /* State of the _asyncio module */
 typedef struct {
+#ifdef Py_GIL_DISABLED
+    PyMutex mutex;
+#endif
     PyTypeObject *FutureIterType;
     PyTypeObject *TaskStepMethWrapper_Type;
     PyTypeObject *FutureType;
@@ -130,9 +141,6 @@ typedef struct {
     /* 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.
@@ -209,6 +217,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]*/
@@ -341,6 +359,8 @@ get_running_loop(asyncio_state *state, PyObject **loop)
             }
         }
 
+        // TODO GH-121621: The should be moved to PyThreadState
+        // for easier and quicker access.
         state->cached_running_loop = rl;
         state->cached_running_loop_tsid = ts_id;
     }
@@ -384,6 +404,9 @@ set_running_loop(asyncio_state *state, PyObject *loop)
         return -1;
     }
 
+
+    // TODO GH-121621: The 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);
 
@@ -1653,26 +1676,19 @@ FutureIter_dealloc(futureiterobject *it)
     // FutureIter is a heap type so any subclass must also be a heap type.
     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.
-    // Due to this restriction, subclasses that belong to a different module
-    // will not be able to use the free list.
-    if (module && _PyModule_GetDef(module) == &_asynciomodule) {
-        state = get_asyncio_state(module);
-    }
-
-    if (state && state->fi_freelist_len < FI_FREELIST_MAXLEN) {
-        state->fi_freelist_len++;
-        it->future = (FutureObj*) state->fi_freelist;
-        state->fi_freelist = it;
+#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;
     }
-    else {
+    else
+#endif
+    {
         PyObject_GC_Del(it);
         Py_DECREF(tp);
     }
@@ -1876,14 +1892,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;
@@ -2018,10 +2038,12 @@ static  PyMethodDef TaskWakeupDef = {
 static void
 register_task(asyncio_state *state, TaskObj *task)
 {
+    ASYNCIO_STATE_LOCK(state);
     assert(Task_Check(state, task));
     assert(task != &state->asyncio_tasks.tail);
     if (task->next != NULL) {
         // already registered
+        ASYNCIO_STATE_UNLOCK(state);
         return;
     }
     assert(task->prev == NULL);
@@ -2030,6 +2052,7 @@ register_task(asyncio_state *state, TaskObj *task)
     task->next = state->asyncio_tasks.head;
     state->asyncio_tasks.head->prev = task;
     state->asyncio_tasks.head = task;
+    ASYNCIO_STATE_UNLOCK(state);
 }
 
 static int
@@ -2039,7 +2062,7 @@ register_eager_task(asyncio_state *state, PyObject *task)
 }
 
 static void
-unregister_task(asyncio_state *state, TaskObj *task)
+unregister_task_lock_held(asyncio_state *state, TaskObj *task)
 {
     assert(Task_Check(state, task));
     assert(task != &state->asyncio_tasks.tail);
@@ -2061,6 +2084,14 @@ unregister_task(asyncio_state *state, TaskObj *task)
     assert(state->asyncio_tasks.head != task);
 }
 
+static void
+unregister_task(asyncio_state *state, TaskObj *task)
+{
+    ASYNCIO_STATE_LOCK(state);
+    unregister_task_lock_held(state, task);
+    ASYNCIO_STATE_UNLOCK(state);
+}
+
 static int
 unregister_eager_task(asyncio_state *state, PyObject *task)
 {
@@ -2213,7 +2244,12 @@ _asyncio_Task___init___impl(TaskObj *self, PyObject *coro, PyObject *loop,
         // optimization: defer task name formatting
         // store the task counter as PyLong in the name
         // for deferred formatting in get_name
-        name = PyLong_FromUnsignedLongLong(++state->task_name_counter);
+#ifdef Py_GIL_DISABLED
+        unsigned long long counter = _Py_atomic_add_uint64(&state->task_name_counter, 1) + 1;
+#else
+        unsigned long long counter = ++state->task_name_counter;
+#endif
+        name = PyLong_FromUnsignedLongLong(counter);
     } else if (!PyUnicode_CheckExact(name)) {
         name = PyObject_Str(name);
     } else {
@@ -3328,6 +3364,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
 
@@ -3743,24 +3797,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)
@@ -3791,12 +3827,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;
 }
 
@@ -3828,7 +3867,8 @@ module_clear(PyObject *mod)
 
     Py_CLEAR(state->context_kwname);
 
-    module_free_freelists(state);
+    _PyAsyncModule_ClearFreeLists(_Py_object_freelists_GET(), 0);
+
 
     return 0;
 }
@@ -3940,6 +3980,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);