gh-128002: optimistically remove tasks from linked list when finished (#129995)

kumaraditya303 · web-flow · commit c7a9d06e063c · 2025-02-13T13:19:53.000+05:30
diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c
@@ -413,12 +413,22 @@ future_ensure_alive(FutureObj *fut)
         }                                                           \
     } while(0);
 
+static void unregister_task(asyncio_state *state, TaskObj *task);
 
 static int
 future_schedule_callbacks(asyncio_state *state, FutureObj *fut)
 {
     _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(fut);
 
+    assert(fut->fut_state != STATE_PENDING);
+
+    if (Task_Check(state, fut)) {
+        // remove task from linked-list of tasks
+        // as it is finished now
+        TaskObj *task = (TaskObj *)fut;
+        unregister_task(state, task);
+    }
+
     if (fut->fut_callback0 != NULL) {
         /* There's a 1st callback */
 
@@ -4030,6 +4040,7 @@ add_tasks_llist(struct llist_node *head, PyListObject *tasks)
     struct llist_node *node;
     llist_for_each_safe(node, head) {
         TaskObj *task = llist_data(node, TaskObj, task_node);
+        assert(task->task_state == STATE_PENDING);
         // The linked list holds borrowed references to task
         // as such it is possible that the task is concurrently
         // deallocated while added to this list.
