Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Await _connect and inline read_messages callback to _connect #350

Merged
merged 1 commit into from
Nov 21, 2020
Merged

Await _connect and inline read_messages callback to _connect #350

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 9 additions & 8 deletions jupyter_server/gateway/handlers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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()
Expand Down Expand Up @@ -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."""
Expand Down