Skip to content

Commit

Permalink
Backport gh-100345 / Port D39555934: Provide C implementation for asy…
Browse files Browse the repository at this point in the history
…ncio.current_task

Summary:
Implementing it in C makes it about 4x-6x faster

Upstream issue: python/cpython#100344
Upstream PR (3.12): python/cpython#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
  • Loading branch information
itamaro authored and Service User committed Dec 27, 2022
1 parent 0bd735a commit 9c3de43
Show file tree
Hide file tree
Showing 4 changed files with 93 additions and 8 deletions.
3 changes: 3 additions & 0 deletions Lib/asyncio/tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
21 changes: 14 additions & 7 deletions Lib/test/test_asyncio/test_tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -2972,6 +2972,7 @@ class CIntrospectionTests(test_utils.TestCase, BaseTaskIntrospectionTests):


class BaseCurrentLoopTests:
current_task = None

def setUp(self):
super().setUp()
Expand All @@ -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)
Expand Down
38 changes: 38 additions & 0 deletions Modules/_asynciomodule.c
Original file line number Diff line number Diff line change
Expand Up @@ -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 **************************/


Expand Down Expand Up @@ -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
Expand Down
39 changes: 38 additions & 1 deletion Modules/clinic/_asynciomodule.c.h

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit 9c3de43

Please sign in to comment.