Skip to content

Commit

Permalink
GH-120974: make _asyncio.all_tasks thread safe (#122801)
Browse files Browse the repository at this point in the history
Make `_asyncio.all_tasks` thread safe, also changes state lock to use critical section.
  • Loading branch information
kumaraditya303 authored Aug 11, 2024
1 parent 363374c commit 3aaed08
Showing 1 changed file with 15 additions and 7 deletions.
22 changes: 15 additions & 7 deletions Modules/_asynciomodule.c
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -1923,15 +1924,15 @@ 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);

task->next = state->asyncio_tasks.head;
state->asyncio_tasks.head->prev = task;
state->asyncio_tasks.head = task;
exit:
ASYNCIO_STATE_UNLOCK(state);
}

Expand All @@ -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) {
Expand All @@ -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);
}

Expand Down Expand Up @@ -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);
Expand All @@ -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);
Expand Down

0 comments on commit 3aaed08

Please sign in to comment.