Eager task factory implementation

python · carljm · May 1, 2023 · Mar 20, 2023 · Mar 20, 2023 · Mar 20, 2023
commit a127f98a7cc88b2a5d2c719648840f3f413ee03b
diff --git a/Include/internal/pycore_global_objects_fini_generated.h b/Include/internal/pycore_global_objects_fini_generated.h
diff --git a/Include/internal/pycore_global_strings.h b/Include/internal/pycore_global_strings.h
@@ -339,6 +339,7 @@ struct _Py_global_strings {
        STRUCT_FOR_ID(copy)
        STRUCT_FOR_ID(copyreg)
        STRUCT_FOR_ID(coro)
+        STRUCT_FOR_ID(coro_result)
        STRUCT_FOR_ID(count)
        STRUCT_FOR_ID(cwd)
        STRUCT_FOR_ID(d)

diff --git a/Include/internal/pycore_runtime_init_generated.h b/Include/internal/pycore_runtime_init_generated.h
diff --git a/Include/internal/pycore_unicodeobject_generated.h b/Include/internal/pycore_unicodeobject_generated.h
diff --git a/Lib/asyncio/taskgroups.py b/Lib/asyncio/taskgroups.py
@@ -163,7 +163,8 @@ def create_task(self, coro, *, name=None, context=None):
            task = self._loop.create_task(coro)
        else:
            task = self._loop.create_task(coro, context=context)
-        tasks._set_task_name(task, name)
+        if name is not None and not task.done():  # If it's done already, it's a future
+            tasks._set_task_name(task, name)
        task.add_done_callback(self._on_task_done)
        self._tasks.add(task)
        return task

@@ -6,6 +6,7 @@
    'wait', 'wait_for', 'as_completed', 'sleep',
    'gather', 'shield', 'ensure_future', 'run_coroutine_threadsafe',
    'current_task', 'all_tasks',
+    'create_eager_task_factory', 'eager_task_factory',
    '_register_task', '_unregister_task', '_enter_task', '_leave_task',
 )

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


+_NOT_SET = object()
+
 class Task(futures._PyFuture):  # Inherit Python Task implementation
                                # from a Python Future implementation.

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

-    def __init__(self, coro, *, loop=None, name=None, context=None):
+    def __init__(self, coro, *, loop=None, name=None, context=None,
+                                coro_result=_NOT_SET):
        super().__init__(loop=loop)
        if self._source_traceback:
            del self._source_traceback[-1]
@@ -117,7 +121,10 @@ def __init__(self, coro, *, loop=None, name=None, context=None):
        else:
            self._context = context

-        self._loop.call_soon(self.__step, context=self._context)
+        if coro_result is _NOT_SET:
+            self._loop.call_soon(self.__step, context=self._context)
+        else:
+            self.__step_handle_result(coro_result)
        _register_task(self)

    def __del__(self):
@@ -287,55 +294,58 @@ def __step(self, exc=None):
        except BaseException as exc:
            super().set_exception(exc)
        else:
-            blocking = getattr(result, '_asyncio_future_blocking', None)
-            if blocking is not None:
+            self.__step_handle_result(result)
+        finally:
+            _leave_task(self._loop, self)
+            self = None  # Needed to break cycles when an exception occurs.
+
+    def __step_handle_result(self, result):
+        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 futures._get_loop(result) is not self._loop:
+                new_exc = RuntimeError(
+                    f'Task {self!r} got Future '
+                    f'{result!r} attached to a different loop')
+                self._loop.call_soon(
+                    self.__step, new_exc, context=self._context)
+            elif blocking:
+                if result is self:
                    new_exc = RuntimeError(
-                        f'Task {self!r} got Future '
-                        f'{result!r} attached to a different loop')
+                        f'Task cannot await on itself: {self!r}')
                    self._loop.call_soon(
                        self.__step, new_exc, context=self._context)
-                elif blocking:
-                    if result is self:
-                        new_exc = RuntimeError(
-                            f'Task cannot await on itself: {self!r}')
-                        self._loop.call_soon(
-                            self.__step, new_exc, context=self._context)
-                    else:
-                        result._asyncio_future_blocking = False
-                        result.add_done_callback(
-                            self.__wakeup, context=self._context)
-                        self._fut_waiter = result
-                        if self._must_cancel:
-                            if self._fut_waiter.cancel(
-                                    msg=self._cancel_message):
-                                self._must_cancel = False
                else:
-                    new_exc = RuntimeError(
-                        f'yield was used instead of yield from '
-                        f'in task {self!r} with {result!r}')
-                    self._loop.call_soon(
-                        self.__step, new_exc, context=self._context)
-
-            elif result is None:
-                # Bare yield relinquishes control for one event loop iteration.
-                self._loop.call_soon(self.__step, context=self._context)
-            elif inspect.isgenerator(result):
-                # Yielding a generator is just wrong.
-                new_exc = RuntimeError(
-                    f'yield was used instead of yield from for '
-                    f'generator in task {self!r} with {result!r}')
-                self._loop.call_soon(
-                    self.__step, new_exc, context=self._context)
+                    result._asyncio_future_blocking = False
+                    result.add_done_callback(
+                        self.__wakeup, context=self._context)
+                    self._fut_waiter = result
+                    if self._must_cancel:
+                        if self._fut_waiter.cancel(
+                                msg=self._cancel_message):
+                            self._must_cancel = False
            else:
-                # Yielding something else is an error.
-                new_exc = RuntimeError(f'Task got bad yield: {result!r}')
+                new_exc = RuntimeError(
+                    f'yield was used instead of yield from '
+                    f'in task {self!r} with {result!r}')
                self._loop.call_soon(
                    self.__step, new_exc, context=self._context)
-        finally:
-            _leave_task(self._loop, self)
-            self = None  # Needed to break cycles when an exception occurs.
+
+        elif result is None:
+            # Bare yield relinquishes control for one event loop iteration.
+            self._loop.call_soon(self.__step, context=self._context)
+        elif inspect.isgenerator(result):
+            # Yielding a generator is just wrong.
+            new_exc = RuntimeError(
+                f'yield was used instead of yield from for '
+                f'generator in task {self!r} with {result!r}')
+            self._loop.call_soon(
+                self.__step, new_exc, context=self._context)
+        else:
+            # Yielding something else is an error.
+            new_exc = RuntimeError(f'Task got bad yield: {result!r}')
+            self._loop.call_soon(
+                self.__step, new_exc, context=self._context)

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


+def create_eager_task_factory(custom_task_constructor):
+
+    def factory(loop, coro, *, name=None, context=None):
+        loop._check_closed()
+        if not loop.is_running():
+            return custom_task_constructor(coro, loop=loop, name=name, context=context)
+
+        try:
+            result = coro.send(None)
+        except StopIteration as si:
+            fut = loop.create_future()
+            fut.set_result(si.value)
+            return fut
+        except Exception as ex:
+            fut = loop.create_future()
+            fut.set_exception(ex)
+            return fut
+        else:
+            task = custom_task_constructor(
+                coro, loop=loop, name=name, context=context, coro_result=result)
+            if task._source_traceback:
+                del task._source_traceback[-1]
+            return task
+
+    return factory
+
+eager_task_factory = create_eager_task_factory(Task)
+
+
 # WeakSet containing all alive tasks.
 _all_tasks = weakref.WeakSet()


@@ -156,6 +156,9 @@ class _asyncio.Future "FutureObj *" "&Future_Type"
 /* Get FutureIter from Future */
 static PyObject * future_new_iter(PyObject *);

+static PyObject *
+task_step_handle_result_impl(asyncio_state *state, TaskObj *task, PyObject *result);
+

 static int
 _is_coroutine(asyncio_state *state, PyObject *coro)
@@ -2025,15 +2028,16 @@ _asyncio.Task.__init__
    loop: object = None
    name: object = None
    context: object = None
+    coro_result: object = NULL

 A coroutine wrapped in a Future.
 [clinic start generated code]*/

 static int
 _asyncio_Task___init___impl(TaskObj *self, PyObject *coro, PyObject *loop,
-                            PyObject *name, PyObject *context)
-/*[clinic end generated code: output=49ac96fe33d0e5c7 input=924522490c8ce825]*/
-
+                            PyObject *name, PyObject *context,
+                            PyObject *coro_result)
+/*[clinic end generated code: output=e241855787412a77 input=3fcd7fb1c00d3f87]*/
 {
    if (future_init((FutureObj*)self, loop)) {
        return -1;
@@ -2081,8 +2085,16 @@ _asyncio_Task___init___impl(TaskObj *self, PyObject *coro, PyObject *loop,
        return -1;
    }

-    if (task_call_step_soon(state, self, NULL)) {
-        return -1;
+    if (coro_result == NULL) {
+        if (task_call_step_soon(state, self, NULL)) {
+            return -1;
+        }
+    }
+    else {
+        PyObject * res = task_step_handle_result_impl(state, self, coro_result);
+        if (res == NULL) {
+            return -1;
+        }
    }
    return register_task(state, (PyObject*)self);
 }
@@ -2822,6 +2834,22 @@ task_step_impl(asyncio_state *state, TaskObj *task, PyObject *exc)
        Py_RETURN_NONE;
    }

+    PyObject *ret = task_step_handle_result_impl(state, task, result);
+    Py_XDECREF(result);
+    return ret;
+
+fail:
+    Py_XDECREF(result);
+    return NULL;
+}
+
+
+static PyObject *
+task_step_handle_result_impl(asyncio_state *state, TaskObj *task, PyObject *result)
+{
+    int res;
+    PyObject *o;
+
    if (result == (PyObject*)task) {
        /* We have a task that wants to await on itself */
        goto self_await;
@@ -2858,7 +2886,8 @@ task_step_impl(asyncio_state *state, TaskObj *task, PyObject *exc)
        Py_DECREF(tmp);

        /* task._fut_waiter = result */
-        task->task_fut_waiter = result;  /* no incref is necessary */
+        Py_INCREF(result);
+        task->task_fut_waiter = result;

        if (task->task_must_cancel) {
            PyObject *r;
@@ -2951,7 +2980,8 @@ task_step_impl(asyncio_state *state, TaskObj *task, PyObject *exc)
        Py_DECREF(tmp);

        /* task._fut_waiter = result */
-        task->task_fut_waiter = result;  /* no incref is necessary */
+        Py_INCREF(result);
+        task->task_fut_waiter = result;

        if (task->task_must_cancel) {
            PyObject *r;
@@ -2986,21 +3016,18 @@ task_step_impl(asyncio_state *state, TaskObj *task, PyObject *exc)
            state, task, PyExc_RuntimeError,
            "yield was used instead of yield from for "
            "generator in task %R with %R", task, result);
-        Py_DECREF(result);
        return o;
    }

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

 self_await:
    o = task_set_error_soon(
        state, task, PyExc_RuntimeError,
        "Task cannot await on itself: %R", task);
-    Py_DECREF(result);
    return o;

 yield_insteadof_yf:
@@ -3009,19 +3036,16 @@ task_step_impl(asyncio_state *state, TaskObj *task, PyObject *exc)
        "yield was used instead of yield from "
        "in task %R with %R",
        task, result);
-    Py_DECREF(result);
    return o;

 different_loop:
    o = task_set_error_soon(
        state, task, PyExc_RuntimeError,
        "Task %R got Future %R attached to a different loop",
        task, result);
-    Py_DECREF(result);
    return o;

 fail:
-    Py_XDECREF(result);
    return NULL;
 }