From ce2491a8644c79ebb25c32864cd110ab4af47562 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Tue, 29 Mar 2022 20:43:25 +0300 Subject: [PATCH 1/6] Make changes --- Lib/asyncio/tasks.py | 6 +- Modules/_asynciomodule.c | 93 ++++++++++++++++++++++++------- Modules/clinic/_asynciomodule.c.h | 34 ++++++++++- 3 files changed, 111 insertions(+), 22 deletions(-) diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index 27fe58da15136a0..fb10a5f52da97d0 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -252,6 +252,10 @@ def uncancel(self): self._num_cancels_requested -= 1 return self._num_cancels_requested + def _check_future(self, fut): + """Return False if task and future loops are not compatible.""" + return futures._get_loop(fut) is self._loop + def __step(self, exc=None): if self.done(): raise exceptions.InvalidStateError( @@ -292,7 +296,7 @@ def __step(self, exc=None): blocking = getattr(result, '_asyncio_future_blocking', None) if blocking is not None: # Yielded Future must come from Future.__iter__(). - if futures._get_loop(result) is not self._loop: + if not self._check_future(result): new_exc = RuntimeError( f'Task {self!r} got Future ' f'{result!r} attached to a different loop') diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c index 5b8555e74b37517..513f73b120c150e 100644 --- a/Modules/_asynciomodule.c +++ b/Modules/_asynciomodule.c @@ -23,6 +23,7 @@ _Py_IDENTIFIER(call_soon); _Py_IDENTIFIER(cancel); _Py_IDENTIFIER(get_event_loop); _Py_IDENTIFIER(throw); +_Py_IDENTIFIER(_check_future); /* State of the _asyncio module */ @@ -1795,6 +1796,7 @@ class _asyncio.Task "TaskObj *" "&Task_Type" static int task_call_step_soon(TaskObj *, PyObject *); static PyObject * task_wakeup(TaskObj *, PyObject *); static PyObject * task_step(TaskObj *, PyObject *); +static int task_check_future(TaskObj *, PyObject *); /* ----- Task._step wrapper */ @@ -2269,7 +2271,6 @@ Returns the remaining number of cancellation requests. static PyObject * _asyncio_Task_uncancel_impl(TaskObj *self) /*[clinic end generated code: output=58184d236a817d3c input=68f81a4b90b46be2]*/ -/*[clinic end generated code]*/ { if (self->task_num_cancels_requested > 0) { self->task_num_cancels_requested -= 1; @@ -2277,6 +2278,25 @@ _asyncio_Task_uncancel_impl(TaskObj *self) return PyLong_FromLong(self->task_num_cancels_requested); } +/*[clinic input] +_asyncio.Task._check_future + + future: object + +Return False if task and future loops are not compatible. +[clinic start generated code]*/ + +static PyObject * +_asyncio_Task__check_future_impl(TaskObj *self, PyObject *future) +/*[clinic end generated code: output=2473f5b55d524bd5 input=76b4fa671a14b681]*/ +{ + int res = task_check_future(self, future); + if (res < 0) { + return NULL; + } + return PyBool_FromLong(res); +} + /*[clinic input] _asyncio.Task.get_stack @@ -2502,6 +2522,7 @@ static PyMethodDef TaskType_methods[] = { _ASYNCIO_TASK_CANCEL_METHODDEF _ASYNCIO_TASK_CANCELLING_METHODDEF _ASYNCIO_TASK_UNCANCEL_METHODDEF + _ASYNCIO_TASK__CHECK_FUTURE_METHODDEF _ASYNCIO_TASK_GET_STACK_METHODDEF _ASYNCIO_TASK_PRINT_STACK_METHODDEF _ASYNCIO_TASK__MAKE_CANCELLED_ERROR_METHODDEF @@ -2569,6 +2590,36 @@ TaskObj_dealloc(PyObject *self) Py_TYPE(task)->tp_free(task); } +static int +task_check_future(TaskObj *task, PyObject *future) +{ + if (Task_CheckExact(task)) { + int res; + if (Future_CheckExact(future) || Task_CheckExact(future)) { + FutureObj *fut = (FutureObj *)future; + res = (fut->fut_loop != task->task_loop); + } else { + PyObject *oloop = get_future_loop(future); + if (oloop == NULL) { + return -1; + } + res = (oloop != task->task_loop); + Py_DECREF(oloop); + } + return res; + } else { + PyObject * ret = _PyObject_CallMethodIdOneArg((PyObject *)task, + &PyId__check_future, + future); + if (ret == NULL) { + return -1; + } + int is_true = PyObject_IsTrue(ret); + Py_DECREF(ret); + return is_true; + } +} + static int task_call_step_soon(TaskObj *task, PyObject *arg) { @@ -2707,22 +2758,22 @@ task_step_impl(TaskObj *task, PyObject *exc) /* The error is StopIteration and that means that the underlying coroutine has resolved */ - PyObject *res; + PyObject *tmp; if (task->task_must_cancel) { // Task is cancelled right before coro stops. task->task_must_cancel = 0; - res = future_cancel((FutureObj*)task, task->task_cancel_msg); + tmp = future_cancel((FutureObj*)task, task->task_cancel_msg); } else { - res = future_set_result((FutureObj*)task, result); + tmp = future_set_result((FutureObj*)task, result); } Py_DECREF(result); - if (res == NULL) { + if (tmp == NULL) { return NULL; } - Py_DECREF(res); + Py_DECREF(tmp); Py_RETURN_NONE; } @@ -2786,11 +2837,15 @@ task_step_impl(TaskObj *task, PyObject *exc) /* Check if `result` is FutureObj or TaskObj (and not a subclass) */ if (Future_CheckExact(result) || Task_CheckExact(result)) { PyObject *wrapper; - PyObject *res; + PyObject *tmp; FutureObj *fut = (FutureObj*)result; /* Check if `result` future is attached to a different loop */ - if (fut->fut_loop != task->task_loop) { + res = task_check_future(task, result); + if (res == -1) { + goto fail; + } + if (res == 0) { goto different_loop; } @@ -2805,13 +2860,13 @@ task_step_impl(TaskObj *task, PyObject *exc) if (wrapper == NULL) { goto fail; } - res = future_add_done_callback( + tmp = future_add_done_callback( (FutureObj*)result, wrapper, task->task_context); Py_DECREF(wrapper); - if (res == NULL) { + if (tmp == NULL) { goto fail; } - Py_DECREF(res); + Py_DECREF(tmp); /* task._fut_waiter = result */ task->task_fut_waiter = result; /* no incref is necessary */ @@ -2853,7 +2908,7 @@ task_step_impl(TaskObj *task, PyObject *exc) if (o != NULL && o != Py_None) { /* `result` is a Future-compatible object */ PyObject *wrapper; - PyObject *res; + PyObject *tmp; int blocking = PyObject_IsTrue(o); Py_DECREF(o); @@ -2862,15 +2917,13 @@ task_step_impl(TaskObj *task, PyObject *exc) } /* Check if `result` future is attached to a different loop */ - PyObject *oloop = get_future_loop(result); - if (oloop == NULL) { + res = task_check_future(task, result); + if (res == -1) { goto fail; } - if (oloop != task->task_loop) { - Py_DECREF(oloop); + if (res == 0) { goto different_loop; } - Py_DECREF(oloop); if (!blocking) { goto yield_insteadof_yf; @@ -2897,13 +2950,13 @@ task_step_impl(TaskObj *task, PyObject *exc) PyObject *stack[2]; stack[0] = wrapper; stack[1] = (PyObject *)task->task_context; - res = PyObject_Vectorcall(add_cb, stack, 1, context_kwname); + tmp = PyObject_Vectorcall(add_cb, stack, 1, context_kwname); Py_DECREF(add_cb); Py_DECREF(wrapper); - if (res == NULL) { + if (tmp == NULL) { goto fail; } - Py_DECREF(res); + Py_DECREF(tmp); /* task._fut_waiter = result */ task->task_fut_waiter = result; /* no incref is necessary */ diff --git a/Modules/clinic/_asynciomodule.c.h b/Modules/clinic/_asynciomodule.c.h index 4b64367a3f6312e..7de733d5214d4a5 100644 --- a/Modules/clinic/_asynciomodule.c.h +++ b/Modules/clinic/_asynciomodule.c.h @@ -466,6 +466,38 @@ _asyncio_Task_uncancel(TaskObj *self, PyObject *Py_UNUSED(ignored)) return _asyncio_Task_uncancel_impl(self); } +PyDoc_STRVAR(_asyncio_Task__check_future__doc__, +"_check_future($self, /, future)\n" +"--\n" +"\n" +"Return False if task and future loops are not compatible."); + +#define _ASYNCIO_TASK__CHECK_FUTURE_METHODDEF \ + {"_check_future", (PyCFunction)(void(*)(void))_asyncio_Task__check_future, METH_FASTCALL|METH_KEYWORDS, _asyncio_Task__check_future__doc__}, + +static PyObject * +_asyncio_Task__check_future_impl(TaskObj *self, PyObject *future); + +static PyObject * +_asyncio_Task__check_future(TaskObj *self, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames) +{ + PyObject *return_value = NULL; + static const char * const _keywords[] = {"future", NULL}; + static _PyArg_Parser _parser = {NULL, _keywords, "_check_future", 0}; + PyObject *argsbuf[1]; + PyObject *future; + + args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, 1, 1, 0, argsbuf); + if (!args) { + goto exit; + } + future = args[0]; + return_value = _asyncio_Task__check_future_impl(self, future); + +exit: + return return_value; +} + PyDoc_STRVAR(_asyncio_Task_get_stack__doc__, "get_stack($self, /, *, limit=None)\n" "--\n" @@ -890,4 +922,4 @@ _asyncio__leave_task(PyObject *module, PyObject *const *args, Py_ssize_t nargs, exit: return return_value; } -/*[clinic end generated code: output=64b3836574e8a18c input=a9049054013a1b77]*/ +/*[clinic end generated code: output=2d8c53296ab165b5 input=a9049054013a1b77]*/ From 617423ed4ccf21f27408d545961670ad3bd919a3 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Wed, 30 Mar 2022 03:13:19 +0300 Subject: [PATCH 2/6] Fix tests --- Modules/_asynciomodule.c | 36 ++++++++++++++++++++++-------------- 1 file changed, 22 insertions(+), 14 deletions(-) diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c index 513f73b120c150e..e0048350276a45f 100644 --- a/Modules/_asynciomodule.c +++ b/Modules/_asynciomodule.c @@ -1797,6 +1797,7 @@ static int task_call_step_soon(TaskObj *, PyObject *); static PyObject * task_wakeup(TaskObj *, PyObject *); static PyObject * task_step(TaskObj *, PyObject *); static int task_check_future(TaskObj *, PyObject *); +static int task_check_future_exact(TaskObj *, PyObject *); /* ----- Task._step wrapper */ @@ -2290,7 +2291,7 @@ static PyObject * _asyncio_Task__check_future_impl(TaskObj *self, PyObject *future) /*[clinic end generated code: output=2473f5b55d524bd5 input=76b4fa671a14b681]*/ { - int res = task_check_future(self, future); + int res = task_check_future_exact(self, future); if (res < 0) { return NULL; } @@ -2590,23 +2591,30 @@ TaskObj_dealloc(PyObject *self) Py_TYPE(task)->tp_free(task); } +static int +task_check_future_exact(TaskObj *task, PyObject *future) +{ + int res; + if (Future_CheckExact(future) || Task_CheckExact(future)) { + FutureObj *fut = (FutureObj *)future; + res = (fut->fut_loop == task->task_loop); + } else { + PyObject *oloop = get_future_loop(future); + if (oloop == NULL) { + return -1; + } + res = (oloop == task->task_loop); + Py_DECREF(oloop); + } + return res; +} + + static int task_check_future(TaskObj *task, PyObject *future) { if (Task_CheckExact(task)) { - int res; - if (Future_CheckExact(future) || Task_CheckExact(future)) { - FutureObj *fut = (FutureObj *)future; - res = (fut->fut_loop != task->task_loop); - } else { - PyObject *oloop = get_future_loop(future); - if (oloop == NULL) { - return -1; - } - res = (oloop != task->task_loop); - Py_DECREF(oloop); - } - return res; + return task_check_future_exact(task, future); } else { PyObject * ret = _PyObject_CallMethodIdOneArg((PyObject *)task, &PyId__check_future, From 924e01176c55b456e0b01a7c5a61876fcd0ce88d Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Wed, 30 Mar 2022 17:10:00 +0300 Subject: [PATCH 3/6] Add tests --- Lib/asyncio/tasks.py | 4 ++-- Lib/test/test_asyncio/test_tasks.py | 10 +++++++++- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index fb10a5f52da97d0..3952b5f2a7743de 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -252,9 +252,9 @@ def uncancel(self): self._num_cancels_requested -= 1 return self._num_cancels_requested - def _check_future(self, fut): + def _check_future(self, future): """Return False if task and future loops are not compatible.""" - return futures._get_loop(fut) is self._loop + return futures._get_loop(future) is self._loop def __step(self, exc=None): if self.done(): diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py index 8df1957bbe9e7a3..80afb2735136201 100644 --- a/Lib/test/test_asyncio/test_tasks.py +++ b/Lib/test/test_asyncio/test_tasks.py @@ -2383,7 +2383,13 @@ def add_done_callback(self, *args, **kwargs): return super().add_done_callback(*args, **kwargs) class Task(CommonFuture, BaseTask): - pass + def __init__(self, *args, **kwargs): + self._check_future_called = 0 + super().__init__(*args, **kwargs) + + def _check_future(self, future): + self._check_future_called += 1 + return super()._check_future(future) class Future(CommonFuture, BaseFuture): pass @@ -2409,6 +2415,8 @@ async def func(): dict(fut.calls), {'add_done_callback': 1}) + self.assertEqual(1, task._check_future_called) + # Add patched Task & Future back to the test case cls.Task = Task cls.Future = Future From e3cc354422bf7c177a45b089d0cfe947de62713e Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Wed, 30 Mar 2022 18:35:59 +0300 Subject: [PATCH 4/6] Add NEWS --- .../NEWS.d/next/Library/2022-03-30-18-35-50.bpo-47167.nCNHsB.rst | 1 + 1 file changed, 1 insertion(+) create mode 100644 Misc/NEWS.d/next/Library/2022-03-30-18-35-50.bpo-47167.nCNHsB.rst diff --git a/Misc/NEWS.d/next/Library/2022-03-30-18-35-50.bpo-47167.nCNHsB.rst b/Misc/NEWS.d/next/Library/2022-03-30-18-35-50.bpo-47167.nCNHsB.rst new file mode 100644 index 000000000000000..a37dd25810f2eb0 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2022-03-30-18-35-50.bpo-47167.nCNHsB.rst @@ -0,0 +1 @@ +Allow overriding a future compliance check in :class:`asyncio.Task`. From 8760a6019fabd9eb60ef2dfa287ea0a24469cdda Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Thu, 31 Mar 2022 16:00:00 +0300 Subject: [PATCH 5/6] Apply notes --- Modules/_asynciomodule.c | 12 ++++-------- Modules/clinic/_asynciomodule.c.h | 11 ++++++++--- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c index e0048350276a45f..d8d3da91cdd8e0a 100644 --- a/Modules/_asynciomodule.c +++ b/Modules/_asynciomodule.c @@ -2280,22 +2280,18 @@ _asyncio_Task_uncancel_impl(TaskObj *self) } /*[clinic input] -_asyncio.Task._check_future +_asyncio.Task._check_future -> bool future: object Return False if task and future loops are not compatible. [clinic start generated code]*/ -static PyObject * +static int _asyncio_Task__check_future_impl(TaskObj *self, PyObject *future) -/*[clinic end generated code: output=2473f5b55d524bd5 input=76b4fa671a14b681]*/ +/*[clinic end generated code: output=a3bfba79295c8d57 input=3b1d6dfd6fe90aa5]*/ { - int res = task_check_future_exact(self, future); - if (res < 0) { - return NULL; - } - return PyBool_FromLong(res); + return task_check_future_exact(self, future); } /*[clinic input] diff --git a/Modules/clinic/_asynciomodule.c.h b/Modules/clinic/_asynciomodule.c.h index 7de733d5214d4a5..163b0f95691b312 100644 --- a/Modules/clinic/_asynciomodule.c.h +++ b/Modules/clinic/_asynciomodule.c.h @@ -475,7 +475,7 @@ PyDoc_STRVAR(_asyncio_Task__check_future__doc__, #define _ASYNCIO_TASK__CHECK_FUTURE_METHODDEF \ {"_check_future", (PyCFunction)(void(*)(void))_asyncio_Task__check_future, METH_FASTCALL|METH_KEYWORDS, _asyncio_Task__check_future__doc__}, -static PyObject * +static int _asyncio_Task__check_future_impl(TaskObj *self, PyObject *future); static PyObject * @@ -486,13 +486,18 @@ _asyncio_Task__check_future(TaskObj *self, PyObject *const *args, Py_ssize_t nar static _PyArg_Parser _parser = {NULL, _keywords, "_check_future", 0}; PyObject *argsbuf[1]; PyObject *future; + int _return_value; args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, 1, 1, 0, argsbuf); if (!args) { goto exit; } future = args[0]; - return_value = _asyncio_Task__check_future_impl(self, future); + _return_value = _asyncio_Task__check_future_impl(self, future); + if ((_return_value == -1) && PyErr_Occurred()) { + goto exit; + } + return_value = PyBool_FromLong((long)_return_value); exit: return return_value; @@ -922,4 +927,4 @@ _asyncio__leave_task(PyObject *module, PyObject *const *args, Py_ssize_t nargs, exit: return return_value; } -/*[clinic end generated code: output=2d8c53296ab165b5 input=a9049054013a1b77]*/ +/*[clinic end generated code: output=fdb7129263a8712e input=a9049054013a1b77]*/ From da248fc9c80170aa15684db8a47904633386e212 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Fri, 1 Apr 2022 01:05:22 +0300 Subject: [PATCH 6/6] Document the method --- Doc/library/asyncio-extending.rst | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/Doc/library/asyncio-extending.rst b/Doc/library/asyncio-extending.rst index 619723e61b5f9e5..215d215bb14fe75 100644 --- a/Doc/library/asyncio-extending.rst +++ b/Doc/library/asyncio-extending.rst @@ -48,16 +48,27 @@ For this purpose the following, *private* constructors are listed: .. method:: Future.__init__(*, loop=None) -Create a built-in future instance. + Create a built-in future instance. -*loop* is an optional event loop instance. + *loop* is an optional event loop instance. .. method:: Task.__init__(coro, *, loop=None, name=None, context=None) -Create a built-in task instance. + Create a built-in task instance. -*loop* is an optional event loop instance. The rest of arguments are described in -:meth:`loop.create_task` description. + *loop* is an optional event loop instance. The rest of arguments are described in + :meth:`loop.create_task` description. + + .. versionchanged:: 3.11 + + *context* argument is added. + +.. method:: Tasl._check_future(future) + + Return ``True`` if *future* is attached to the same loop as the task, ``False`` + otherwise. + + .. versionadded:: 3.11 Task lifetime support