From 9c3de43984938423a2059b3059e1cf3a588900eb Mon Sep 17 00:00:00 2001 From: Itamar Ostricher Date: Tue, 27 Dec 2022 07:51:48 -0800 Subject: [PATCH] Backport gh-100345 / Port D39555934: Provide C implementation for asyncio.current_task Summary: Implementing it in C makes it about 4x-6x faster Upstream issue: https://github.com/python/cpython/issues/100344 Upstream PR (3.12): https://github.com/python/cpython/pull/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 --- Lib/asyncio/tasks.py | 3 +++ Lib/test/test_asyncio/test_tasks.py | 21 ++++++++++------ Modules/_asynciomodule.c | 38 ++++++++++++++++++++++++++++ Modules/clinic/_asynciomodule.c.h | 39 ++++++++++++++++++++++++++++- 4 files changed, 93 insertions(+), 8 deletions(-) diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index 953dd3aa0ff..541532fd35f 100644 --- 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 index 2102c86e862..2b1b60ba97f 100644 --- 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 index 498902dcd96..ab84756f598 100644 --- 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 index af90d73ab03..feb60e680cf 100644 --- a/Modules/clinic/_asynciomodule.c.h +++ b/Modules/clinic/_asynciomodule.c.h @@ -1534,4 +1534,41 @@ _asyncio__AwaitingFuture_set_exception(_AwaitingFutureObj *self, PyObject *const exit: return return_value; } -/*[clinic end generated code: output=56abd79114e8e9b4 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)(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; +} +/*[clinic end generated code: output=8ed6665799b57e65 input=a9049054013a1b77]*/