micropython · smurfix · Jun 20, 2022 · Jun 22, 2022 · Sep 2, 2023 · Sep 2, 2023
diff --git a/extmod/asyncio/__init__.py b/extmod/asyncio/__init__.py
@@ -14,8 +14,10 @@
     "Lock": "lock",
     "open_connection": "stream",
     "start_server": "stream",
+    "run_server": "stream",
     "StreamReader": "stream",
     "StreamWriter": "stream",
+    "TaskGroup": "taskgroup",
 }
 
 

diff --git a/extmod/asyncio/core.py b/extmod/asyncio/core.py
@@ -15,6 +15,25 @@
 # Exceptions
 
 
+class BaseExceptionGroup(BaseException):
+    def split(self, typ):
+        a, b = [], []
+        if isinstance(typ, (BaseException, tuple)):
+            for err in self.args[1]:
+                (a if isinstance(err, typ) else b).append(err)
+        else:
+            for err in self.args[1]:
+                (a if typ(err) else b).append(err)
+        return a, b
+
+
+class ExceptionGroup(Exception):  # TODO cannot also inherit from BaseExceptionGroup
+    pass
+
+
+ExceptionGroup.split = BaseExceptionGroup.split
+
+
 class CancelledError(BaseException):
     pass
 
@@ -32,7 +51,7 @@ class TimeoutError(Exception):
 
 
 # "Yield" once, then raise StopIteration
-class SingletonGenerator:
+class SleepHandler:
     def __init__(self):
         self.state = None
         self.exc = StopIteration()
@@ -51,9 +70,10 @@ def __next__(self):
 
 
 # Pause task execution for the given time (integer in milliseconds, uPy extension)
-# Use a SingletonGenerator to do it without allocating on the heap
-def sleep_ms(t, sgen=SingletonGenerator()):
-    assert sgen.state is None
+# Try not to allocate a SleepHandler on the heap if possible
+def sleep_ms(t, sgen=SleepHandler()):
+    if sgen.state is not None:  # the static one is busy
+        sgen = SleepHandler()
     sgen.state = ticks_add(ticks(), max(0, t))
     return sgen
 
@@ -222,6 +242,9 @@ def run_until_complete(main_task=None):
                 _exc_context["exception"] = exc
                 _exc_context["future"] = t
                 Loop.call_exception_handler(_exc_context)
+                # XXX if we do await it later,
+                # leaving t.data as None will cause a fault.
+                t.data = exc
 
 
 # Create a new task from a coroutine and run it until it finishes

diff --git a/extmod/asyncio/manifest.py b/extmod/asyncio/manifest.py
@@ -9,6 +9,7 @@
         "funcs.py",
         "lock.py",
         "stream.py",
+        "taskgroup.py",
     ),
     base_path="..",
     opt=3,

diff --git a/extmod/asyncio/stream.py b/extmod/asyncio/stream.py
@@ -17,7 +17,8 @@ async def __aenter__(self):
         return self
 
     async def __aexit__(self, exc_type, exc, tb):
-        await self.close()
+        self.s.close()
+        pass
 
     def close(self):
         pass
@@ -152,7 +153,7 @@ async def _serve(self, s, cb):
 
 
 # Helper function to start a TCP stream server, running as a new task
-# TODO could use an accept-callback on socket read activity instead of creating a task
+# DOES NOT USE TASKGROUPS. Use run_server instead
 async def start_server(cb, host, port, backlog=5):
     import socket
 
@@ -170,6 +171,47 @@ async def start_server(cb, host, port, backlog=5):
     return srv
 
 
+# Helper task to run a TCP stream server.
+# Callbacks (i.e. connection handlers) may run in a different taskgroup.
+async def run_server(cb, host, port, backlog=5, taskgroup=None):
+    import socket
+
+    # Create and bind server socket.
+    host = socket.getaddrinfo(host, port)[0]  # TODO this is blocking!
+    s = socket.socket()
+    s.setblocking(False)
+    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    s.bind(host[-1])
+    s.listen(backlog)
+
+    if taskgroup is None:
+        from . import TaskGroup
+
+        async with TaskGroup() as tg:
+            await _run_server(tg, s, cb)
+    else:
+        await _run_server(taskgroup, s, cb)
+
+
+async def _run_server(tg, s, cb):
+    while True:
+        try:
+            yield core._io_queue.queue_read(s)
+        except core.CancelledError:
+            # Shutdown server
+            s.close()
+            return
+        try:
+            s2, addr = s.accept()
+        except Exception:
+            # Ignore a failed accept
+            continue
+
+        s2.setblocking(False)
+        s2s = Stream(s2, {"peername": addr})
+        tg.create_task(cb(s2s, s2s))
+
+
 ################################################################################
 # Legacy uasyncio compatibility