class KernelGatewayWSClient(LoggingConfigurable):
"""Proxy web socket connection to a kernel gateway."""
def __init__(self):
self.ws = None
self.ws_future = Future()
@gen.coroutine
def _connect(self, kernel_id):
ws_url = url_path_join(KG_URL.replace("http", "ws"), "/api/kernels", url_escape(kernel_id), "channels")
self.log.info("Connecting to {}".format(ws_url))
request = HTTPRequest(ws_url, headers=KG_HEADERS, validate_cert=VALIDATE_KG_CERT)
self.ws_future = websocket_connect(request)
self.ws = yield self.ws_future
# TODO: handle connection errors/timeout
def _disconnect(self):
if self.ws is not None:
# Close connection
self.ws.close()
elif not self.ws_future.done():
# Cancel pending connection
self.ws_future.cancel()
@gen.coroutine
def _read_messages(self, callback):
"""Read messages from server."""
while True:
message = yield self.ws.read_message()
if message is None:
break # TODO: handle socket close
callback(message)
def on_open(self, kernel_id, message_callback, **kwargs):
"""Web socket connection open."""
self._connect(kernel_id)
loop = IOLoop.current()
loop.add_future(self.ws_future, lambda future: self._read_messages(message_callback))
def on_message(self, message):
"""Send message to server."""
if self.ws is None:
loop = IOLoop.current()
loop.add_future(self.ws_future, lambda future: self._write_message(message))
else:
self._write_message(message)
def _write_message(self, message):
"""Send message to server."""
self.ws.write_message(message)
def on_close(self):
"""Web socket closed event."""
self._disconnect()
def test_future_set_result_unless_cancelled(self):
fut = Future()
future_set_result_unless_cancelled(fut, 42)
self.assertEqual(fut.result(), 42)
self.assertFalse(fut.cancelled())
fut = Future()
fut.cancel()
is_cancelled = fut.cancelled()
future_set_result_unless_cancelled(fut, 42)
self.assertEqual(fut.cancelled(), is_cancelled)
if not is_cancelled:
self.assertEqual(fut.result(), 42)
def wait(self, timeout: Union[float, datetime.timedelta] = None) -> "Future[None]":
"""Block until the internal flag is true.
Returns a Future, which raises `tornado.util.TimeoutError` after a
timeout.
"""
fut = Future() # type: Future[None]
if self._value:
fut.set_result(None)
return fut
self._waiters.add(fut)
fut.add_done_callback(lambda fut: self._waiters.remove(fut))
if timeout is None:
return fut
else:
timeout_fut = gen.with_timeout(
timeout, fut, quiet_exceptions=(CancelledError,)
)
# This is a slightly clumsy workaround for the fact that
# gen.with_timeout doesn't cancel its futures. Cancelling
# fut will remove it from the waiters list.
timeout_fut.add_done_callback(
lambda tf: fut.cancel() if not fut.done() else None
)
return timeout_fut
def wait(
self,
timeout: Union[float,
datetime.timedelta] = None) -> Awaitable[None]:
"""Block until the internal flag is true.
Returns an awaitable, which raises `tornado.util.TimeoutError` after a
timeout.
"""
fut = Future() # type: Future[None]
if self._value:
fut.set_result(None)
return fut
self._waiters.add(fut)
fut.add_done_callback(lambda fut: self._waiters.remove(fut))
if timeout is None:
return fut
else:
timeout_fut = gen.with_timeout(timeout,
fut,
quiet_exceptions=(CancelledError, ))
# This is a slightly clumsy workaround for the fact that
# gen.with_timeout doesn't cancel its futures. Cancelling
# fut will remove it from the waiters list.
timeout_fut.add_done_callback(lambda tf: fut.cancel()
if not fut.done() else None)
return timeout_fut
class GatewayWebSocketClient(LoggingConfigurable):
"""Proxy web socket connection to a kernel/enterprise gateway."""
def __init__(self, **kwargs):
super(GatewayWebSocketClient, self).__init__(**kwargs)
self.kernel_id = None
self.ws = None
self.ws_future = Future()
self.ws_future_cancelled = False
@gen.coroutine
def _connect(self, kernel_id):
self.kernel_id = kernel_id
ws_url = url_path_join(GatewayClient.instance().ws_url,
GatewayClient.instance().kernels_endpoint,
url_escape(kernel_id), 'channels')
self.log.info('Connecting to {}'.format(ws_url))
kwargs = {}
kwargs = GatewayClient.instance().load_connection_args(**kwargs)
request = HTTPRequest(ws_url, **kwargs)
self.ws_future = websocket_connect(request)
self.ws_future.add_done_callback(self._connection_done)
def _connection_done(self, fut):
if not self.ws_future_cancelled: # prevent concurrent.futures._base.CancelledError
self.ws = fut.result()
self.log.debug("Connection is ready: ws: {}".format(self.ws))
else:
self.log.warning(
"Websocket connection has been cancelled via client disconnect before its establishment. "
"Kernel with ID '{}' may not be terminated on GatewayClient: {}"
.format(self.kernel_id,
GatewayClient.instance().url))
def _disconnect(self):
if self.ws is not None:
# Close connection
self.ws.close()
elif not self.ws_future.done():
# Cancel pending connection. Since future.cancel() is a noop on tornado, we'll track cancellation locally
self.ws_future.cancel()
self.ws_future_cancelled = True
self.log.debug("_disconnect: ws_future_cancelled: {}".format(
self.ws_future_cancelled))
@gen.coroutine
def _read_messages(self, callback):
"""Read messages from gateway server."""
while True:
message = None
if not self.ws_future_cancelled:
try:
message = yield self.ws.read_message()
except Exception as e:
self.log.error(
"Exception reading message from websocket: {}".format(
e)) # , exc_info=True)
if message is None:
break
callback(
message
) # pass back to notebook client (see self.on_open and WebSocketChannelsHandler.open)
else: # ws cancelled - stop reading
break
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))
def on_message(self, message):
"""Send message to gateway server."""
if self.ws is None:
loop = IOLoop.current()
loop.add_future(self.ws_future,
lambda future: self._write_message(message))
else:
self._write_message(message)
def _write_message(self, message):
"""Send message to gateway server."""
try:
if not self.ws_future_cancelled:
self.ws.write_message(message)
except Exception as e:
self.log.error("Exception writing message to websocket: {}".format(
e)) # , exc_info=True)
def on_close(self):
"""Web socket closed event."""
self._disconnect()
class KernelGatewayWSClient(LoggingConfigurable):
"""Proxy web socket connection to a kernel/enterprise gateway."""
def __init__(self, **kwargs):
super(KernelGatewayWSClient, self).__init__(**kwargs)
self.kernel_id = None
self.ws = None
self.ws_future = Future()
self.ws_future_cancelled = False
@gen.coroutine
def _connect(self, kernel_id):
self.kernel_id = kernel_id
ws_url = url_path_join(
os.getenv('KG_WS_URL', KG_URL.replace('http', 'ws')),
'/api/kernels', url_escape(kernel_id), 'channels')
self.log.info('Connecting to {}'.format(ws_url))
parameters = {
"headers": KG_HEADERS,
"validate_cert": VALIDATE_KG_CERT,
"connect_timeout": KG_CONNECT_TIMEOUT,
"request_timeout": KG_REQUEST_TIMEOUT
}
if KG_HTTP_USER:
parameters["auth_username"] = KG_HTTP_USER
if KG_HTTP_PASS:
parameters["auth_password"] = KG_HTTP_PASS
if KG_CLIENT_KEY:
parameters["client_key"] = KG_CLIENT_KEY
parameters["client_cert"] = KG_CLIENT_CERT
if KG_CLIENT_CA:
parameters["ca_certs"] = KG_CLIENT_CA
request = HTTPRequest(ws_url, **parameters)
self.ws_future = websocket_connect(request)
self.ws_future.add_done_callback(self._connection_done)
def _connection_done(self, fut):
if not self.ws_future_cancelled: # prevent concurrent.futures._base.CancelledError
self.ws = fut.result()
self.log.debug("Connection is ready: ws: {}".format(self.ws))
else:
self.log.warning(
"Websocket connection has been cancelled via client disconnect before its establishment. "
"Kernel with ID '{}' may not be terminated on Gateway: {}".
format(self.kernel_id, KG_URL))
def _disconnect(self):
if self.ws is not None:
# Close connection
self.ws.close()
elif not self.ws_future.done():
# Cancel pending connection. Since future.cancel() is a noop on tornado, we'll track cancellation locally
self.ws_future.cancel()
self.ws_future_cancelled = True
self.log.debug("_disconnect: ws_future_cancelled: {}".format(
self.ws_future_cancelled))
@gen.coroutine
def _read_messages(self, callback):
"""Read messages from gateway server."""
while True:
message = None
if not self.ws_future_cancelled:
try:
message = yield self.ws.read_message()
except Exception as e:
self.log.error(
"Exception reading message from websocket: {}".format(
e)) # , exc_info=True)
if message is None:
break
callback(
message
) # pass back to notebook client (see self.on_open and WebSocketChannelsHandler.open)
else: # ws cancelled - stop reading
break
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))
def on_message(self, message):
"""Send message to gateway server."""
if self.ws is None:
loop = IOLoop.current()
loop.add_future(self.ws_future,
lambda future: self._write_message(message))
else:
self._write_message(message)
def _write_message(self, message):
"""Send message to gateway server."""
try:
if not self.ws_future_cancelled:
self.ws.write_message(message)
except Exception as e:
self.log.error("Exception writing message to websocket: {}".format(
e)) # , exc_info=True)
def on_close(self):
"""Web socket closed event."""
self._disconnect()
class ProjectEvents(BaseHandler):
response_cancelled = False
polling_clients = set()
PROM_POLLING_CLIENTS.set_function(
lambda: len(ProjectEvents.polling_clients)
)
@api_auth
@PROM_REQUESTS.async_('events')
async def get(self, project_id):
ProjectEvents.polling_clients.add(self.request.remote_ip)
tornado.log.access_log.info(
"started %s %s (%s) (%s)",
self.request.method,
self.request.uri,
self.request.remote_ip,
self.current_user,
)
from_id = int(self.get_query_argument('from'))
project, _ = self.get_project(project_id)
self.project_id = int(project_id)
# Limit over which we won't send update but rather reload the frontend
LIMIT = 20
# Check for immediate update
cmds = (
self.db.query(database.Command)
.filter(database.Command.id > from_id)
.filter(database.Command.project_id == project.id)
.limit(LIMIT)
).all()
if len(cmds) == LIMIT:
return await self.send_json({'reload': True})
if cmds:
# Convert to JSON
cmds_json = [cmd.to_json() for cmd in cmds]
else:
# Wait for an event (which comes in JSON)
self.wait_future = Future()
self.application.observe_project(project.id, self.wait_future)
self.db.expire_all()
# Close DB connection to not overflow the connection pool
self.close_db_connection()
try:
cmds_json = [await self.wait_future]
except asyncio.CancelledError:
return
# Remove 'project_id' from each event
def _change_cmd_json(old):
new = dict(old)
new.pop('project_id')
return new
cmds_json = [_change_cmd_json(cmd) for cmd in cmds_json]
return await self.send_json({'events': cmds_json})
def on_connection_close(self):
self.response_cancelled = True
self.wait_future.cancel()
self.application.unobserve_project(self.project_id, self.wait_future)
def on_finish(self):
super(ProjectEvents, self).on_finish()
ProjectEvents.polling_clients.discard(self.request.remote_ip)
def _log(self):
if not self.response_cancelled:
self.application.log_request(self)
else:
tornado.log.access_log.info(
"aborted %s %s (%s) %.2fms",
self.request.method,
self.request.uri,
self.request.remote_ip,
1000.0 * self.request.request_time(),
)
class KernelGatewayWSClient(LoggingConfigurable):
'''Proxy web socket connection to a kernel gateway.'''
def __init__(self):
self.ws = None
self.ws_future = Future()
@gen.coroutine
def _connect(self, kernel_id):
ws_url = url_path_join(KG_URL.replace('http', 'ws'), '/api/kernels',
url_escape(kernel_id), 'channels')
self.log.info('Connecting to {}'.format(ws_url))
parameters = {
"headers": KG_HEADERS,
"validate_cert": VALIDATE_KG_CERT,
"auth_username": KG_HTTP_USER,
"auth_password": KG_HTTP_PASS,
"connect_timeout": KG_CONNECT_TIMEOUT,
"request_timeout": KG_REQUEST_TIMEOUT
}
if KG_CLIENT_KEY:
parameters["client_key"] = KG_CLIENT_KEY
parameters["client_cert"] = KG_CLIENT_CERT
if KG_CLIENT_CA:
parameters["ca_certs"] = KG_CLIENT_CA
request = HTTPRequest(ws_url, **parameters)
self.ws_future = websocket_connect(request)
self.ws = yield self.ws_future
# TODO: handle connection errors/timeout
def _disconnect(self):
if self.ws is not None:
# Close connection
self.ws.close()
elif not self.ws_future.done():
# Cancel pending connection
self.ws_future.cancel()
@gen.coroutine
def _read_messages(self, callback):
'''Read messages from server.'''
while True:
message = yield self.ws.read_message()
if message is None: break # TODO: handle socket close
callback(message)
def on_open(self, kernel_id, message_callback, **kwargs):
'''Web socket connection open.'''
self._connect(kernel_id)
loop = IOLoop.current()
loop.add_future(self.ws_future,
lambda future: self._read_messages(message_callback))
def on_message(self, message):
'''Send message to server.'''
if self.ws is None:
loop = IOLoop.current()
loop.add_future(self.ws_future,
lambda future: self._write_message(message))
else:
self._write_message(message)
def _write_message(self, message):
'''Send message to server.'''
self.ws.write_message(message)
def on_close(self):
'''Web socket closed event.'''
self._disconnect()
class GatewayWebSocketClient(LoggingConfigurable):
"""Proxy web socket connection to a kernel/enterprise gateway."""
def __init__(self, **kwargs):
super(GatewayWebSocketClient, self).__init__(**kwargs)
self.kernel_id = None
self.ws = None
self.ws_future = Future()
self.ws_future_cancelled = False
@gen.coroutine
def _connect(self, kernel_id):
self.kernel_id = kernel_id
ws_url = url_path_join(
GatewayClient.instance().ws_url,
GatewayClient.instance().kernels_endpoint, url_escape(kernel_id), 'channels'
)
self.log.info('Connecting to {}'.format(ws_url))
kwargs = {}
kwargs = GatewayClient.instance().load_connection_args(**kwargs)
request = HTTPRequest(ws_url, **kwargs)
self.ws_future = websocket_connect(request)
self.ws_future.add_done_callback(self._connection_done)
def _connection_done(self, fut):
if not self.ws_future_cancelled: # prevent concurrent.futures._base.CancelledError
self.ws = fut.result()
self.log.debug("Connection is ready: ws: {}".format(self.ws))
else:
self.log.warning("Websocket connection has been cancelled via client disconnect before its establishment. "
"Kernel with ID '{}' may not be terminated on GatewayClient: {}".
format(self.kernel_id, GatewayClient.instance().url))
def _disconnect(self):
if self.ws is not None:
# Close connection
self.ws.close()
elif not self.ws_future.done():
# Cancel pending connection. Since future.cancel() is a noop on tornado, we'll track cancellation locally
self.ws_future.cancel()
self.ws_future_cancelled = True
self.log.debug("_disconnect: ws_future_cancelled: {}".format(self.ws_future_cancelled))
@gen.coroutine
def _read_messages(self, callback):
"""Read messages from gateway server."""
while True:
message = None
if not self.ws_future_cancelled:
try:
message = yield self.ws.read_message()
except Exception as e:
self.log.error("Exception reading message from websocket: {}".format(e)) # , exc_info=True)
if message is None:
break
callback(message) # pass back to notebook client (see self.on_open and WebSocketChannelsHandler.open)
else: # ws cancelled - stop reading
break
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)
)
def on_message(self, message):
"""Send message to gateway server."""
if self.ws is None:
loop = IOLoop.current()
loop.add_future(
self.ws_future,
lambda future: self._write_message(message)
)
else:
self._write_message(message)
def _write_message(self, message):
"""Send message to gateway server."""
try:
if not self.ws_future_cancelled:
self.ws.write_message(message)
except Exception as e:
self.log.error("Exception writing message to websocket: {}".format(e)) # , exc_info=True)
def on_close(self):
"""Web socket closed event."""
self._disconnect()
class ProjectEvents(BaseHandler):
PROM_API.labels('events').inc(0)
response_cancelled = False
polling_clients = set()
PROM_POLLING_CLIENTS.set_function(
lambda: len(ProjectEvents.polling_clients))
@api_auth
async def get(self, project_id):
PROM_API.labels('events').inc()
ProjectEvents.polling_clients.add(self.request.remote_ip)
tornado.log.access_log.info(
"started %s %s (%s)",
self.request.method,
self.request.uri,
self.request.remote_ip,
)
from_id = int(self.get_query_argument('from'))
project, _ = self.get_project(project_id)
self.project_id = int(project_id)
# Check for immediate update
cmd = (self.db.query(
database.Command).filter(database.Command.id > from_id).filter(
database.Command.project_id == project.id).limit(1)
).one_or_none()
# Wait for an event
if cmd is None:
self.wait_future = Future()
self.application.observe_project(project.id, self.wait_future)
self.db.expire_all()
try:
cmd = await self.wait_future
except asyncio.CancelledError:
return
payload = dict(cmd.payload)
type_ = payload.pop('type', None)
if type_ == 'project_meta':
result = {'project_meta': payload}
elif type_ == 'document_add':
payload['id'] = cmd.document_id
result = {'document_add': [payload]}
elif type_ == 'document_delete':
result = {'document_delete': [cmd.document_id]}
elif type_ == 'highlight_add':
result = {'highlight_add': {cmd.document_id: [payload]}}
elif type_ == 'highlight_delete':
result = {'highlight_delete': {cmd.document_id: [payload['id']]}}
elif type_ == 'tag_add':
result = {
'tag_add': [payload],
}
elif type_ == 'tag_delete':
result = {
'tag_delete': [payload['id']],
}
elif type_ == 'tag_merge':
result = {
'tag_merge': [
{
'src': payload['src'],
'dest': payload['dest']
},
],
}
elif type_ == 'member_add':
result = {
'member_add': [{
'member': payload['member'],
'privileges': payload['privileges']
}]
}
elif type_ == 'member_remove':
result = {'member_remove': [payload['member']]}
else:
raise ValueError("Unknown command type %r" % type_)
if cmd.tag_count_changes is not None:
result['tag_count_changes'] = cmd.tag_count_changes
result['id'] = cmd.id
return self.send_json(result)
def on_connection_close(self):
self.response_cancelled = True
self.wait_future.cancel()
self.application.unobserve_project(self.project_id, self.wait_future)
def on_finish(self):
ProjectEvents.polling_clients.discard(self.request.remote_ip)
def _log(self):
if not self.response_cancelled:
self.application.log_request(self)
else:
tornado.log.access_log.info(
"aborted %s %s (%s) %.2fms",
self.request.method,
self.request.uri,
self.request.remote_ip,
1000.0 * self.request.request_time(),
)
class GatewayWebSocketClient(LoggingConfigurable):
"""Proxy web socket connection to a kernel/enterprise gateway."""
def __init__(self, **kwargs):
super(GatewayWebSocketClient, self).__init__(**kwargs)
self.kernel_id = None
self.ws = None
self.ws_future = Future()
self.disconnected = False
self.retry = 0
async def _connect(self, kernel_id, message_callback):
# websocket is initialized before connection
self.ws = None
self.kernel_id = kernel_id
ws_url = url_path_join(
GatewayClient.instance().ws_url,
GatewayClient.instance().kernels_endpoint,
url_escape(kernel_id),
"channels",
)
self.log.info("Connecting to {}".format(ws_url))
kwargs = {}
kwargs = GatewayClient.instance().load_connection_args(**kwargs)
request = HTTPRequest(ws_url, **kwargs)
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()
self.retry = 0
self.log.debug("Connection is ready: ws: {}".format(self.ws))
else:
self.log.warning(
"Websocket connection has been closed via client disconnect or due to error. "
"Kernel with ID '{}' may not be terminated on GatewayClient: {}"
.format(self.kernel_id,
GatewayClient.instance().url))
def _disconnect(self):
self.disconnected = True
if self.ws is not None:
# Close connection
self.ws.close()
elif not self.ws_future.done():
# Cancel pending connection. Since future.cancel() is a noop on tornado, we'll track cancellation locally
self.ws_future.cancel()
self.log.debug(
"_disconnect: future cancelled, disconnected: {}".format(
self.disconnected))
async def _read_messages(self, callback):
"""Read messages from gateway server."""
while self.ws is not None:
message = None
if not self.disconnected:
try:
message = await self.ws.read_message()
except Exception as e:
self.log.error(
"Exception reading message from websocket: {}".format(
e)) # , exc_info=True)
if message is None:
if not self.disconnected:
self.log.warning(
"Lost connection to Gateway: {}".format(
self.kernel_id))
break
callback(
message
) # pass back to notebook client (see self.on_open and WebSocketChannelsHandler.open)
else: # ws cancelled - stop reading
break
# NOTE(esevan): if websocket is not disconnected by client, try to reconnect.
if not self.disconnected and self.retry < GatewayClient.instance(
).gateway_retry_max:
jitter = random.randint(10, 100) * 0.01
retry_interval = (min(
GatewayClient.instance().gateway_retry_interval *
(2**self.retry),
GatewayClient.instance().gateway_retry_interval_max,
) + jitter)
self.retry += 1
self.log.info(
"Attempting to re-establish the connection to Gateway in %s secs (%s/%s): %s",
retry_interval,
self.retry,
GatewayClient.instance().gateway_retry_max,
self.kernel_id,
)
await asyncio.sleep(retry_interval)
loop = IOLoop.current()
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."""
loop = IOLoop.current()
loop.spawn_callback(self._connect, kernel_id, message_callback)
def on_message(self, message):
"""Send message to gateway server."""
if self.ws is None:
loop = IOLoop.current()
loop.add_future(self.ws_future,
lambda future: self._write_message(message))
else:
self._write_message(message)
def _write_message(self, message):
"""Send message to gateway server."""
try:
if not self.disconnected and self.ws is not None:
self.ws.write_message(message)
except Exception as e:
self.log.error("Exception writing message to websocket: {}".format(
e)) # , exc_info=True)
def on_close(self):
"""Web socket closed event."""
self._disconnect()
