Add basic eager async evaluation to Tasks (Python only) and TaskGroups

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

Author: DinoV

Lib/asyncio/taskgroups.py

@@ -8,6 +8,12 @@
 from . import events
 from . import exceptions
 from . import tasks
+from .futures import Future
+
+
+class PyCoroEagerResult:
+    def __init__(self, value):
+        self.value = value
 
 
 class TaskGroup:
@@ -24,6 +30,7 @@ def __init__(self):
         self._errors = []
         self._base_error = None
         self._on_completed_fut = None
+        self._enqueues = {}
 
     def __repr__(self):
         info = ['']
@@ -57,6 +64,16 @@ async def __aenter__(self):
 
         return self
 
+    def eager_eval(self, coro):
+        try:
+            fut = coro.send(None)
+        except StopIteration as e:
+            return PyCoroEagerResult(e.args[0] if e.args else None)
+        else:
+            task = self.create_task(coro)
+            task._set_fut_awaiter(fut)
+            return task
+
     async def __aexit__(self, et, exc, tb):
         self._exiting = True
         propagate_cancellation_error = None
@@ -89,6 +106,23 @@ async def __aexit__(self, et, exc, tb):
                 #
                 self._abort()
 
+        if self._enqueues:
+            for coro in self._enqueues:
+                res = self.eager_eval(coro)
+                if isinstance(res, PyCoroEagerResult):
+                    fut = self._enqueues[coro]
+                    if fut is not None:
+                        fut.set_result(res.value)
+                else:
+                    def queue_callback():
+                        fut = self._enqueues[coro]
+                        if fut is not None:
+                            res.add_done_callback(lambda task: fut.set_result(task.result()))
+                    queue_callback()
+
+            self._unfinished_tasks -= len(self._enqueues)
+            self._enqueues.clear()
+
         # We use while-loop here because "self._on_completed_fut"
         # can be cancelled multiple times if our parent task
         # is being cancelled repeatedly (or even once, when
@@ -153,6 +187,31 @@ def create_task(self, coro, *, name=None, context=None):
         self._tasks.add(task)
         return task
 
+    def enqueue(self, coro, no_future=True):
+        if not self._entered:
+            raise RuntimeError(f"TaskGroup {self!r} has not been entered")
+
+        if coro in self._enqueues:
+            return self._enqueues[coro]
+
+        if not self._enqueues:
+            # if the looop starts running because someone awaits, we want
+            # to run the co-routines which are enqueued as well.
+            self._loop.call_soon(self._enqueue_enqueus)
+
+        self._unfinished_tasks += 1
+        if no_future:
+            fut = self._enqueues[coro] = None
+        else:
+            fut = self._enqueues[coro] = Future(loop=self._loop)
+        return fut
+
+    def _enqueue_enqueus(self):
+        for coro in self._enqueues:
+            self.create_task(coro)
+        self._unfinished_tasks -= len(self._enqueues)
+        self._enqueues.clear()
+
     # Since Python 3.8 Tasks propagate all exceptions correctly,
     # except for KeyboardInterrupt and SystemExit which are
     # still considered special.

Lib/asyncio/tasks.py

@@ -117,7 +117,9 @@ def __init__(self, coro, *, loop=None, name=None, context=None):
         else:
             self._context = context
 
-        self._loop.call_soon(self.__step, context=self._context)
+        if not getattr(coro, "cr_suspended", False):
+            self._loop.call_soon(self.__step, context=self._context)
+
         _register_task(self)
 
     def __del__(self):
@@ -293,55 +295,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:
-                # Yielded Future must come from Future.__iter__().
-                if not self._check_future(result):
-                    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 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)
+            self._set_fut_awaiter(result)
+        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.
+    def _set_fut_awaiter(self, result):
+        blocking = getattr(result, '_asyncio_future_blocking', None)
+        if blocking is not None:
+            # Yielded Future must come from Future.__iter__().
+            if not self._check_future(result):
                 new_exc = RuntimeError(
-                    f'yield was used instead of yield from for '
-                    f'generator in task {self!r} with {result!r}')
+                    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 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:
-                # 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:
@@ -369,8 +374,8 @@ def __wakeup(self, future):
     pass
 else:
     # _CTask is needed for tests.
-    Task = _CTask = _asyncio.Task
-
+    #Task = _CTask = _asyncio.Task
+    pass
 
 def create_task(coro, *, name=None, context=None):
     """Schedule the execution of a coroutine object in a spawn task.

Lib/test/test_asyncio/test_taskgroups.py

@@ -722,6 +722,101 @@ async def coro(val):
             await t2
             self.assertEqual(2, ctx.get(cvar))
 
+    async def test_taskgroup_enqueue_01(self):
+
+        async def foo1():
+            await asyncio.sleep(0.1)
+            return 42
+
+        async def eager():
+            return 11
+
+        async with taskgroups.TaskGroup() as g:
+            t1 = g.enqueue(foo1(), no_future=False)
+            t2 = g.enqueue(eager(), no_future=False)
+
+        self.assertEqual(t1.result(), 42)
+        self.assertEqual(t2.result(), 11)
+
+    async def test_taskgroup_enqueue_02(self):
+
+        async def foo1():
+            return 42
+
+        async def eager():
+            return 11
+
+        async with taskgroups.TaskGroup() as g:
+            t1 = g.enqueue(foo1(), no_future=False)
+            t2 = g.enqueue(eager(), no_future=False)
+
+        self.assertEqual(t1.result(), 42)
+        self.assertEqual(t2.result(), 11)
+
+    async def test_taskgroup_fanout_task(self):
+        async def step(i):
+            if i == 0:
+                return
+            async with taskgroups.TaskGroup() as g:
+                for _ in range(6):
+                    g.create_task(step(i - 1))
+
+        import time
+        s = time.perf_counter()
+        await step(6)
+        e = time.perf_counter()
+        print(e-s)
+
+    async def test_taskgroup_fanout_enqueue(self):
+        async def step(i):
+            if i == 0:
+                return
+            async with taskgroups.TaskGroup() as g:
+                for _ in range(6):
+                    g.enqueue(step(i - 1))
+
+        import time
+        s = time.perf_counter()
+        await step(6)
+        e = time.perf_counter()
+        print(e-s)
+
+    async def test_taskgroup_fanout_enqueue_02(self):
+        async def intermediate2(i):
+            return await intermediate(i)
+
+        async def intermediate(i):
+            async with taskgroups.TaskGroup() as g:
+                for _ in range(6):
+                    g.enqueue(step(i - 1))
+
+        async def step(i):
+            if i == 0:
+                return
+
+            return await intermediate2(i)
+
+
+        import time
+        s = time.perf_counter()
+        await step(6)
+        e = time.perf_counter()
+        print(e-s)
+
+
+    async def test_taskgroup_fanout_enqueue_future(self):
+        async def step(i):
+            if i == 0:
+                return
+            async with taskgroups.TaskGroup() as g:
+                for _ in range(6):
+                    g.enqueue(step(i - 1), no_future=False)
+
+        import time
+        s = time.perf_counter()
+        await step(6)
+        e = time.perf_counter()
+        print(e-s)
 
 if __name__ == "__main__":
     unittest.main()