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

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

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