8000 Eager task factory implementation · python/cpython@26536f2 · GitHub
[go: up one dir, main page]
Skip to content

Commit 26536f2

Browse files
itamaroitamaro
committed
Eager task factory implementation
1 parent 5c75b7a commit 26536f2

8 files changed

+145
-66
lines changed

‎Include/internal/pycore_global_objects_fini_generated.h

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Include/internal/pycore_global_strings.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -339,6 +339,7 @@ struct _Py_global_strings {
339339
STRUCT_FOR_ID(copy)
340340
STRUCT_FOR_ID(copyreg)
341341
STRUCT_FOR_ID(coro)
342+
STRUCT_FOR_ID(coro_result)
342343
STRUCT_FOR_ID(count)
343344
STRUCT_FOR_ID(cwd)
344345
STRUCT_FOR_ID(d)

‎Include/internal/pycore_runtime_init_generated.h

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Include/internal/pycore_unicodeobject_generated.h

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Lib/asyncio/taskgroups.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -163,7 +163,8 @@ def create_task(self, coro, *, name=None, context=None):
163163
task = self._loop.create_task(coro)
164164
else:
165165
task = self._loop.create_task(coro, context=context)
166-
tasks._set_task_name(task, name)
166+
if name is not None and not task.done(): # If it's done already, it's a future
167+
tasks._set_task_name(task, name)
167168
task.add_done_callback(self._on_task_done)
168169
self._tasks.add(task)
169170
return task

‎Lib/asyncio/tasks.py

Lines changed: 82 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
'wait', 'wait_for', 'as_completed', 'sleep',
77
'gather', 'shield', 'ensure_future', 'run_coroutine_threadsafe',
88
'current_task', 'all_tasks',
9+
'create_eager_task_factory', 'eager_task_factory',
910
'_register_task', '_unregister_task', '_enter_task', '_leave_task',
1011
)
1112

@@ -75,6 +76,8 @@ def _set_task_name(task, name):
7576
set_name(name)
7677

7778

79+
_NOT_SET = object()
80+
7881
class Task(futures._PyFuture): # Inherit Python Task implementation
7982
# from a Python Future implementation.
8083

@@ -93,7 +96,8 @@ class Task(futures._PyFuture): # Inherit Python Task implementation
9396
# status is still pending
9497
_log_destroy_pending = True
9598

96-
def __init__(self, coro, *, loop=None, name=None, context=None):
99+
def __init__(self, coro, *, loop=None, name=None, context=None,
100+
coro_result=_NOT_SET):
97101
super().__init__(loop=loop)
98102
if self._source_traceback:
99103
del self._source_traceback[-1]
@@ -117,7 +121,10 @@ def __init__(self, coro, *, loop=None, name=None, context=None):
117121
else:
118122
self._context = context
119123

120-
self._loop.call_soon(self.__step, context=self._context)
124+
if coro_result is _NOT_SET:
125+
self._loop.call_soon(self.__step, context=self._context)
126+
else:
127+
self.__step_handle_result(coro_result)
121128
_register_task(self)
122129

123130
def __del__(self):
@@ -287,55 +294,58 @@ def __step(self, exc=None):
287294
except BaseException as exc:
288295
super().set_exception(exc)
289296
else:
290-
blocking = getattr(result, '_asyncio_future_blocking', None)
291-
if blocking is not None:
297+
self.__step_handle_result(result)
298+
finally:
299+
_leave_task(self._loop, self)
300+
self = None # Needed to break cycles when an exception occurs.
301+
302+
def __step_handle_result(self, result):
303+
blocking = getattr(result, '_asyncio_future_blocking', None)
304+
if blocking is not None:
292305
# Yielded Future must come from Future.__iter__().
293-
if futures._get_loop(result) is not self._loop:
306+
if futures._get_loop(result) is not self._loop:
307+
new_exc = RuntimeError(
308+
f'Task {self!r} got Future '
309+
f'{result!r} attached to a different loop')
310+
self._loop.call_soon(
311+
self.__step, new_exc, context=self._context)
312+
elif blocking:
313+
if result is self:
294314
new_exc = RuntimeError(
295-
f'Task {self!r} got Future '
296-
f'{result!r} attached to a different loop')
315+
f'Task cannot await on itself: {self!r}')
297316
self._loop.call_soon(
298317
self.__step, new_exc, context=self._context)
299-
elif blocking:
300-
if result is self:
301-
new_exc = RuntimeError(
302-
f'Task cannot await on itself: {self!r}')
303-
self._loop.call_soon(
304-
self.__step, new_exc, context=self._context)
305-
else:
306-
result._asyncio_future_blocking = False
307-
result.add_done_callback(
308-
self.__wakeup, context=self._context)
309-
self._fut_waiter = result
310-
if self._must_cancel:
311-
if self._fut_waiter.cancel(
312-
msg=self._cancel_message):
313-
self._must_cancel = False
314318
else:
315-
new_exc = RuntimeError(
316-
f'yield was used instead of yield from '
317-
f'in task {self!r} with {result!r}')
318-
self._loop.call_soon(
319-
self.__step, new_exc, context=self._context)
320-
321-
elif result is None:
322-
# Bare yield relinquishes control for one event loop iteration.
323-
self._loop.call_soon(self.__step, context=self._context)
324-
elif inspect.isgenerator(result):
325-
# Yielding a generator is just wrong.
326-
new_exc = RuntimeError(
327-
f'yield was used instead of yield from for '
328-
f'generator in task {self!r} with {result!r}')
329-
self._loop.call_soon(
330-
self.__step, new_exc, context=self._context)
319+
result._asyncio_future_blocking = False
320+
result.add_done_callback(
321+
self.__wakeup, context=self._context)
322+
self._fut_waiter = result
323+
if self._must_cancel:
324+
if self._fut_waiter.cancel(
325+
msg=self._cancel_message):
326+
self._must_cancel = False
331327
else:
332-
# Yielding something else is an error.
333-
new_exc = RuntimeError(f'Task got bad yield: {result!r}')
328+
new_exc = RuntimeError(
329+
f'yield was used instead of yield from '
330+
f'in task {self!r} with {result!r}')
334331
self._loop.call_soon(
335332
self.__step, new_exc, context=self._context)
336-
finally:
337-
_leave_task(self._loop, self)
338-
self = None # Needed to break cycles when an exception occurs.
333+
334+
elif result is None:
335+
# Bare yield relinquishes control for one event loop iteration.
336+
self._loop.call_soon(self.__step, context=self._context)
337+
elif inspect.isgenerator(result):
338+
# Yielding a generator is just wrong.
339+
new_exc = RuntimeError(
340+
f'yield was used instead of yield from for '
341+
f'generator in task {self!r} with {result!r}')
342+
self._loop.call_soon(
343+
self.__step, new_exc, context=self._context)
344+
else:
345+
# Yielding something else is an error.
346+
new_exc = RuntimeError(f'Task got bad yield: {result!r}')
347+
self._loop.call_soon(
348+
self.__step, new_exc, context=self._context)
339349

340350
def __wakeup(self, future):
341351
try:
@@ -897,6 +907,35 @@ def callback():
897907
return future
898908

899909

910+
def create_eager_task_factory(custom_task_constructor):
911+
912+
def factory(loop, coro, *, name=None, context=None):
913+
loop._check_closed()
914+
if not loop.is_running():
915+
return custom_task_constructor(coro, loop=loop, name=name, context=context)
916+
917+
try:
918+
result = coro.send(None)
919+
except StopIteration as si:
920+
fut = loop.create_future()
921+
fut.set_result(si.value)
922+
return fut
923+
except Exception as ex:
924+
fut = loop.create_future()
925+
fut.set_exception(ex)
926+
return fut
927+
else:
928+
task = custom_task_constructor(
929+
coro, loop=loop, name=name, context=context, coro_result=result)
930+
if task._source_traceback:
931+
del task._source_traceback[-1]
932+
return task
933+
934+
return factory
935+
936+
eager_task_factory = create_eager_task_factory(Task)
937+
938+
900939
# WeakSet containing all alive tasks.
901940
_all_tasks = weakref.WeakSet()
902941

‎Modules/_asynciomodule.c

Lines changed: 37 additions & 13 deletions
< 325D td data-grid-cell-id="diff-6bd9e39980b88a721d902bcd915bbb3f24762f7f253430c45e52c42a2c5afd01-2832-2860-1" data-selected="false" role="gridcell" style="background-color:var(--bgColor-default);text-align:center" tabindex="-1" valign="top" class="focusable-grid-cell diff-line-number position-relative diff-line-number-neutral left-side">2860
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,9 @@ class _asyncio.Future "FutureObj *" "&Future_Type"
156156
/* Get FutureIter from Future */
157157
static PyObject * future_new_iter(PyObject *);
158158

159+
static PyObject *
160+
task_step_handle_result_impl(asyncio_state *state, TaskObj *task, PyObject *result);
161+
159162

160163
static int
161164
_is_coroutine(asyncio_state *state, PyObject *coro)
@@ -2032,15 +2035,16 @@ _asyncio.Task.__init__
20322035
loop: object = None
20332036
name: object = None
20342037
context: object = None
2038+
coro_result: object = NULL
20352039
20362040
A coroutine wrapped in a Future.
20372041
[clinic start generated code]*/
20382042

20392043
static int
20402044
_asyncio_Task___init___impl(TaskObj *self, PyObject *coro, PyObject *loop,
2041-
PyObject *name, PyObject *context)
2042-
/*[clinic end generated code: output=49ac96fe33d0e5c7 input=924522490c8ce825]*/
2043-
2045+
PyObject *name, PyObject *context,
2046+
PyObject *coro_result)
2047+
/*[clinic end generated code: output=e241855787412a77 input=3fcd7fb1c00d3f87]*/
20442048
{
20452049
if (future_init((FutureObj*)self, loop)) {
20462050
return -1;
@@ -2088,8 +2092,16 @@ _asyncio_Task___init___impl(TaskObj *self, PyObject *coro, PyObject *loop,
20882092
return -1;
20892093
}
20902094

2091-
if (task_call_step_soon(state, self, NULL)) {
2092-
return -1;
2095+
if (coro_result == NULL) {
2096+
if (task_call_step_soon(state, self, NULL)) {
2097+
return -1;
2098+
}
2099+
}
2100+
else {
2101+
PyObject * res = task_step_handle_result_impl(state, self, coro_result);
2102+
if (res == NULL) {
2103+
return -1;
2104+
}
20932105
}
20942106
return register_task(state, (PyObject*)self);
20952107
}
@@ -2827,6 +2839,22 @@ task_step_impl(asyncio_state *state, TaskObj *task, PyObject *exc)
28272839
Py_RETURN_NONE;
28282840
}
28292841

2842+
PyObject *ret = task_step_handle_result_impl(state, task, result);
2843+
Py_XDECREF(result);
2844+
return ret;
2845+
2846+
fail:
2847+
Py_XDECREF(result);
2848+
return NULL;
2849+
}
2850+
2851+
2852+
static PyObject *
2853+
task_step_handle_result_impl(asyncio_state *state, TaskObj *task, PyObject *result)
2854+
{
2855+
int res;
2856+
PyObject *o;
2857+
28302858
if (result == (PyObject*)task) {
28312859
/* We have a task that wants to await on itself */
2832
goto self_await;
@@ -2863,7 +2891,8 @@ task_step_impl(asyncio_state *state, TaskObj *task, PyObject *exc)
28632891
Py_DECREF(tmp);
28642892

28652893
/* task._fut_waiter = result */
2866-
task->task_fut_waiter = result; /* no incref is necessary */
2894+
Py_INCREF(result);
2895+
task->task_fut_waiter = result;
28672896

28682897
if (task->task_must_cancel) {
28692898
PyObject *r;
@@ -2956,7 +2985,8 @@ task_step_impl(asyncio_state *state, TaskObj *task, PyObject *exc)
29562985
Py_DECREF(tmp);
29572986

29582987
/* task._fut_waiter = result */
2959-
task->task_fut_waiter = result; /* no incref is necessary */
2988+
Py_INCREF(result);
2989+
task->task_fut_waiter = result;
29602990

29612991
if (task->task_must_cancel) {
29622992
PyObject *r;
@@ -2991,21 +3021,18 @@ task_step_impl(asyncio_state *state, TaskObj *task, PyObject *exc)
29913021
state, task, PyExc_RuntimeError,
29923022
"yield was used instead of yield from for "
29933023
"generator in task %R with %R", task, result);
2994-
Py_DECREF(result);
29953024
return o;
29963025
}
29973026

29983027
/* The `result` is none of the above */
29993028
o = task_set_error_soon(
30003029
state, task, PyExc_RuntimeError, "Task got bad yield: %R", result);
3001-
Py_DECREF(result);
30023030
return o;
30033031

30043032
self_await:
30053033
o = task_set_error_soon(
30063034
state, task, PyExc_RuntimeError,
30073035
"Task cannot await on itself: %R", task);
3008-
Py_DECREF(result);
30093036
return o;
30103037

30113038
yield_insteadof_yf:
@@ -3014,19 +3041,16 @@ task_step_impl(asyncio_state *state, TaskObj *task, PyObject *exc)
30143041
"yield was used instead of yield from "
30153042
"in task %R with %R",
30163043
task, result);
3017-
Py_DECREF(result);
30183044
return o;
30193045

30203046
different_loop:
30213047
o = task_set_error_soon(
30223048
state, task, PyExc_RuntimeError,
30233049
"Task %R got Future %R attached to a different loop",
30243050
task, result);
3025-
Py_DECREF(result);
30263051
return o;
30273052

30283053
fail:
3029-
Py_XDECREF(result);
30303054
return NULL;
30313055
}
30323056

0 commit comments

Comments
 (0)
0