[bitbake-devel][PATCH v2 01/10] asyncrpc: Add Task Group

Adds a Task Group object which can be used to manage multiple
asynchronous tasks and correctly cancel them if an exception is raised

Signed-off-by: Joshua Watt <[email protected]>
---
 lib/bb/asyncrpc/__init__.py  |  1 +
 lib/bb/asyncrpc/taskgroup.py | 52 ++++++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+)
 create mode 100644 lib/bb/asyncrpc/taskgroup.py

diff --git a/lib/bb/asyncrpc/__init__.py b/lib/bb/asyncrpc/__init__.py
index a4371643d..2f0956dfa 100644
--- a/lib/bb/asyncrpc/__init__.py
+++ b/lib/bb/asyncrpc/__init__.py
@@ -14,3 +14,4 @@ from .exceptions import (
     ConnectionClosedError,
     InvokeError,
 )
+from .taskgroup import TaskGroup
diff --git a/lib/bb/asyncrpc/taskgroup.py b/lib/bb/asyncrpc/taskgroup.py
new file mode 100644
index 000000000..b61f40385
--- /dev/null
+++ b/lib/bb/asyncrpc/taskgroup.py
@@ -0,0 +1,52 @@
+#
+# Copyright BitBake Contributors
+#
+# SPDX-License-Identifier: GPL-2.0-only
+#
+import asyncio
+import logging
+
+logger = logging.getLogger("asyncio.TaskGroup")
+
+
+class TaskGroup(object):
+    def __init__(self):
+        self._tasks = []
+
+    def create_task(self, coro, **kwargs):
+        self._tasks.append(asyncio.create_task(coro, **kwargs))
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        try:
+            if exc is None:
+                while self._tasks:
+                    done, pending = await asyncio.wait(
+                        self._tasks, return_when=asyncio.FIRST_COMPLETED
+                    )
+                    self._tasks = pending
+
+                    for t in done:
+                        try:
+                            await t
+                        except asyncio.CancelledError:
+                            pass
+
+        finally:
+            for t in self._tasks:
+                t.cancel()
+                try:
+                    await t
+                except:
+                    # Ignore exceptions
+                    pass
+
+        return False
+
+    @classmethod
+    async def run(cls, *coros):
+        async with cls() as group:
+            for c in coros:
+                group.create_task(c)
-- 
2.54.0