Backport gh-100345 / Port D39555934: Provide C implementation for asyncio.current_task

itamaro · Service User · commit 9c3de4398493 · 2022-12-27T07:51:49.000-08:00
Summary: Implementing it in C makes it about 4x-6x faster Upstream issue: python/cpython#100344 Upstream PR (3.12): python/cpython#100345 Original summary from D39555934: Provide a native implementation of current_task as it is called frequently in the RequestContext scope guards. Test Plan: CI ``` ./python -m test test_asyncio.test_tasks -v -m "*current_task*" -R 3:3 ... ---------------------------------------------------------------------- Ran 33 tests in 0.045s OK (skipped=3) . == Tests result: SUCCESS == 1 test OK. Total duration: 2.4 sec Tests result: SUCCESS ``` Reviewers: alexanderm, pranavtbhat, #cinder Reviewed By: alexanderm Differential Revision: https://phabricator.intern.facebook.com/D42229719 Tasks: T132307964 Tags: commitClose, cinder-310-upstream-complete, accept2ship, cinder-310-standalone-upstream-candidate
diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
@@ -971,12 +971,14 @@ def _unregister_task(task):
 _py_enter_task = _enter_task
 _py_leave_task = _leave_task
 _py_all_tasks = all_tasks
+_py_current_task = current_task
 
 try:
     from _asyncio import (_register_task, _unregister_task,
                           _enter_task, _leave_task,
                           _current_tasks,
                           all_tasks,
+                          current_task,
                           AsyncLazyValue as _ASYNC_LAZY_VALUE_TYPE,
                           gather,
                           _is_coro_suspended,
@@ -989,3 +991,4 @@ def _unregister_task(task):
     _c_enter_task = _enter_task
     _c_leave_task = _leave_task
     _c_all_tasks = all_tasks
+    _c_current_task = current_task
diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
@@ -2972,6 +2972,7 @@ class CIntrospectionTests(test_utils.TestCase, BaseTaskIntrospectionTests):
 
 
 class BaseCurrentLoopTests:
+    current_task = None
 
     def setUp(self):
         super().setUp()
@@ -2982,33 +2983,39 @@ def new_task(self, coro):
         raise NotImplementedError
 
     def test_current_task_no_running_loop(self):
-        self.assertIsNone(asyncio.current_task(loop=self.loop))
+        self.assertIsNone(self.current_task(loop=self.loop))
 
     def test_current_task_no_running_loop_implicit(self):
         with self.assertRaises(RuntimeError):
-            asyncio.current_task()
+            self.current_task()
 
     def test_current_task_with_implicit_loop(self):
         async def coro():
-            self.assertIs(asyncio.current_task(loop=self.loop), task)
+            self.assertIs(self.current_task(loop=self.loop), task)
 
-            self.assertIs(asyncio.current_task(None), task)
-            self.assertIs(asyncio.current_task(), task)
+            self.assertIs(self.current_task(None), task)
+            self.assertIs(self.current_task(), task)
 
         task = self.new_task(coro())
         self.loop.run_until_complete(task)
-        self.assertIsNone(asyncio.current_task(loop=self.loop))
+        self.assertIsNone(self.current_task(loop=self.loop))
 
 
 class PyCurrentLoopTests(BaseCurrentLoopTests, test_utils.TestCase):
+    current_task = staticmethod(tasks._py_current_task)
 
     def new_task(self, coro):
         return tasks._PyTask(coro, loop=self.loop)
 
 
-@unittest.skipUnless(hasattr(tasks, '_CTask'),
+@unittest.skipUnless(hasattr(tasks, '_CTask') and
+                     hasattr(tasks, '_c_current_task'),
                      'requires the C _asyncio module')
 class CCurrentLoopTests(BaseCurrentLoopTests, test_utils.TestCase):
+    if hasattr(tasks, '_c_current_task'):
+        current_task = staticmethod(tasks._c_current_task)
+    else:
+        current_task = None
 
     def new_task(self, coro):
         return getattr(tasks, '_CTask')(coro, loop=self.loop)
diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c
@@ -8344,6 +8344,43 @@ static PyTypeObject _AwaitingFuture_Type = {
 };
 
 
+/*[clinic input]
+_asyncio.current_task
+
+    loop: object = None
+
+Return a currently executed task.
+
+[clinic start generated code]*/
+
+static PyObject *
+_asyncio_current_task_impl(PyObject *module, PyObject *loop)
+/*[clinic end generated code: output=fe15ac331a7f981a input=58910f61a5627112]*/
+{
+    PyObject *ret;
+
+    if (loop == Py_None) {
+        loop = _asyncio_get_running_loop_impl(module);
+        if (loop == NULL) {
+            return NULL;
+        }
+    } else {
+        Py_INCREF(loop);
+    }
+
+    ret = PyDict_GetItemWithError(current_tasks, loop);
+    Py_DECREF(loop);
+    if (ret == NULL && PyErr_Occurred()) {
+        return NULL;
+    }
+    else if (ret == NULL) {
+        Py_RETURN_NONE;
+    }
+    Py_INCREF(ret);
+    return ret;
+}
+
+
 /*********************** Module **************************/
 
 
@@ -8534,6 +8571,7 @@ module_init(void)
 PyDoc_STRVAR(module_doc, "Accelerator module for asyncio");
 
 static PyMethodDef asyncio_methods[] = {
+    _ASYNCIO_CURRENT_TASK_METHODDEF
     _ASYNCIO_GET_EVENT_LOOP_METHODDEF
     _ASYNCIO__GET_EVENT_LOOP_METHODDEF
     _ASYNCIO_GET_RUNNING_LOOP_METHODDEF
diff --git a/Modules/clinic/_asynciomodule.c.h b/Modules/clinic/_asynciomodule.c.h
