Fix WebSocket race condition

Author: tibdex

jupyter_server_proxy/handlers.py

@@ -231,6 +231,8 @@ async def proxy_open(self, host, port, proxied_path=''):
 
         client_uri = self.get_client_uri('ws', host, port, proxied_path)
         headers = self.request.headers
+        current_loop = ioloop.IOLoop.current()
+        ws_connected = current_loop.asyncio_loop.create_future()
 
         def message_cb(message):
             """
@@ -261,10 +263,15 @@ async def start_websocket_connection():
             request = httpclient.HTTPRequest(url=client_uri, headers=headers)
             self.ws = await pingable_ws_connect(request=request,
                 on_message_callback=message_cb, on_ping_callback=ping_cb)
+            ws_connected.set_result(True)
             self._record_activity()
             self.log.info('Websocket connection established to {}'.format(client_uri))
 
-        ioloop.IOLoop.current().add_callback(start_websocket_connection)
+        current_loop.add_callback(start_websocket_connection)
+        # Wait for the WebSocket to be connected before resolving.
+        # Otherwise, messages sent by the client before the
+        # WebSocket successful connection would be dropped.
+        await ws_connected
 
 
     def proxy_request_headers(self):