await _connect and inline read_messages callback to _connect

Author: ricklamers

jupyter_server/gateway/handlers.py

@@ -134,7 +134,7 @@ def __init__(self, **kwargs):
         self.ws_future = Future()
         self.disconnected = False
 
-    async def _connect(self, kernel_id):
+    async def _connect(self, kernel_id, message_callback):
         # websocket is initialized before connection
         self.ws = None
         self.kernel_id = kernel_id
@@ -150,6 +150,12 @@ async def _connect(self, kernel_id):
         self.ws_future = websocket_connect(request)
         self.ws_future.add_done_callback(self._connection_done)
 
+        loop = IOLoop.current()
+        loop.add_future(
+            self.ws_future,
+            lambda future: self._read_messages(message_callback)
+        )
+
     def _connection_done(self, fut):
         if not self.disconnected and fut.exception() is None:  # prevent concurrent.futures._base.CancelledError
             self.ws = fut.result()
@@ -188,18 +194,13 @@ async def _read_messages(self, callback):
 
         if not self.disconnected: # if websocket is not disconnected by client, attept to reconnect to Gateway
             self.log.info("Attempting to re-establish the connection to Gateway: {}".format(self.kernel_id))
-            self._connect(self.kernel_id)
             loop = IOLoop.current()
-            loop.add_future(self.ws_future, lambda future: self._read_messages(callback))
+            loop.spawn_callback(self._connect, self.kernel_id, callback)
 
     def on_open(self, kernel_id, message_callback, **kwargs):
         """Web socket connection open against gateway server."""
-        self._connect(kernel_id)
         loop = IOLoop.current()
-        loop.add_future(
-            self.ws_future,
-            lambda future: self._read_messages(message_callback)
-        )
+        loop.spawn_callback(self._connect, kernel_id, message_callback)
 
     def on_message(self, message):
         """Send message to gateway server."""