From fb554ad68db415ab149d4596113b8fc74f99bcc6 Mon Sep 17 00:00:00 2001 From: Kumar Aditya Date: Mon, 15 Dec 2025 16:57:51 +0530 Subject: [PATCH] gh-140414: streamline thread state access in `asyncio` (#142742) --- Include/internal/pycore_context.h | 3 ++ Modules/_asynciomodule.c | 62 +++++++++++++++---------------- Python/context.c | 4 +- 3 files changed, 36 insertions(+), 33 deletions(-) diff --git a/Include/internal/pycore_context.h b/Include/internal/pycore_context.h index c77ef7910c09..a833f790a621 100644 --- a/Include/internal/pycore_context.h +++ b/Include/internal/pycore_context.h @@ -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 */ diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c index 0e6a1e93e04f..0b2a5d7093e1 100644 --- a/Modules/_asynciomodule.c +++ b/Modules/_asynciomodule.c @@ -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 */ @@ -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); } @@ -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; @@ -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; @@ -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; @@ -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) { @@ -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; @@ -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; } @@ -3452,7 +3445,9 @@ 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; } @@ -3460,12 +3455,12 @@ task_step(asyncio_state *state, TaskObj *task, PyObject *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; } @@ -3476,10 +3471,10 @@ 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; } @@ -3487,9 +3482,9 @@ task_eager_start(asyncio_state *state, TaskObj *task) // 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; } @@ -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; @@ -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; } @@ -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 @@ -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; } @@ -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; @@ -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; @@ -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); } diff --git a/Python/context.c b/Python/context.c index 2f978b1c0abc..620e78ab1f9e 100644 --- a/Python/context.c +++ b/Python/context.c @@ -190,7 +190,7 @@ context_switched(PyThreadState *ts) } -static int +int _PyContext_Enter(PyThreadState *ts, PyObject *octx) { ENSURE_Context(octx, -1) @@ -220,7 +220,7 @@ PyContext_Enter(PyObject *octx) } -static int +int _PyContext_Exit(PyThreadState *ts, PyObject *octx) { ENSURE_Context(octx, -1) -- 2.47.3