extmod/asyncio: Fix Stream.wait_closed.

extmod/asyncio/stream.py

@@ -2,13 +2,15 @@
 # MIT license; Copyright (c) 2019-2020 Damien P. George
 
 from . import core
+from .funcs import wait_for_ms
 
 
 class Stream:
     def __init__(self, s, e={}):
         self.s = s
         self.e = e
         self.out_buf = b""
+        self._server = False
 
     def get_extra_info(self, v):
         return self.e[v]
@@ -17,7 +19,15 @@ def close(self):
         pass
 
     async def wait_closed(self):
-        # TODO yield?
+        while True and self._server:
+            try:
+                # this makes sure there is no pending data in the socket
+                # to avoid ECONNRESET at client side.
+                if not await wait_for_ms(self.read(), 10):
+                    break
+            except Exception:  # ECONNRESET, TimeoutError
+                break
+
         self.s.close()
 
     # async
@@ -171,6 +181,7 @@ async def _serve(self, s, cb, ssl):
                     continue
             s2.setblocking(False)
             s2s = Stream(s2, {"peername": addr})
+            s2s._server = True
             core.create_task(cb(s2s, s2s))
 
 

tests/multi_net/asyncio_tcp_server_client_drain.py

@@ -0,0 +1,59 @@
+# Test asyncio TCP server and client using start_server() and open_connection()
+
+try:
+    import asyncio
+except ImportError:
+    print("SKIP")
+    raise SystemExit
+
+PORT = 8000
+
+
+async def handle_connection(reader, writer):
+    # Test that peername exists (but don't check its value, it changes)
+    writer.get_extra_info("peername")
+
+    data = await reader.read(11)  # leaving bytes unread triggers the error
+    # data = await reader.read(13)  # reading everything
+    print("echo:", data.replace(b"\r\n", b""))
+    writer.write(data)
+    await writer.drain()
+
+    print("close")
+    writer.close()
+    await writer.wait_closed()
+
+    print("done")
+    ev.set()
+
+
+async def tcp_server():
+    global ev
+    ev = asyncio.Event()
+    server = await asyncio.start_server(handle_connection, "0.0.0.0", PORT)
+    print("server running")
+    multitest.next()
+    async with server:
+        await asyncio.wait_for(ev.wait(), 10)
+
+
+async def tcp_client(message):
+    reader, writer = await asyncio.open_connection(IP, PORT)
+    print("write:", message.replace(b"\r\n", b""))
+    writer.write(message)
+    await writer.drain()
+    data = await reader.read()  # if the client doesn't read exactly or is fast
+    # enough e.g.
+    # data = await reader.read(11)  # this works on unix because is
+    # fast enough, however esp32 fails.
+    print("read:", data.replace(b"\r\n", b""))
+
+
+def instance0():
+    multitest.globals(IP=multitest.get_network_ip())
+    asyncio.run(tcp_server())
+
+
+def instance1():
+    multitest.next()
+    asyncio.run(tcp_client(b"client data\r\n"))

tests/multi_net/asyncio_tcp_server_client_drain.py.exp

@@ -0,0 +1,8 @@
+--- instance0 ---
+server running
+echo: b'client data'
+close
+done
+--- instance1 ---
+write: b'client data'
+read: b'client data'