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

Author: kumaraditya303

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

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

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)