Skip to content

gh-140414: streamline thread state access in asyncio #142742

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kumaraditya303 merged 1 commit into from
Dec 15, 2025
+36 −33
Merged

gh-140414: streamline thread state access in asyncio #142742

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions Include/internal/pycore_context.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,5 +55,8 @@ struct _pycontexttokenobject {
// Export for '_testcapi' shared extension
PyAPI_FUNC(PyObject*) _PyContext_NewHamtForTests(void);

PyAPI_FUNC(int) _PyContext_Enter(PyThreadState *ts, PyObject *octx);
PyAPI_FUNC(int) _PyContext_Exit(PyThreadState *ts, PyObject *octx);


#endif /* !Py_INTERNAL_CONTEXT_H */
62 changes: 31 additions & 31 deletions Modules/_asynciomodule.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2076,8 +2076,8 @@ class _asyncio.Task "TaskObj *" "&Task_Type"

static int task_call_step_soon(asyncio_state *state, TaskObj *, PyObject *);
static PyObject *task_wakeup(PyObject *op, PyObject *arg);
static PyObject * task_step(asyncio_state *, TaskObj *, PyObject *);
static int task_eager_start(asyncio_state *state, TaskObj *task);
static PyObject *task_step(asyncio_state *, TaskObj *, PyObject *);
static int task_eager_start(_PyThreadStateImpl *ts, asyncio_state *state, TaskObj *task);

/* ----- Task._step wrapper */

Expand Down Expand Up @@ -2195,15 +2195,14 @@ static PyMethodDef TaskWakeupDef = {
/* ----- Task introspection helpers */

static void
register_task(TaskObj *task)
register_task(_PyThreadStateImpl *ts, TaskObj *task)
{
if (task->task_node.next != NULL) {
// already registered
assert(task->task_node.prev != NULL);
return;
}
_PyThreadStateImpl *tstate = (_PyThreadStateImpl *) _PyThreadState_GET();
struct llist_node *head = &tstate->asyncio_tasks_head;
struct llist_node *head = &ts->asyncio_tasks_head;
llist_insert_tail(head, &task->task_node);
}

Expand Down Expand Up @@ -2241,10 +2240,8 @@ unregister_task(TaskObj *task)
}

static int
enter_task(PyObject *loop, PyObject *task)
enter_task(_PyThreadStateImpl *ts, PyObject *loop, PyObject *task)
{
_PyThreadStateImpl *ts = (_PyThreadStateImpl *)_PyThreadState_GET();

if (ts->asyncio_running_loop != loop) {
PyErr_Format(PyExc_RuntimeError, "loop %R is not the running loop", loop);
return -1;
Expand All @@ -2264,10 +2261,8 @@ enter_task(PyObject *loop, PyObject *task)
}

static int
leave_task(PyObject *loop, PyObject *task)
leave_task(_PyThreadStateImpl *ts, PyObject *loop, PyObject *task)
{
_PyThreadStateImpl *ts = (_PyThreadStateImpl *)_PyThreadState_GET();

if (ts->asyncio_running_loop != loop) {
PyErr_Format(PyExc_RuntimeError, "loop %R is not the running loop", loop);
return -1;
Expand All @@ -2286,10 +2281,8 @@ leave_task(PyObject *loop, PyObject *task)
}

static PyObject *
swap_current_task(PyObject *loop, PyObject *task)
swap_current_task(_PyThreadStateImpl *ts, PyObject *loop, PyObject *task)
{
_PyThreadStateImpl *ts = (_PyThreadStateImpl *)_PyThreadState_GET();

if (ts->asyncio_running_loop != loop) {
PyErr_Format(PyExc_RuntimeError, "loop %R is not the running loop", loop);
return NULL;
Expand Down Expand Up @@ -2384,7 +2377,7 @@ _asyncio_Task___init___impl(TaskObj *self, PyObject *coro, PyObject *loop,
if (self->task_name == NULL) {
return -1;
}

_PyThreadStateImpl *ts = (_PyThreadStateImpl *)_PyThreadState_GET();
if (eager_start) {
PyObject *res = PyObject_CallMethodNoArgs(loop, &_Py_ID(is_running));
if (res == NULL) {
Expand All @@ -2393,7 +2386,7 @@ _asyncio_Task___init___impl(TaskObj *self, PyObject *coro, PyObject *loop,
int is_loop_running = Py_IsTrue(res);
Py_DECREF(res);
if (is_loop_running) {
if (task_eager_start(state, self)) {
if (task_eager_start(ts, state, self)) {
return -1;
}
return 0;
Expand All @@ -2408,7 +2401,7 @@ _asyncio_Task___init___impl(TaskObj *self, PyObject *coro, PyObject *loop,
// works correctly in non-owning threads.
_PyObject_SetMaybeWeakref((PyObject *)self);
#endif
register_task(self);
register_task(ts, self);
return 0;
}

Expand Down Expand Up @@ -3452,20 +3445,22 @@ task_step(asyncio_state *state, TaskObj *task, PyObject *exc)
{
PyObject *res;

if (enter_task(task->task_loop, (PyObject*)task) < 0) {
_PyThreadStateImpl *ts = (_PyThreadStateImpl *)_PyThreadState_GET();

if (enter_task(ts, task->task_loop, (PyObject*)task) < 0) {
return NULL;
}

res = task_step_impl(state, task, exc);

if (res == NULL) {
PyObject *exc = PyErr_GetRaisedException();
leave_task(task->task_loop, (PyObject*)task);
leave_task(ts, task->task_loop, (PyObject*)task);
_PyErr_ChainExceptions1(exc);
return NULL;
}
else {
if (leave_task(task->task_loop, (PyObject*)task) < 0) {
if (leave_task(ts, task->task_loop, (PyObject*)task) < 0) {
Py_DECREF(res);
return NULL;
}
Expand All @@ -3476,20 +3471,20 @@ task_step(asyncio_state *state, TaskObj *task, PyObject *exc)
}

static int
task_eager_start(asyncio_state *state, TaskObj *task)
task_eager_start(_PyThreadStateImpl *ts, asyncio_state *state, TaskObj *task)
{
assert(task != NULL);
PyObject *prevtask = swap_current_task(task->task_loop, (PyObject *)task);
PyObject *prevtask = swap_current_task(ts, task->task_loop, (PyObject *)task);
if (prevtask == NULL) {
return -1;
}
// register the task into the linked list of tasks
// if the task completes eagerly (without suspending) then it will unregister itself
// in future_schedule_callbacks when done, otherwise
// it will continue as a regular (non-eager) asyncio task
register_task(task);
register_task(ts, task);

if (PyContext_Enter(task->task_context) == -1) {
if (_PyContext_Enter(&ts->base, task->task_context) == -1) {
Py_DECREF(prevtask);
return -1;
}
Expand All @@ -3508,7 +3503,7 @@ task_eager_start(asyncio_state *state, TaskObj *task)
Py_DECREF(stepres);
}

PyObject *curtask = swap_current_task(task->task_loop, prevtask);
PyObject *curtask = swap_current_task(ts, task->task_loop, prevtask);
Py_DECREF(prevtask);
if (curtask == NULL) {
retval = -1;
Expand All @@ -3517,7 +3512,7 @@ task_eager_start(asyncio_state *state, TaskObj *task)
Py_DECREF(curtask);
}

if (PyContext_Exit(task->task_context) == -1) {
if (_PyContext_Exit(&ts->base, task->task_context) == -1) {
retval = -1;
}

Expand Down Expand Up @@ -3712,7 +3707,8 @@ _asyncio__register_task_impl(PyObject *module, PyObject *task)
if (Task_Check(state, task)) {
// task is an asyncio.Task instance or subclass, use efficient
// linked-list implementation.
register_task((TaskObj *)task);
_PyThreadStateImpl *ts = (_PyThreadStateImpl *)_PyThreadState_GET();
register_task(ts, (TaskObj *)task);
Py_RETURN_NONE;
}
// As task does not inherit from asyncio.Task, fallback to less efficient
Expand Down Expand Up @@ -3745,7 +3741,8 @@ _asyncio__register_eager_task_impl(PyObject *module, PyObject *task)
if (Task_Check(state, task)) {
// task is an asyncio.Task instance or subclass, use efficient
// linked-list implementation.
register_task((TaskObj *)task);
_PyThreadStateImpl *ts = (_PyThreadStateImpl *)_PyThreadState_GET();
register_task(ts, (TaskObj *)task);
Py_RETURN_NONE;
}

Expand Down Expand Up @@ -3832,7 +3829,8 @@ static PyObject *
_asyncio__enter_task_impl(PyObject *module, PyObject *loop, PyObject *task)
/*[clinic end generated code: output=a22611c858035b73 input=de1b06dca70d8737]*/
{
if (enter_task(loop, task) < 0) {
_PyThreadStateImpl *ts = (_PyThreadStateImpl *)_PyThreadState_GET();
if (enter_task(ts, loop, task) < 0) {
return NULL;
}
Py_RETURN_NONE;
Expand All @@ -3856,7 +3854,8 @@ static PyObject *
_asyncio__leave_task_impl(PyObject *module, PyObject *loop, PyObject *task)
/*[clinic end generated code: output=0ebf6db4b858fb41 input=51296a46313d1ad8]*/
{
if (leave_task(loop, task) < 0) {
_PyThreadStateImpl *ts = (_PyThreadStateImpl *)_PyThreadState_GET();
if (leave_task(ts, loop, task) < 0) {
return NULL;
}
Py_RETURN_NONE;
Expand All @@ -3880,7 +3879,8 @@ _asyncio__swap_current_task_impl(PyObject *module, PyObject *loop,
PyObject *task)
/*[clinic end generated code: output=9f88de958df74c7e input=c9c72208d3d38b6c]*/
{
return swap_current_task(loop, task);
_PyThreadStateImpl *ts = (_PyThreadStateImpl *)_PyThreadState_GET();
return swap_current_task(ts, loop, task);
}


Expand Down
4 changes: 2 additions & 2 deletions Python/context.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -190,7 +190,7 @@ context_switched(PyThreadState *ts)
}


static int
int
_PyContext_Enter(PyThreadState *ts, PyObject *octx)
{
ENSURE_Context(octx, -1)
Expand Down Expand Up @@ -220,7 +220,7 @@ PyContext_Enter(PyObject *octx)
}


static int
int
_PyContext_Exit(PyThreadState *ts, PyObject *octx)
{
ENSURE_Context(octx, -1)
Expand Down
Loading