From 3aaed083a3f5eb7e490495c460b3dc1ce7451ce8 Mon Sep 17 00:00:00 2001
From: Kumar Aditya <kumaraditya@python.org>
Date: Sun, 11 Aug 2024 20:49:31 +0530
Subject: [PATCH] GH-120974: make `_asyncio.all_tasks` thread safe (#122801)

Make `_asyncio.all_tasks` thread safe, also changes state lock to use critical section.
---
 Modules/_asynciomodule.c | 22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c
index c6eb43f044f..0a769c46b87 100644
--- a/Modules/_asynciomodule.c
+++ b/Modules/_asynciomodule.c
@@ -3,6 +3,7 @@
 #endif
 
 #include "Python.h"
+#include "pycore_critical_section.h"  // Py_BEGIN_CRITICAL_SECTION_MUT()
 #include "pycore_dict.h"          // _PyDict_GetItem_KnownHash()
 #include "pycore_freelist.h"      // _Py_FREELIST_POP()
 #include "pycore_modsupport.h"    // _PyArg_CheckPositional()
@@ -77,8 +78,8 @@ typedef struct {
 #define Task_Check(state, obj) PyObject_TypeCheck(obj, state->TaskType)
 
 #ifdef Py_GIL_DISABLED
-#   define ASYNCIO_STATE_LOCK(state) PyMutex_Lock(&state->mutex)
-#   define ASYNCIO_STATE_UNLOCK(state) PyMutex_Unlock(&state->mutex)
+#   define ASYNCIO_STATE_LOCK(state) Py_BEGIN_CRITICAL_SECTION_MUT(&state->mutex)
+#   define ASYNCIO_STATE_UNLOCK(state) Py_END_CRITICAL_SECTION()
 #else
 #   define ASYNCIO_STATE_LOCK(state) ((void)state)
 #   define ASYNCIO_STATE_UNLOCK(state) ((void)state)
@@ -1923,8 +1924,7 @@ register_task(asyncio_state *state, TaskObj *task)
     assert(task != &state->asyncio_tasks.tail);
     if (task->next != NULL) {
         // already registered
-        ASYNCIO_STATE_UNLOCK(state);
-        return;
+        goto exit;
     }
     assert(task->prev == NULL);
     assert(state->asyncio_tasks.head != NULL);
@@ -1932,6 +1932,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;
+exit:
     ASYNCIO_STATE_UNLOCK(state);
 }
 
@@ -1951,8 +1952,7 @@ unregister_task(asyncio_state *state, TaskObj *task)
         // not registered
         assert(task->prev == NULL);
         assert(state->asyncio_tasks.head != task);
-        ASYNCIO_STATE_UNLOCK(state);
-        return;
+        goto exit;
     }
     task->next->prev = task->prev;
     if (task->prev == NULL) {
@@ -1964,6 +1964,7 @@ unregister_task(asyncio_state *state, TaskObj *task)
     task->next = NULL;
     task->prev = NULL;
     assert(state->asyncio_tasks.head != task);
+exit:
     ASYNCIO_STATE_UNLOCK(state);
 }
 
@@ -3628,6 +3629,8 @@ _asyncio_all_tasks_impl(PyObject *module, PyObject *loop)
         Py_DECREF(item);
     }
     Py_DECREF(eager_iter);
+    int err = 0;
+    ASYNCIO_STATE_LOCK(state);
     TaskObj *head = state->asyncio_tasks.head;
     Py_INCREF(head);
     assert(head != NULL);
@@ -3639,11 +3642,16 @@ _asyncio_all_tasks_impl(PyObject *module, PyObject *loop)
             Py_DECREF(tasks);
             Py_DECREF(loop);
             Py_DECREF(head);
-            return NULL;
+            err = 1;
+            break;
         }
         Py_INCREF(head->next);
         Py_SETREF(head, head->next);
     }
+    ASYNCIO_STATE_UNLOCK(state);
+    if (err) {
+        return NULL;
+    }
     PyObject *scheduled_iter = PyObject_GetIter(state->non_asyncio_tasks);
     if (scheduled_iter == NULL) {
         Py_DECREF(tasks);