fix flaky asgi tests and wait for the client's close frame in asgi websockets

rolo/asgi.py

@@ -345,6 +345,10 @@ class ASGIWebSocketAdapter(rolows.WebSocketAdapter):
     _receive: "ASGIReceiveCallable"
     _send: "ASGISendCallable"
 
+    close_timeout: float = 5
+    """Seconds to wait for the client to echo the server's close frame, before the ASGI app returns and leaves
+    the TCP connection to the server."""
+
     def __init__(
         self,
         scope: "WebsocketScope",
@@ -357,12 +361,35 @@ def __init__(
         self._send = send
         self._loop = loop
         self._last_received_at = time.monotonic()
+        self._close_sent = False
+        self._disconnected = False
 
     async def asgi_send_async(self, event: "_WebsocketResponse"):
         await self._send(event)
 
     async def asgi_receive_async(self) -> "_WebsocketRequest":
-        return await self._receive()
+        event = await self._receive()
+        if event["type"] == "websocket.disconnect":
+            self._disconnected = True
+        return event
+
+    async def wait_closed(self) -> None:
+        """
+        Waits for the client to echo the close frame the server sent, so the close handshake completes before the
+        ASGI app returns. An ASGI server can terminate the TCP connection as soon as the app returns, and the
+        client's close frame then arriving on a closed socket makes the kernel reset the connection (RFC 6455
+        section 5.5.1). Frames the client sends in the meantime are discarded.
+        """
+        if not self._close_sent or self._disconnected:
+            return
+        try:
+            await asyncio.wait_for(self._receive_until_disconnect(), self.close_timeout)
+        except asyncio.TimeoutError:
+            pass
+
+    async def _receive_until_disconnect(self) -> None:
+        while not self._disconnected:
+            await self.asgi_receive_async()
 
     def asgi_send(self, event: "_WebsocketResponse", timeout: float = None) -> None:
         """
@@ -496,6 +523,7 @@ def accept(
         )
 
     def close(self, code: int = 1001, reason: str = None, timeout: float = None):
+        self._close_sent = True
         self.asgi_send(
             {
                 "type": "websocket.close",
@@ -660,6 +688,8 @@ async def handle_websocket(
             )
             raise
 
+        await environ["rolo.websocket"].wait_closed()
+
     async def handle_lifespan(
         self, scope: "HTTPScope", receive: "ASGIReceiveCallable", send: "ASGISendCallable"
     ):

tests/test_asgi.py

@@ -175,10 +175,14 @@ class IterableResponse:
         def __init__(self, data: list[bytes]):
             self.data = data
             self.closed = False
+            # holds the server after the last packet, so it can't close the iterable before
+            # the client has checked that it's still open
+            self.release = threading.Event()
 
         def __iter__(self):
             for packet in self.data:
                 yield packet
+            assert self.release.wait(timeout=5)
 
         def close(self):
             # should be called through the werkzeug layers
@@ -198,6 +202,7 @@ def app(request: Request) -> Response:
     assert next(gen) == b"foo"
     assert next(gen) == b"bar"
     assert not iterable.closed
+    iterable.release.set()
 
     with pytest.raises(StopIteration):
         next(gen)

tests/websocket/test_websockets.py

@@ -221,7 +221,11 @@ def app(request: WebSocketRequest):
 
     frame = client.recv_frame()
     assert frame.opcode == websocket.ABNF.OPCODE_CLOSE
+    # a server that terminates the TCP connection without waiting for the echo makes the kernel reset the
+    # connection once the echo arrives, and the reset then fails the read below
+    time.sleep(0.1)
     client.send_close(status=struct.unpack("!H", frame.data[:2])[0])
+    time.sleep(0.1)
 
     client.sock.settimeout(5)
     assert client.sock.recv(1) == b"", "expected the server to terminate the TCP connection"
@@ -273,6 +277,7 @@ def echo_headers(request: WebSocketRequest):
     assert headers["Connection"] == "Upgrade"
     assert headers["Authorization"] == "Basic let-me-in"
     assert headers["CasedHeader"] == "hello"
+    client.close()
 
 
 def test_websocket_reject(serve_websocket_listener):
@@ -318,6 +323,7 @@ def echo_headers(request: WebSocketRequest):
 
     assert received.get(timeout=5) == "textbar"
     assert received.get(timeout=5) == b"bar"
+    client.close()
 
 
 def test_last_received_at_message(serve_websocket_listener):
@@ -436,6 +442,7 @@ def echo_headers(request: WebSocketRequest):
 
     e = match.get(timeout=5)
     assert e.match("Cannot send data type <class 'dict'> over websocket")
+    client.close()
 
 
 def test_router_integration(serve_websocket_listener):
@@ -457,6 +464,7 @@ def _handler(request: WebSocketRequest, request_args: dict):
     assert client.recv() == "foo"
     assert client.recv() == "id=bar"
     assert "CasedHeader" in json.loads(client.recv())
+    client.close()
 
 
 def test_send_many_messages(serve_websocket_listener):