proof of concept for deterministic close race

Author: belm0

tests/issue_96.py

@@ -81,7 +81,7 @@ async def handler(request):
     try:
         ws = await request.accept()
         await ws.send_message('issue96')
-        await trio.sleep(random()/1000)
+        await ws._for_testing_peer_closed_connection.wait()
         await trio.aclose_forcefully(ws._stream)
     except ConnectionClosed:
         pass
@@ -93,13 +93,14 @@ async def main():
             ssl_context=None)
         await nursery.start(serve)
 
-        for n in range(1000):
+        for n in range(1):
             logging.info('Connection %d', n)
             try:
                 connection = await connect_websocket(nursery, 'localhost', 8000,
                     '/', use_ssl=False)
                 await connection.get_message()
                 await connection.aclose()
+                await trio.sleep(.1)
             except ConnectionClosed:
                 pass
 

trio_websocket/_impl.py

@@ -544,6 +544,9 @@ def __init__(self, stream, wsproto, *, path=None,
         # Set once a WebSocket closed handshake takes place, i.e after a close
         # frame has been sent and a close frame has been received.
         self._close_handshake = trio.Event()
+        # Set immediately upon receiving closed event from peer.  Used to
+        # test close race conditions between client and server.
+        self._for_testing_peer_closed_connection = trio.Event()
 
     @property
     def closed(self):
@@ -807,6 +810,8 @@ async def _handle_connection_closed_event(self, event):
 
         :param event:
         '''
+        self._for_testing_peer_closed_connection.set()
+        await trio.sleep(0)
         await self._write_pending()
         await self._close_web_socket(event.code, event.reason or None)
         self._close_handshake.set()