extmod/asyncio/stream.py: Fix cancellation handling of start_server.

The following code:

  server = await asyncio.start_server(...)
  async with server:
    ... code that raises ...

would lose the original exception because the server's task would not have
had a chance to be scheduled yet, and so awaiting the task in wait_closed
would raise the cancellation instead of the original exception.

Additionally, ensures that explicitly cancelling the parent task delivers
the cancellation correctly (previously was masked by the server loop), now
this only happens if the server was closed, not when the task was
cancelled.

Signed-off-by: Jim Mussared <jim.mussared@gmail.com>

1 files changed, 23 insertions, 3 deletions

diff --git a/extmod/asyncio/stream.py b/extmod/asyncio/stream.py
index c47c48cf0..5547bfbd5 100644
--- a/extmod/asyncio/stream.py
+++ b/extmod/asyncio/stream.py
@@ -127,20 +127,30 @@ class Server:
         await self.wait_closed()
 
     def close(self):
+        # Note: the _serve task must have already started by now due to the sleep
+        # in start_server, so `state` won't be clobbered at the start of _serve.
+        self.state = True
         self.task.cancel()
 
     async def wait_closed(self):
         await self.task
 
     async def _serve(self, s, cb):
+        self.state = False
         # Accept incoming connections
         while True:
             try:
                 yield core._io_queue.queue_read(s)
-            except core.CancelledError:
-                # Shutdown server
+            except core.CancelledError as er:
+                # The server task was cancelled, shutdown server and close socket.
                 s.close()
-                return
+                if self.state:
+                    # If the server was explicitly closed, ignore the cancellation.
+                    return
+                else:
+                    # Otherwise e.g. the parent task was cancelled, propagate
+                    # cancellation.
+                    raise er
             try:
                 s2, addr = s.accept()
             except:
@@ -167,6 +177,16 @@ async def start_server(cb, host, port, backlog=5):
     # Create and return server object and task.
     srv = Server()
     srv.task = core.create_task(srv._serve(s, cb))
+    try:
+        # Ensure that the _serve task has been scheduled so that it gets to
+        # handle cancellation.
+        await core.sleep_ms(0)
+    except core.CancelledError as er:
+        # If the parent task is cancelled during this first sleep, then
+        # we will leak the task and it will sit waiting for the socket, so
+        # cancel it.
+        srv.task.cancel()
+        raise er
     return srv