From ef8c620d36d0eca8ff74c540717cfcf73c36651a Mon Sep 17 00:00:00 2001
From: Pranav Thulasiram Bhat <pranavtbhat@fb.com>
Date: Mon, 19 Sep 2022 09:51:44 -0700
Subject: [PATCH] Optimize asyncio.current_task

Summary: Provide a native implementation of current_task as it is called frequently in the RequestContext scope guards.

Reviewed By: czardoz

Differential Revision: D39555934

fbshipit-source-id: f6dc3bf896e2509171835d333c6fa454cc6d2a3b
---
 Lib/asyncio/tasks.py              |  1 +
 Modules/_asynciomodule.c          | 35 +++++++++++++++++++++++++++
 Modules/clinic/_asynciomodule.c.h | 39 ++++++++++++++++++++++++++++++-
 3 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index 990661190c3..2248bda0927 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -975,6 +975,7 @@ def _unregister_task(task):
     from _asyncio import (_register_task, _unregister_task,
                           _enter_task, _leave_task,
                           _current_tasks,
+                          current_task,
                           all_tasks,
                           Task as CTask,
                           AsyncLazyValue as _ASYNC_LAZY_VALUE_TYPE,
diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c
index b0fee884a08..1788eb895f4 100644
--- a/Modules/_asynciomodule.c
+++ b/Modules/_asynciomodule.c
@@ -7838,6 +7838,40 @@ _asyncio__leave_task_impl(PyObject *module, PyObject *loop, PyObject *task)
     Py_RETURN_NONE;
 }
 
+/*[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 = get_event_loop();
+    }
+
+    if (loop == NULL) {
+        return NULL;
+    }
+
+    ret = PyDict_GetItemWithError(current_tasks, loop);
+    if (ret == NULL && PyErr_Occurred()) {
+        return NULL;
+    } else if (ret == NULL) {
+        Py_RETURN_NONE;
+    } else {
+        Py_INCREF(ret);
+        return ret;
+    }
+}
+
 /*[clinic input]
 _asyncio.all_tasks
 
@@ -8525,6 +8559,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_RUNNING_LOOP_METHODDEF
     _ASYNCIO__GET_RUNNING_LOOP_METHODDEF
diff --git a/Modules/clinic/_asynciomodule.c.h b/Modules/clinic/_asynciomodule.c.h
index 22148c97616..b854b8bcd2d 100644
--- a/Modules/clinic/_asynciomodule.c.h
+++ b/Modules/clinic/_asynciomodule.c.h
@@ -1129,6 +1129,43 @@ _asyncio__leave_task(PyObject *module, PyObject *const *args, Py_ssize_t nargs,
     return return_value;
 }
 
+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)(void(*)(void))_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;
+    static const char * const _keywords[] = {"loop", NULL};
+    static _PyArg_Parser _parser = {NULL, _keywords, "current_task", 0};
+    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;
+}
+
 PyDoc_STRVAR(_asyncio_all_tasks__doc__,
 "all_tasks($module, /, loop=None)\n"
 "--\n"
@@ -1409,4 +1446,4 @@ _asyncio__AwaitingFuture_set_exception(_AwaitingFutureObj *self, PyObject *const
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=b49fddf1bc7c5445 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=e47571d03f1229ab input=a9049054013a1b77]*/