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

This commit is contained in:
Kumar Aditya 2025-12-15 16:57:51 +05:30 committed by GitHub
parent 15c9d9027e
commit fb554ad68d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 36 additions and 33 deletions

View file

@ -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 */

View file

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

View file

@ -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)