async def safe_handle(self, websocket, payload): message = None try: message = json.loads(payload) self.log.debug(f"Rpc call <- {message['command']}") response = await self.ws_api(message) # Only respond if we return something from api call if response is not None: log.debug(f"Rpc response -> {message['command']}") # Set success to true automatically (unless it's already set) if "success" not in response: response["success"] = True await websocket.send_str(format_response(message, response)) except Exception as e: tb = traceback.format_exc() self.log.warning(f"Error while handling message: {tb}") if len(e.args) > 0: error = {"success": False, "error": f"{e.args[0]}"} else: error = {"success": False, "error": f"{e}"} if message is None: return await websocket.send_str(format_response(message, error))
async def safe_handle(self, websocket, payload): message = None try: message = json.loads(payload) response = await self.ws_api(message) if response is not None: await websocket.send_str(format_response(message, response)) except BaseException as e: tb = traceback.format_exc() self.log.error(f"Error while handling message: {tb}") error = {"success": False, "error": f"{e}"} if message is None: return await websocket.send_str(format_response(message, error))
async def handle_message(self, websocket, message): """ This function gets called when new message is received via websocket. """ command = message["command"] destination = message["destination"] if destination != "daemon": await self.message_for_service(message) return data = None if "data" in message: data = message["data"] if command == "ping": response = await self.ping() elif command == "start_service": response = await self.start_service(data) elif command == "stop_service": response = await self.stop_service(data) elif command == "is_running": response = await self.is_running(data) elif command == "exit": response = await self.exit() elif command == "register_service": response = await self.register_service(websocket, data) else: response = { "success": False, "error": f"unknown_command {command}" } full_response = format_response(message, response) await websocket.send(full_response)
async def handle_message( self, websocket: WebSocketServerProtocol, message: Dict[str, Any]) -> Tuple[Optional[str], List[Any]]: """ This function gets called when new message is received via websocket. """ command = message["command"] destination = message["destination"] if destination != "daemon": destination = message["destination"] if destination in self.connections: sockets = self.connections[destination] return dict_to_json_str(message), sockets return None, [] data = None if "data" in message: data = message["data"] commands_with_data = [ "start_service", "start_plotting", "stop_plotting", "stop_service", "is_running", "register_service", ] if not isinstance(data, dict) and command in commands_with_data: response = { "success": False, "error": f'{command} requires "data"' } elif command == "ping": response = await ping() elif command == "start_service": response = await self.start_service(cast(Dict[str, Any], data)) elif command == "start_plotting": response = await self.start_plotting(cast(Dict[str, Any], data)) elif command == "stop_plotting": response = await self.stop_plotting(cast(Dict[str, Any], data)) elif command == "stop_service": response = await self.stop_service(cast(Dict[str, Any], data)) elif command == "is_running": response = await self.is_running(cast(Dict[str, Any], data)) elif command == "exit": response = await self.stop() elif command == "register_service": response = await self.register_service(websocket, cast(Dict[str, Any], data)) else: self.log.error(f"UK>> {message}") response = { "success": False, "error": f"unknown_command {command}" } full_response = format_response(message, response) return full_response, [websocket]
async def safe_handle(self, websocket: WebSocketServerProtocol, path: str): service_name = "" try: async for message in websocket: try: decoded = json.loads(message) if "data" not in decoded: decoded["data"] = {} response, sockets_to_use = await self.handle_message( websocket, decoded) except Exception as e: tb = traceback.format_exc() self.log.error(f"Error while handling message: {tb}") error = {"success": False, "error": f"{e}"} response = format_response(decoded, error) sockets_to_use = [] if len(sockets_to_use) > 0: for socket in sockets_to_use: try: await socket.send(response) except Exception as e: tb = traceback.format_exc() self.log.error( f"Unexpected exception trying to send to websocket: {e} {tb}" ) self.remove_connection(socket) await socket.close() except Exception as e: tb = traceback.format_exc() service_name = "Unknown" if websocket in self.remote_address_map: service_name = self.remote_address_map[websocket] if isinstance(e, ConnectionClosedOK): self.log.info( f"ConnectionClosedOk. Closing websocket with {service_name} {e}" ) elif isinstance(e, WebSocketException): self.log.info( f"Websocket exception. Closing websocket with {service_name} {e} {tb}" ) else: self.log.error(f"Unexpected exception in websocket: {e} {tb}") finally: self.remove_connection(websocket) await websocket.close()
async def handle_message(self, websocket, message) -> Tuple[Optional[str], List[Any]]: """ This function gets called when new message is received via websocket. """ command = message["command"] destination = message["destination"] if destination != "daemon": destination = message["destination"] if destination in self.connections: sockets = self.connections[destination] return dict_to_json_str(message), sockets return None, [] data = None if "data" in message: data = message["data"] if command == "ping": response = await ping() elif command == "start_service": response = await self.start_service(data) elif command == "start_plotting": response = await self.start_plotting(data) elif command == "stop_plotting": response = await self.stop_plotting(data) elif command == "stop_service": response = await self.stop_service(data) elif command == "is_running": response = await self.is_running(data) elif command == "exit": response = await self.stop() elif command == "register_service": response = await self.register_service(websocket, data) else: self.log.error(f"UK>> {message}") response = { "success": False, "error": f"unknown_command {command}" } full_response = format_response(message, response) return full_response, [websocket]
async def safe_handle(self, websocket, path): async for message in websocket: try: decoded = json.loads(message) # self.log.info(f"Message received: {decoded}") await self.handle_message(websocket, decoded) except (BaseException, websockets.exceptions.ConnectionClosed) as e: if isinstance(e, websockets.exceptions.ConnectionClosed): service_name = self.remote_address_map[websocket.remote_address[1]] self.log.info( f"ConnectionClosed. Closing websocket with {service_name}" ) if service_name in self.connections: self.connections.pop(service_name) await websocket.close() else: tb = traceback.format_exc() self.log.error(f"Error while handling message: {tb}") error = {"success": False, "error": f"{e}"} await websocket.send(format_response(message, error))
async def safe_handle(self, websocket, path): service_name = "" try: async for message in websocket: try: decoded = json.loads(message) response, socket_to_use = await self.handle_message( websocket, decoded) except Exception as e: tb = traceback.format_exc() self.log.error(f"Error while handling message: {tb}") error = {"success": False, "error": f"{e}"} response = format_response(message, error) if socket_to_use is not None: await socket_to_use.send(response) except websockets.exceptions.ConnectionClosedOK as e: if websocket.remote_address[1] in self.remote_address_map: service_name = self.remote_address_map[ websocket.remote_address[1]] self.log.info( f"ConnectionClosedOk. Closing websocket with {service_name} {e}" ) except websockets.exceptions.WebSocketException as e: if websocket.remote_address[1] in self.remote_address_map: service_name = self.remote_address_map[ websocket.remote_address[1]] self.log.info( f"Websocket exception. Closing websocket with {service_name} {e}" ) except Exception as e: tb = traceback.format_exc() self.log.error(f"Unexpected exception in websocket: {e} {tb}") finally: if websocket.remote_address[1] in self.remote_address_map: service_name = self.remote_address_map[ websocket.remote_address[1]] if service_name in self.connections: self.connections.pop(service_name) await websocket.close()