def __init__(self, url=None, ssl_opts=None, connect_immediately=True, max_wait=2): self.ws = None self.url = url or 'ws://127.0.0.1:{}/wsrpc'.format(config['cherrypy']['server.socket_port']) self._lock = RLock() self._callbacks = {} self._counter = count() self.ssl_opts = ssl_opts self._reconnect_attempts = 0 self._last_poll, self._last_reconnect_attempt = None, None self._dispatcher = Caller(self._dispatch, threads=1) self._checker = DaemonTask(self._check, interval=1) if connect_immediately: self.connect(max_wait=max_wait)
pass class WebSocketChecker(WebSocketTool): def __init__(self): cherrypy.Tool.__init__(self, 'before_handler', self.upgrade) def upgrade(self, **kwargs): try: kwargs['handler_cls'].check_authentication() except: raise cherrypy.HTTPError( 401, 'You must be logged in to establish a websocket connection.') else: return WebSocketTool.upgrade(self, **kwargs) cherrypy.tools.websockets = WebSocketChecker() websocket_plugin = WebSocketPlugin(cherrypy.engine) if hasattr(WebSocketPlugin.start, '__func__'): WebSocketPlugin.start.__func__.priority = 66 else: WebSocketPlugin.start.priority = 66 websocket_plugin.subscribe() broadcaster = Caller(WebSocketDispatcher.broadcast) responder = Caller(WebSocketDispatcher.handle_message, threads=config['ws.thread_pool'])
class WebSocket(object): """ Utility class for making websocket connections. This improves on the ws4py websocket client classes mainly by adding several features: - automatically detecting dead connections and re-connecting - utility methods for making synchronous rpc calls and for making asynchronous subscription calls with callbacks - adding locking to make sending messages thread-safe """ poll_method = 'sideboard.poll' WebSocketDispatcher = _WebSocketClientDispatcher def __init__(self, url=None, ssl_opts=None, connect_immediately=True, max_wait=2): self.ws = None self.url = url or 'ws://127.0.0.1:{}/wsrpc'.format(config['cherrypy']['server.socket_port']) self._lock = RLock() self._callbacks = {} self._counter = count() self.ssl_opts = ssl_opts self._reconnect_attempts = 0 self._last_poll, self._last_reconnect_attempt = None, None self._dispatcher = Caller(self._dispatch, threads=1) self._checker = DaemonTask(self._check, interval=1) if connect_immediately: self.connect(max_wait=max_wait) def __enter__(self): return self def __exit__(self, exc_type, exc_value, traceback): self.close() @property def _should_reconnect(self): interval = min(config['ws.reconnect_interval'], 2 ** self._reconnect_attempts) cutoff = datetime.now() - timedelta(seconds=interval) return not self.connected and (self._reconnect_attempts == 0 or self._last_reconnect_attempt < cutoff) @property def _should_poll(self): cutoff = datetime.now() - timedelta(seconds=config['ws.poll_interval']) return self.connected and (self._last_poll is None or self._last_poll < cutoff) def _check(self): if self._should_reconnect: self._reconnect() if self._should_poll: self._poll() def _poll(self): assert self.ws and self.ws.connected, 'cannot poll while websocket is not connected' try: self.call(self.poll_method) except: log.warning('no poll response received from {!r}, closing connection, will attempt to reconnect', self.url, exc_info=True) self.ws.close() else: self._last_poll = datetime.now() def _refire_subscriptions(self): try: for cb in self._callbacks.values(): if 'client' in cb: self._send(method=cb['method'], params=cb['params'], client=cb['client']) except: pass # self._send() already closes and logs on error def _reconnect(self): with self._lock: assert not self.connected, 'connection is still active' try: self.ws = self.WebSocketDispatcher(self._dispatcher, self.url, ssl_opts=self.ssl_opts) self.ws.connect() except Exception as e: log.warn('failed to connect to {}: {}', self.url, str(e)) self._last_reconnect_attempt = datetime.now() self._reconnect_attempts += 1 else: self._reconnect_attempts = 0 self._refire_subscriptions() def _next_id(self, prefix): return '{}-{}'.format(prefix, next(self._counter)) def _send(self, **kwargs): log.debug('sending {}', kwargs) with self._lock: assert self.connected, 'tried to send data on closed websocket {!r}'.format(self.url) try: return self.ws.send(kwargs) except: log.warn('failed to send {!r} on {!r}, closing websocket and will attempt to reconnect', kwargs, self.url) self.ws.close() raise def _dispatch(self, message): log.debug('dispatching {}', message) try: assert isinstance(message, Mapping), 'incoming message is not a dictionary' assert 'client' in message or 'callback' in message, 'no callback or client in message {}'.format(message) id = message.get('client') or message.get('callback') assert id in self._callbacks, 'unknown dispatchee {}'.format(id) except AssertionError: self.fallback(message) else: if 'error' in message: self._callbacks[id]['errback'](message['error']) else: self._callbacks[id]['callback'](message.get('data')) def fallback(self, message): """ Handler method which is called for incoming websocket messages which aren't valid responses to an outstanding call or subscription. By default this just logs an error message. You can override this by subclassing this class, or just by assigning a hander method, e.g. >>> ws = WebSocket() >>> ws.fallback = some_handler_function >>> ws.connect() """ _, exc, _ = sys.exc_info() log.error('no callback registered for message {!r}, message ignored: {}', message, exc) @property def connected(self): """boolean indicating whether or not this connection is currently active""" return bool(self.ws) and self.ws.connected def connect(self, max_wait=0): """ Start the background threads which connect this websocket and handle RPC dispatching. This method is safe to call even if the websocket is already connected. You may optionally pass a max_wait parameter if you want to wait for up to that amount of time for the connection to go through; if that amount of time elapses without successfully connecting, a warning message is logged. """ self._checker.start() self._dispatcher.start() for i in range(10 * max_wait): if not self.connected: stopped.wait(0.1) else: break else: if max_wait: log.warn('websocket {!r} not connected after {} seconds', self.url, max_wait) def close(self): """ Closes the underlying websocket connection and stops background tasks. This method is always safe to call; exceptions will be swallowed and logged, and calling close on an already-closed websocket is a no-op. """ self._checker.stop() self._dispatcher.stop() if self.ws: self.ws.close() def subscribe(self, callback, method, *args, **kwargs): """ Send a websocket request which you expect to subscribe you to a channel with a callback which will be called every time there is new data, and return the client id which uniquely identifies this subscription. Callback may be either a function or a dictionary in the form { 'callback': <function>, 'errback': <function> } Both callback and errback take a single argument; for callback, this is the return value of the method, for errback it is the error message returning. If no errback is specified, we will log errors at the ERROR level and do nothing further. The positional and keyword arguments passed to this function will be used as the arguments to the remote method. """ client = self._next_id('client') if isinstance(callback, Mapping): assert 'callback' in callback and 'errback' in callback, 'callback and errback are required' client = callback.setdefault('client', client) self._callbacks[client] = callback else: self._callbacks[client] = { 'client': client, 'callback': callback, 'errback': lambda result: log.error('{}(*{}, **{}) returned an error: {!r}', method, args, kwargs, result) } self._callbacks[client].update({ 'method': method, 'params': args or kwargs }) try: self._send(method=method, params=args or kwargs, client=client) except: log.warn('initial subscription to {} at {!r} failed, will retry on reconnect', method, self.url) return client def unsubscribe(self, client): """ Cancel the websocket subscription identified by the specified client id. This id is returned from the subscribe() method, e.g. >>> client = ws.subscribe(some_callback, 'foo.some_function') >>> ws.unsubscribe(client) """ self._callbacks.pop(client, None) try: self._send(action='unsubscribe', client=client) except: pass def call(self, method, *args, **kwargs): """ Send a websocket rpc method call, then wait for and return the eventual response, or raise an exception if we get back an error. This method will raise an AssertionError after 10 seconds if no response of any kind was received. The positional and keyword arguments to this method are used as the arguments to the rpc function call. """ result, error = [], [] callback = self._next_id('callback') self._callbacks[callback] = { 'callback': result.append, 'errback': error.append } try: self._send(method=method, params=args or kwargs, callback=callback) except: self._callbacks.pop(callback, None) raise for i in range(10 * config['ws.call_timeout']): stopped.wait(0.1) if stopped.is_set() or result or error: break self._callbacks.pop(callback, None) assert not stopped.is_set(), 'websocket closed before response was received' assert result, error[0] if error else 'no response received for 10 seconds' return result[0] def make_caller(self, method): """ Returns a function which calls the specified method; useful for creating callbacks, e.g. >>> authenticate = ws.make_caller('auth.authenticate') >>> authenticate('username', 'password') True """ originating_ws = sideboard.lib.threadlocal.get('websocket') client = sideboard.lib.threadlocal.get('message', {}).get('client') if client and originating_ws: return _Subscriber(client=client, src_ws=originating_ws, dest_ws=self, method=method) else: return lambda *args, **kwargs: self.call(method, *args, **kwargs)
class WebSocket(object): """ Utility class for making websocket connections. This improves on the ws4py websocket client classes mainly by adding several features: - automatically detecting dead connections and re-connecting - utility methods for making synchronous rpc calls and for making asynchronous subscription calls with callbacks - adding locking to make sending messages thread-safe """ poll_method = 'sideboard.poll' WebSocketDispatcher = _WebSocketClientDispatcher def __init__(self, url=None, ssl_opts=None, connect_immediately=True, max_wait=2): self.ws = None self.url = url or 'ws://127.0.0.1:{}/wsrpc'.format( config['cherrypy']['server.socket_port']) self._lock = RLock() self._callbacks = {} self._counter = count() self.ssl_opts = ssl_opts self._reconnect_attempts = 0 self._last_poll, self._last_reconnect_attempt = None, None self._dispatcher = Caller(self._dispatch, threads=1) self._checker = DaemonTask(self._check, interval=1) if connect_immediately: self.connect(max_wait=max_wait) def __enter__(self): return self def __exit__(self, exc_type, exc_value, traceback): self.close() def preprocess(self, method, params): """ Each message we send has its parameters passed to this function and the actual parameters sent are whatever this function returns. By default this just returns the message unmodified, but plugins can override this to add whatever logic is needed. We pass the method name in its full "service.method" form in case the logic depends on the service being invoked. """ return params @property def _should_reconnect(self): interval = min(config['ws.reconnect_interval'], 2**self._reconnect_attempts) cutoff = datetime.now() - timedelta(seconds=interval) return not self.connected and (self._reconnect_attempts == 0 or self._last_reconnect_attempt < cutoff) @property def _should_poll(self): cutoff = datetime.now() - timedelta(seconds=config['ws.poll_interval']) return self.connected and (self._last_poll is None or self._last_poll < cutoff) def _check(self): if self._should_reconnect: self._reconnect() if self._should_poll: self._poll() def _poll(self): assert self.ws and self.ws.connected, 'cannot poll while websocket is not connected' try: self.call(self.poll_method) except: log.warning( 'no poll response received from {!r}, closing connection, will attempt to reconnect', self.url, exc_info=True) self.ws.close() else: self._last_poll = datetime.now() def _refire_subscriptions(self): try: for cb in self._callbacks.values(): if 'client' in cb: params = cb['paramback']( ) if 'paramback' in cb else cb['params'] self._send(method=cb['method'], params=params, client=cb['client']) except: pass # self._send() already closes and logs on error def _reconnect(self): with self._lock: assert not self.connected, 'connection is still active' try: self.ws = self.WebSocketDispatcher(self._dispatcher, self.url, ssl_opts=self.ssl_opts) self.ws.connect() except Exception as e: log.warn('failed to connect to {}: {}', self.url, str(e)) self._last_reconnect_attempt = datetime.now() self._reconnect_attempts += 1 else: self._reconnect_attempts = 0 self._refire_subscriptions() def _next_id(self, prefix): return '{}-{}'.format(prefix, next(self._counter)) def _send(self, **kwargs): log.debug('sending {}', kwargs) with self._lock: assert self.connected, 'tried to send data on closed websocket {!r}'.format( self.url) try: return self.ws.send(kwargs) except: log.warn( 'failed to send {!r} on {!r}, closing websocket and will attempt to reconnect', kwargs, self.url) self.ws.close() raise def _dispatch(self, message): log.debug('dispatching {}', message) try: assert isinstance(message, Mapping), 'incoming message is not a dictionary' assert 'client' in message or 'callback' in message, 'no callback or client in message {}'.format( message) id = message.get('client') or message.get('callback') assert id in self._callbacks, 'unknown dispatchee {}'.format(id) except AssertionError: self.fallback(message) else: if 'error' in message: self._callbacks[id]['errback'](message['error']) else: self._callbacks[id]['callback'](message.get('data')) def fallback(self, message): """ Handler method which is called for incoming websocket messages which aren't valid responses to an outstanding call or subscription. By default this just logs an error message. You can override this by subclassing this class, or just by assigning a hander method, e.g. >>> ws = WebSocket() >>> ws.fallback = some_handler_function >>> ws.connect() """ _, exc, _ = sys.exc_info() log.error( 'no callback registered for message {!r}, message ignored: {}', message, exc) @property def connected(self): """boolean indicating whether or not this connection is currently active""" return bool(self.ws) and self.ws.connected def connect(self, max_wait=0): """ Start the background threads which connect this websocket and handle RPC dispatching. This method is safe to call even if the websocket is already connected. You may optionally pass a max_wait parameter if you want to wait for up to that amount of time for the connection to go through; if that amount of time elapses without successfully connecting, a warning message is logged. """ self._checker.start() self._dispatcher.start() for i in range(10 * max_wait): if not self.connected: stopped.wait(0.1) else: break else: if max_wait: log.warn('websocket {!r} not connected after {} seconds', self.url, max_wait) def close(self): """ Closes the underlying websocket connection and stops background tasks. This method is always safe to call; exceptions will be swallowed and logged, and calling close on an already-closed websocket is a no-op. """ self._checker.stop() self._dispatcher.stop() if self.ws: self.ws.close() def subscribe(self, callback, method, *args, **kwargs): """ Send a websocket request which you expect to subscribe you to a channel with a callback which will be called every time there is new data, and return the client id which uniquely identifies this subscription. Callback may be either a function or a dictionary in the form { 'callback': <function>, 'errback': <function>, # optional 'paramback: <function>, # optional 'client': <string> # optional } Both callback and errback take a single argument; for callback, this is the return value of the method, for errback it is the error message returning. If no errback is specified, we will log errors at the ERROR level and do nothing further. The paramback function exists for subscriptions where we might want to pass different parameters every time we reconnect. This might be used for e.g. time-based parameters. This function takes no arguments and returns the parameters which should be passed every time we connect and fire (or re-fire) all of our subscriptions. The client id is automatically generated if omitted, and you should not set this yourself unless you really know what you're doing. The positional and keyword arguments passed to this function will be used as the arguments to the remote method, unless paramback is passed, in which case that will be used to generate the params, and args/kwargs will be ignored. """ client = self._next_id('client') if isinstance(callback, Mapping): assert 'callback' in callback, 'callback is required' client = callback.setdefault('client', client) self._callbacks[client] = callback else: self._callbacks[client] = {'client': client, 'callback': callback} paramback = self._callbacks[client].get('paramback') params = self.preprocess( method, paramback() if paramback else (args or kwargs)) self._callbacks[client].setdefault( 'errback', lambda result: log.error('{}(*{}, **{}) returned an error: {!r}', method, args, kwargs, result)) self._callbacks[client].update({'method': method, 'params': params}) try: self._send(method=method, params=params, client=client) except: log.warn( 'initial subscription to {} at {!r} failed, will retry on reconnect', method, self.url) return client def unsubscribe(self, client): """ Cancel the websocket subscription identified by the specified client id. This id is returned from the subscribe() method, e.g. >>> client = ws.subscribe(some_callback, 'foo.some_function') >>> ws.unsubscribe(client) """ self._callbacks.pop(client, None) try: self._send(action='unsubscribe', client=client) except: pass def call(self, method, *args, **kwargs): """ Send a websocket rpc method call, then wait for and return the eventual response, or raise an exception if we get back an error. This method will raise an AssertionError after 10 seconds if no response of any kind was received. The positional and keyword arguments to this method are used as the arguments to the rpc function call. """ finished = Event() result, error = [], [] callback = self._next_id('callback') self._callbacks[callback] = { 'callback': lambda response: (result.append(response), finished.set()), 'errback': lambda response: (error.append(response), finished.set()) } params = self.preprocess(method, args or kwargs) try: self._send(method=method, params=params, callback=callback) except: self._callbacks.pop(callback, None) raise wait_until = datetime.now() + timedelta( seconds=config['ws.call_timeout']) while datetime.now() < wait_until: finished.wait(0.1) if stopped.is_set() or result or error: break self._callbacks.pop(callback, None) assert not stopped.is_set( ), 'websocket closed before response was received' assert result, error[ 0] if error else 'no response received for 10 seconds' return result[0] def make_caller(self, method): """ Returns a function which calls the specified method; useful for creating callbacks, e.g. >>> authenticate = ws.make_caller('auth.authenticate') >>> authenticate('username', 'password') True Sideboard supports "passthrough subscriptions", e.g. -> a browser makes a subscription for the "foo.bar" method -> the server has "foo" registered as a remote service -> the server creates its own subscription to "foo.bar" on the remote service and passes all results back to the client as they arrive This method implements that by checking whether it was called from a thread with an active websocket as part of a subscription request. If so then in addition to returning a callable, it also registers the new subscription with the client websocket so it can be cleaned up when the client websocket closes and/or when its subscription is canceled. """ client = sideboard.lib.threadlocal.get_client() originating_ws = sideboard.lib.threadlocal.get('websocket') if client and originating_ws: sub = originating_ws.passthru_subscriptions.get(client) if sub: sub.method = method else: sub = _Subscriber(method=method, src_client=client, dst_client=self._next_id('client'), src_ws=originating_ws, dest_ws=self) originating_ws.passthru_subscriptions[client] = sub return sub else: return lambda *args, **kwargs: self.call(method, *args, **kwargs)