From 76cbbc636f9f60a6595aca37270cce30640433b1 Mon Sep 17 00:00:00 2001 From: Itamar Ostricher Date: Sun, 18 Sep 2022 14:44:30 -0700 Subject: [PATCH 1/6] gh-100344: Provide C implementation for asyncio.current_task MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implementing it in C makes it about 4x-6x faster Microbenchmark: ``` # bench.py import time import timeit import asyncio ITERS: int = 10**6 NANO: int = 10**9 NANO_PER_ITER: float = NANO / ITERS async def main(): # avoid load attr noise py_current_task = asyncio.tasks._py_current_task c_current_task = asyncio.tasks._c_current_task asyncio.current_task() # warmup py_current_task() # warmup c_current_task() # warmup print( "current_task: {}ns".format(timeit.timeit(py_current_task, number=ITERS, timer=time.process_time) * NANO_PER_ITER) ) print( "current_task: {}ns".format(timeit.timeit(c_current_task, number=ITERS, timer=time.process_time) * NANO_PER_ITER) ) asyncio.run(main()) ``` a few runs on MacBook Pro 2.4 GHz 8-Core Intel Core i9 64 GB 2667 MHz DDR4: debug build: ``` ~/work/pyexe/main-dbg ⌚ 9:57:34 $ ./python.exe bench.py [py] current_task: 606.234ns [c] current_task: 104.47699999999993ns ~/work/pyexe/main-dbg ⌚ 9:57:59 $ ./python.exe bench.py [py] current_task: 631.856ns [c] current_task: 110.22500000000002ns ~/work/pyexe/main-dbg ⌚ 9:58:08 $ ./python.exe bench.py [py] current_task: 637.746ns [c] current_task: 105.03899999999999ns ~/work/pyexe/main-dbg ⌚ 9:58:16 $ ./python.exe bench.py [py] current_task: 621.3169999999999ns [c] current_task: 103.01300000000002ns ``` opt build: ``` ~/work/pyexe/main-opt ⌚ 10:33:17 $ ./python.exe bench.py [py] current_task: 128.743ns [c] current_task: 31.997999999999998ns ~/work/pyexe/main-opt ⌚ 10:33:24 $ ./python.exe bench.py [py] current_task: 126.388ns [c] current_task: 32.64599999999998ns ~/work/pyexe/main-opt ⌚ 10:33:26 $ ./python.exe bench.py [py] current_task: 137.053ns [c] current_task: 32.066999999999986ns ~/work/pyexe/main-opt ⌚ 10:33:28 $ ./python.exe bench.py [py] current_task: 131.17800000000003ns [c] current_task: 32.06600000000001ns ``` --- Doc/whatsnew/3.12.rst | 3 + Lib/asyncio/tasks.py | 5 +- ...-12-19-12-18-28.gh-issue-100344.lfCqpE.rst | 2 + Modules/_asynciomodule.c | 39 ++++++++++++ Modules/clinic/_asynciomodule.c.h | 62 ++++++++++++++++++- 5 files changed, 109 insertions(+), 2 deletions(-) create mode 100644 Misc/NEWS.d/next/Library/2022-12-19-12-18-28.gh-issue-100344.lfCqpE.rst diff --git a/Doc/whatsnew/3.12.rst b/Doc/whatsnew/3.12.rst index 0cc4471364b671..d480be27542c3e 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 fa853283c0c5e4..e78719de216fd0 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 @@ def _unregister_task(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/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 00000000000000..d55f6888dbde63 --- /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 60369d89dc39c9..16da59a05a4a1f 100644 --- a/Modules/_asynciomodule.c +++ b/Modules/_asynciomodule.c @@ -3344,6 +3344,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; + } + + ret = PyDict_GetItemWithError(state->current_tasks, loop); + if (ret == NULL && PyErr_Occurred()) { + return NULL; + } + else if (ret == NULL) { + Py_RETURN_NONE; + } + else { + Py_INCREF(ret); + return ret; + } +} + + /*********************** PyRunningLoopHolder ********************/ @@ -3599,6 +3637,7 @@ module_init(asyncio_state *state) 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 f2fbb352c2c69b..43c5d771798634 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]*/ From 5f33beb5f53827d1c0984b875ae70e7ee84e9a4b Mon Sep 17 00:00:00 2001 From: Itamar Ostricher Date: Tue, 20 Dec 2022 09:43:04 +0200 Subject: [PATCH 2/6] Fix refleak in C impl of current_task --- Modules/_asynciomodule.c | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c index 16da59a05a4a1f..880de2cd9c39ef 100644 --- a/Modules/_asynciomodule.c +++ b/Modules/_asynciomodule.c @@ -3362,23 +3362,23 @@ _asyncio_current_task_impl(PyObject *module, PyObject *loop) if (loop == Py_None) { loop = _asyncio_get_running_loop_impl(module); - } - - if (loop == NULL) { - return NULL; + 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; } - else { - Py_INCREF(ret); - return ret; - } + Py_INCREF(ret); + return ret; } From 37072a1581d7a664bd71fd4a142bec66e84eb844 Mon Sep 17 00:00:00 2001 From: Itamar Ostricher Date: Tue, 20 Dec 2022 22:10:48 +0200 Subject: [PATCH 3/6] Optimize also the common case of the Python implementation of asyncio.current_task for ~40% speedup --- Lib/asyncio/tasks.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index e78719de216fd0..bdd323d0e94ffc 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -36,7 +36,10 @@ def current_task(loop=None): """Return a currently executed task.""" if loop is None: loop = events.get_running_loop() - return _current_tasks.get(loop) + try: + return _current_tasks[loop] + except: + return None def all_tasks(loop=None): From 8f0f80a77f9fe41511770f31353a481332220809 Mon Sep 17 00:00:00 2001 From: Itamar Ostricher Date: Wed, 21 Dec 2022 08:34:34 +0200 Subject: [PATCH 4/6] Catch more specific KeyError exception --- Lib/asyncio/tasks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index bdd323d0e94ffc..b1d5e2b9537a26 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -38,7 +38,7 @@ def current_task(loop=None): loop = events.get_running_loop() try: return _current_tasks[loop] - except: + except KeyError: return None From c96db2c3073f1366e51aa17d924ad9ad8b4a7264 Mon Sep 17 00:00:00 2001 From: Itamar Ostricher Date: Wed, 21 Dec 2022 16:19:42 +0200 Subject: [PATCH 5/6] Adjust CurrentLoopTests to cover both C and Python implementations of current_task --- Lib/test/test_asyncio/test_tasks.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py index 5168b8250ef0a2..e533d5273e9f38 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 @@ 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.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) From 391db59ca8cabcba9b8ba20ee0ea83e6613ccb7b Mon Sep 17 00:00:00 2001 From: Itamar Ostricher Date: Thu, 22 Dec 2022 15:13:02 +0200 Subject: [PATCH 6/6] Revert optimizations in Python implementation of asyncio.current_task --- Lib/asyncio/tasks.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index b1d5e2b9537a26..e78719de216fd0 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -36,10 +36,7 @@ def current_task(loop=None): """Return a currently executed task.""" if loop is None: loop = events.get_running_loop() - try: - return _current_tasks[loop] - except KeyError: - return None + return _current_tasks.get(loop) def all_tasks(loop=None):