gh-100344: Add C implementation for `asyncio.current_task` (#100345)

Co-authored-by: pranavtbhat

6 files changed, 123 insertions, 9 deletions

diff --git a/Doc/whatsnew/3.12.rst b/Doc/whatsnew/3.12.rst
index 0cc4471..d480be2 100644
--- a/Doc/whatsnew/3.12.rst
+++ b/Doc/whatsnew/3.12.rst
@@ -225,6 +225,9 @@ asyncio
   a custom event loop factory.
   (Contributed by Kumar Aditya in :gh:`99388`.)
 
+* Add C implementation of :func:`asyncio.current_task` for 4x-6x speedup.
+  (Contributed by Itamar Ostricher and Pranav Thulasiram Bhat in :gh:`100344`.)
+
 inspect
 -------
 
diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index fa85328..e78719d 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -964,6 +964,7 @@ def _unregister_task(task):
     _all_tasks.discard(task)
 
 
+_py_current_task = current_task
 _py_register_task = _register_task
 _py_unregister_task = _unregister_task
 _py_enter_task = _enter_task
@@ -973,10 +974,12 @@ _py_leave_task = _leave_task
 try:
     from _asyncio import (_register_task, _unregister_task,
                           _enter_task, _leave_task,
-                          _all_tasks, _current_tasks)
+                          _all_tasks, _current_tasks,
+                          current_task)
 except ImportError:
     pass
 else:
+    _c_current_task = current_task
     _c_register_task = _register_task
     _c_unregister_task = _unregister_task
     _c_enter_task = _enter_task
diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
index 5168b82..e533d52 100644
--- a/Lib/test/test_asyncio/test_tasks.py
+++ b/Lib/test/test_asyncio/test_tasks.py
@@ -2804,6 +2804,7 @@ class CIntrospectionTests(test_utils.TestCase, BaseTaskIntrospectionTests):
 
 
 class BaseCurrentLoopTests:
+    current_task = None
 
     def setUp(self):
         super().setUp()
@@ -2814,33 +2815,39 @@ class BaseCurrentLoopTests:
         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.assertRaisesRegex(RuntimeError, 'no running event loop'):
-            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/Misc/NEWS.d/next/Library/2022-12-19-12-18-28.gh-issue-100344.lfCqpE.rst b/Misc/NEWS.d/next/Library/2022-12-19-12-18-28.gh-issue-100344.lfCqpE.rst
new file mode 100644
index 0000000..d55f688
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2022-12-19-12-18-28.gh-issue-100344.lfCqpE.rst
@@ -0,0 +1,2 @@
+Provide C implementation for :func:`asyncio.current_task` for a 4x-6x
+speedup.
diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c
index 32be537..6fe4ca4 100644
--- a/Modules/_asynciomodule.c
+++ b/Modules/_asynciomodule.c
@@ -3314,6 +3314,44 @@ _asyncio__leave_task_impl(PyObject *module, PyObject *loop, PyObject *task)
 }
 
 
+/*[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;
+    asyncio_state *state = get_asyncio_state(module);
+
+    if (loop == Py_None) {
+        loop = _asyncio_get_running_loop_impl(module);
+        if (loop == NULL) {
+            return NULL;
+        }
+    } else {
+        Py_INCREF(loop);
+    }
+
+    ret = PyDict_GetItemWithError(state->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 **************************/
 
 
@@ -3494,6 +3532,7 @@ fail:
 PyDoc_STRVAR(module_doc, "Accelerator module for asyncio");
 
 static PyMethodDef asyncio_methods[] = {
+    _ASYNCIO_CURRENT_TASK_METHODDEF
     _ASYNCIO_GET_EVENT_LOOP_METHODDEF
     _ASYNCIO_GET_RUNNING_LOOP_METHODDEF
     _ASYNCIO__GET_RUNNING_LOOP_METHODDEF
diff --git a/Modules/clinic/_asynciomodule.c.h b/Modules/clinic/_asynciomodule.c.h
index f2fbb35..43c5d77 100644
--- a/Modules/clinic/_asynciomodule.c.h
+++ b/Modules/clinic/_asynciomodule.c.h
@@ -1242,4 +1242,64 @@ _asyncio__leave_task(PyObject *module, PyObject *const *args, Py_ssize_t nargs,
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=83580c190031241c input=a9049054013a1b77]*/
+
+PyDoc_STRVAR(_asyncio_current_task__doc__,
+"current_task($module, /, loop=None)\n"
+"--\n"
+"\n"
+"Return a currently executed task.");
+
+#define _ASYNCIO_CURRENT_TASK_METHODDEF    \
+    {"current_task", _PyCFunction_CAST(_asyncio_current_task), METH_FASTCALL|METH_KEYWORDS, _asyncio_current_task__doc__},
+
+static PyObject *
+_asyncio_current_task_impl(PyObject *module, PyObject *loop);
+
+static PyObject *
+_asyncio_current_task(PyObject *module, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
+{
+    PyObject *return_value = NULL;
+    #if defined(Py_BUILD_CORE) && !defined(Py_BUILD_CORE_MODULE)
+
+    #define NUM_KEYWORDS 1
+    static struct {
+        PyGC_Head _this_is_not_used;
+        PyObject_VAR_HEAD
+        PyObject *ob_item[NUM_KEYWORDS];
+    } _kwtuple = {
+        .ob_base = PyVarObject_HEAD_INIT(&PyTuple_Type, NUM_KEYWORDS)
+        .ob_item = { &_Py_ID(loop), },
+    };
+    #undef NUM_KEYWORDS
+    #define KWTUPLE (&_kwtuple.ob_base.ob_base)
+
+    #else  // !Py_BUILD_CORE
+    #  define KWTUPLE NULL
+    #endif  // !Py_BUILD_CORE
+
+    static const char * const _keywords[] = {"loop", NULL};
+    static _PyArg_Parser _parser = {
+        .keywords = _keywords,
+        .fname = "current_task",
+        .kwtuple = KWTUPLE,
+    };
+    #undef KWTUPLE
+    PyObject *argsbuf[1];
+    Py_ssize_t noptargs = nargs + (kwnames ? PyTuple_GET_SIZE(kwnames) : 0) - 0;
+    PyObject *loop = Py_None;
+
+    args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, 0, 1, 0, argsbuf);
+    if (!args) {
+        goto exit;
+    }
+    if (!noptargs) {
+        goto skip_optional_pos;
+    }
+    loop = args[0];
+skip_optional_pos:
+    return_value = _asyncio_current_task_impl(module, loop);
+
+exit:
+    return return_value;
+}
+/*[clinic end generated code: output=00f494214f2fd008 input=a9049054013a1b77]*/