code review

python · ambv · Feb 6, 2025 · Jan 11, 2025 · Jan 12, 2025 · Jan 14, 2025
commit 5b33c45ed57705722b3c2185ff48497d10fe486d
@@ -231,6 +231,8 @@ struct _is {
     // states gets added here and removed from the corresponding
     // thread state's list.
     struct llist_node asyncio_tasks_head;
+    PyMutex asyncio_tasks_lock;
+
     // Per-interpreter state for the obmalloc allocator.  For the main
     // interpreter and for all interpreters that don't have their
     // own obmalloc state, this points to the static structure in

@@ -3999,17 +3999,9 @@ add_one_task(asyncio_state *state, PyObject *tasks, PyObject *task, PyObject *lo
 }
 
 static inline int
-add_tasks_interp(PyInterpreterState *interp, PyListObject *tasks)
+add_tasks_llist(struct llist_node *head, PyListObject *tasks)
 {
-#ifdef Py_GIL_DISABLED
-    assert(interp->stoptheworld.world_stopped);
-#endif
-    // Start traversing from interpreter's linked list
-    struct llist_node *head = &interp->asyncio_tasks_head;
-    _PyThreadStateImpl *thead = (_PyThreadStateImpl *)interp->threads.head;
-
     struct llist_node *node;
-traverse:
     llist_for_each_safe(node, head) {
         TaskObj *task = llist_data(node, TaskObj, task_node);
         // The linked list holds borrowed references to task
@@ -4021,17 +4013,36 @@ add_tasks_interp(PyInterpreterState *interp, PyListObject *tasks)
         // otherwise it gets added to the list.
         if (_Py_TryIncref((PyObject *)task)) {
             if (_PyList_AppendTakeRef((PyListObject *)tasks, (PyObject *)task) < 0) {
-                // do not call any escaping calls here while holding the runtime lock.
+                // do not call any escaping calls here while the world is stopped.
                 return -1;
             }
         }
     }
-    // traverse the linked lists of thread states
-    if (thead != NULL) {
-        head = &thead->asyncio_tasks_head;
-        thead = (_PyThreadStateImpl *)thead->base.next;
-        goto traverse;
+    return 0;
+}
+
+static inline int
+add_tasks_interp(PyInterpreterState *interp, PyListObject *tasks)
+{
+#ifdef Py_GIL_DISABLED
+    assert(interp->stoptheworld.world_stopped);
+#endif
+    // Start traversing from interpreter's linked list
+    struct llist_node *head = &interp->asyncio_tasks_head;
+
+    if (add_tasks_llist(head, tasks) < 0) {
+        return -1;
+    }
+
+    // traverse the task lists of thread states
+    _Py_FOR_EACH_TSTATE_BEGIN(interp, p) {
+        _PyThreadStateImpl *ts = (_PyThreadStateImpl *)p;
+        head = &ts->asyncio_tasks_head;
+        if (add_tasks_llist(head, tasks) < 0) {
+            return -1;
+        }
     }
+
     return 0;
 }
 
@@ -4089,13 +4100,10 @@ _asyncio_all_tasks_impl(PyObject *module, PyObject *loop)
     // This design allows for lock free register/unregister of tasks
     // of loops running concurrently in different threads (general case).
     _PyEval_StopTheWorld(interp);
-    HEAD_LOCK(interp->runtime);
     int ret = add_tasks_interp(interp, (PyListObject *)tasks);
-    HEAD_UNLOCK(interp->runtime);
     _PyEval_StartTheWorld(interp);
     if (ret < 0) {
-        // call any escaping calls after releasing the runtime lock
-        // and starting the world to avoid any deadlocks.
+        // call any escaping calls after starting the world to avoid any deadlocks.
         Py_DECREF(tasks);
         Py_DECREF(loop);
         return NULL;

@@ -1702,12 +1702,12 @@ PyThreadState_Clear(PyThreadState *tstate)
     Py_CLEAR(((_PyThreadStateImpl *)tstate)->asyncio_running_task);
 
 
-    _PyEval_StopTheWorld(tstate->interp);
+    PyMutex_LockFlags(&tstate->interp->asyncio_tasks_lock, _Py_LOCK_DONT_DETACH);
     // merge any lingering tasks from thread state to interpreter's
     // tasks list
     llist_concat(&tstate->interp->asyncio_tasks_head,
                  &((_PyThreadStateImpl *)tstate)->asyncio_tasks_head);
-    _PyEval_StartTheWorld(tstate->interp);
+    PyMutex_Unlock(&tstate->interp->asyncio_tasks_lock);
 
     Py_CLEAR(tstate->dict);
     Py_CLEAR(tstate->async_exc);