diff --git a/BlocksScreen/lib/moonrakerComm.py b/BlocksScreen/lib/moonrakerComm.py index ba0df898..a69bc001 100644 --- a/BlocksScreen/lib/moonrakerComm.py +++ b/BlocksScreen/lib/moonrakerComm.py @@ -22,7 +22,7 @@ class OneShotTokenError(Exception): """Raised when unable to get oneshot token to connect to a websocket""" def __init__(self, message="Unable to get oneshot token", errors=None) -> None: - super(OneShotTokenError).__init__(message, errors) + super().__init__(message, errors) self.errors = errors self.message = message @@ -31,10 +31,6 @@ class MoonWebSocket(QtCore.QObject, threading.Thread): """MoonWebSocket class object for creating a websocket connection to Moonraker.""" QUERY_KLIPPY_TIMEOUT: int = 2 - connected = False - connecting = False - callback_table = {} - _reconnect_count = 0 max_retries = 3 timeout = 3 @@ -53,9 +49,19 @@ class MoonWebSocket(QtCore.QObject, threading.Thread): ) def __init__(self, parent: QtCore.QObject) -> None: + """Initialize the websocket thread, timers, and Moonraker API helper.""" super().__init__(parent) self.daemon = True + self.connected = False + self.connecting = False + self.disconnected = False + self._reconnect_count = 0 + self.callback_table: dict = {} + + self._state_lock = threading.RLock() + self._request_lock = threading.Lock() + self._host = parent.config.get("host", parser=str, default="localhost") self._port = parent.config.get("port", parser=int, default=7125) @@ -63,10 +69,11 @@ def __init__(self, parent: QtCore.QObject) -> None: self._callback = None self._wst = None self._request_id = 0 - self.request_table = {} + self.request_table: dict = {} + self._klippy_retry_count = 0 self._moonRest = MoonRest(host=self._host, port=self._port) self.api: MoonAPI = MoonAPI(self) - self._retry_timer: RepeatedTimer + self._retry_timer: RepeatedTimer | None = None websocket.setdefaulttimeout(self.timeout) self._intentional_disconnect: bool = False @@ -78,42 +85,57 @@ def __init__(self, parent: QtCore.QObject) -> None: self.klippy_state_signal.connect(self.api.request_printer_info) logger.info("Websocket object initialized") + @property + def moonRest(self) -> MoonRest: + """Returns the current moonrestAPI object""" + return self._moonRest + @QtCore.pyqtSlot(name="retry_wb_conn") def retry_wb_conn(self): """Retry websocket connection""" - if self.connecting is True and self.connected is False: - return False - self._reconnect_count = 0 + with self._state_lock: + if self.connecting is True and self.connected is False: + return False + self._reconnect_count = 0 self.try_connection() @QtCore.pyqtSlot(name="try_connection") def try_connection(self): """Try connecting to websocket""" - if self.connecting: - return - self.connecting = True + with self._state_lock: + self.connecting = True + if self._retry_timer is not None: + self._retry_timer.stopTimer() self._retry_timer = RepeatedTimer(self.timeout, self.reconnect) return self.connect() def reconnect(self) -> bool: """Reconnect to websocket""" - if self.connected: - return True - if self._reconnect_count >= self.max_retries: - self._reconnect_count = 0 + with self._state_lock: + if self.connected: + return True + over_limit = self._reconnect_count >= self.max_retries + + if over_limit: + if self._retry_timer is not None: + self._retry_timer.stopTimer() unable_to_connect_event = WebSocketError( data="Unable to establish connection to Websocket" ) self.connecting_signal[int].emit(0) + with self._state_lock: + self.connecting = False try: instance = QtWidgets.QApplication.instance() if instance is not None: - instance.sendEvent(self.parent(), unable_to_connect_event) + instance.postEvent(self.parent(), unable_to_connect_event) else: raise TypeError("QApplication.instance expected ad non-None value") except Exception as e: logger.error( - f"Error on sending Event {unable_to_connect_event.__class__.__name__} | Error message: {e}" + "Error on sending Event %s | Error message: %s", + unable_to_connect_event.__class__.__name__, + e, ) logger.warning( "Maximum number of connection retries reached, Unable to establish connection with Moonraker" @@ -122,17 +144,19 @@ def reconnect(self) -> bool: def connect(self) -> bool: """Connect to websocket""" - if self.connected: - logger.info("Connection established") - return True - self._reconnect_count += 1 - self.connecting_signal[int].emit(int(self._reconnect_count)) + with self._state_lock: + if self.connected: + logger.info("Connection established") + return True + self._reconnect_count += 1 + _count_snapshot = self._reconnect_count + self.connecting_signal[int].emit(int(_count_snapshot)) logger.debug( - f"Establishing connection to Moonraker...\n Try number {self._reconnect_count}" + "Establishing connection to Moonraker...\n Try number %d", + _count_snapshot, ) - # TODO Handle if i cannot connect to moonraker, request server.info and see if i get a result try: - _oneshot_token = self._moonRest.get_oneshot_token() + _oneshot_token = self.moonRest.get_oneshot_token() if _oneshot_token is None: raise OneShotTokenError("Unable to retrieve oneshot token") except Exception as e: @@ -159,7 +183,7 @@ def connect(self) -> bool: logger.debug(self.ws.url) self._wst.start() except Exception as e: - logger.info(f"Unexpected while starting websocket {self._wst.name}: {e}") + logger.info("Unexpected while starting websocket %s: %s", self._wst.name, e) return False return True @@ -169,17 +193,17 @@ def wb_disconnect(self) -> None: self._intentional_disconnect = True self.ws.close() if self._wst.is_alive(): - self._wst.join() + self._wst.join(timeout=self.timeout + 1) logger.info("Websocket closed") def on_error(self, *args) -> None: """Websocket error callback""" # First argument is ws second is error message - # TODO: Handle error messages _error = args[1] if len(args) == 2 else args[0] - logger.info(f"Websocket error, disconnected: {_error}") - self.connected = False - self.disconnected = True + logger.error("Websocket error, disconnected: %s", _error) + with self._state_lock: + self.connected = False + self.disconnected = True def on_close(self, *args) -> None: """Websocket on close callback @@ -192,7 +216,8 @@ def on_close(self, *args) -> None: return _close_status_code = args[1] if len(args) == 3 else None _close_message = args[2] if len(args) == 3 else None - self.connected = False + with self._state_lock: + self.connected = False self.ws.keep_running = False self.connection_lost[str].emit( f"code: {_close_status_code} | message {_close_message}" @@ -233,9 +258,11 @@ def on_open(self, *args) -> None: TypeError: When QApplication.instance `is` None """ _ws = args[0] if len(args) == 1 else None - self.connecting = False - self.connected = True - self._reconnect_count = 0 + with self._state_lock: + self.connecting = False + self.connected = True + self._klippy_retry_count = 0 + self._reconnect_count = 0 self.evaluate_klippy_status() open_event = WebSocketOpen(data="Connected") try: @@ -245,11 +272,12 @@ def on_open(self, *args) -> None: else: raise TypeError("QApplication.instance expected non None value") except Exception as e: - logger.info(f"Unexpected error opening websocket: {e}") + logger.info("Unexpected error opening websocket: %s", e) self.connected_signal.emit() - self._retry_timer.stopTimer() - logger.info(f"Connection to websocket achieved on {_ws}") + if self._retry_timer is not None: + self._retry_timer.stopTimer() + logger.info("Connection to websocket achieved on %s", _ws) def on_message(self, *args) -> None: """Websocket on message callback @@ -261,22 +289,44 @@ def on_message(self, *args) -> None: args[1] if len(args) == 2 else args[0] ) # First argument is ws second is message - response: dict = json.loads(_message) + try: + response: dict = json.loads(_message) + except json.JSONDecodeError as e: + logger.error("Failed to decode websocket message: %s", e) + return message_event = None - if "id" in response and response["id"] in self.request_table: - _entry = self.request_table.pop(response["id"]) + if "id" in response: + with self._request_lock: + _entry = self.request_table.pop(response["id"], None) + else: + _entry = None + + if _entry is not None: if "server.info" in _entry[0]: - _klippy_state = response["result"]["klippy_state"] + if "error" in response: + return + _result = response.get("result", {}) + _klippy_state = _result.get("klippy_state") + if not _klippy_state: + return if _klippy_state == "ready": self.query_klippy_status_timer.stopTimer() + self._klippy_retry_count = 0 self.api.update_status() # Request update status immediately after klippy ready DEVDEBT - elif _klippy_state == "startup": - # request server.info in 2 seconds - if not self.query_klippy_status_timer.running: - self.query_klippy_status_timer.startTimer() - elif _klippy_state == "disconnected": - if not self.query_klippy_status_timer.running: + elif _klippy_state in ("startup", "disconnected"): + self._klippy_retry_count += 1 + if self._klippy_retry_count >= 30: + self.query_klippy_status_timer.stopTimer() + logger.error( + "Klippy startup sequence timed out after %d retries (state=%s)", + self._klippy_retry_count, + _klippy_state, + ) + elif not self.query_klippy_status_timer.running: self.query_klippy_status_timer.startTimer() + self.klippy_connected_signal.emit( + _result.get("klippy_connected", False) + ) self.klippy_state_signal.emit(_klippy_state) return else: @@ -298,11 +348,16 @@ def on_message(self, *args) -> None: else: message_event = WebSocketMessageReceived( method=str(_entry[0]), - data=response["result"], + data=response.get("result", {}), metadata=_entry, ) elif "method" in response: - if response["method"] in self._KLIPPY_NOTIFY_METHODS: + if str(response["method"]).lower() in ( + "notify_klippy_disconnected", + "notify_klippy_shutdown", + ): + self.klippy_state_signal.emit("disconnected") + self._klippy_retry_count = 0 self.evaluate_klippy_status() message_event = ( WebSocketMessageReceived( # mainly used to pass websocket notifications @@ -311,6 +366,8 @@ def on_message(self, *args) -> None: metadata=None, ) ) + else: + return if message_event is None: return @@ -325,30 +382,40 @@ def on_message(self, *args) -> None: "Unexpected error while creating websocket message event: %s", e ) - def send_request(self, method: str, params: dict = {}, callback=None) -> bool: + def send_request( + self, method: str, params: dict | None = None, callback=None + ) -> bool: """Send a request over the websocket Args: method (str): Websocket method name - params (dict, optional): parameters for the websocket method. Defaults to {}. + params (dict, optional): parameters for the websocket method. Defaults to None. callback (callable, optional): Called with ``response["result"]`` when the response arrives. If None, the response is routed as a ``WebSocketMessageReceived`` event to the parent widget. Defaults to None. Returns: bool: Whether the method finished and a request was sent """ - if not self.connected or self.ws is None: + if params is None: + params = {} + with self._state_lock: + _connected = self.connected + _ws = self.ws + if not _connected or _ws is None: return False - self._request_id += 1 - self.request_table[self._request_id] = [method, params, callback] + with self._request_lock: + self._request_id += 1 + _rid = self._request_id + self.request_table[_rid] = [method, params, callback] + packet = { "jsonrpc": "2.0", "method": method, "params": params, - "id": self._request_id, + "id": _rid, } - self.ws.send(json.dumps(packet)) + _ws.send(json.dumps(packet)) return True @@ -471,11 +538,11 @@ def restart_service(self, service): @QtCore.pyqtSlot(name="firmware_restart") def firmware_restart(self): - """Request Klipper firmware restart + """`POST MoonrakerAPI` /printer/firmware_restart + Firmware restart to Klipper - HTTP_REQUEST: POST /printer/firmware_restart - - JSON_RPC_REQUEST: printer.firmware_restart + Returns: + str: Returns an 'ok' from Moonraker """ return self._ws.send_request(method="printer.firmware_restart") @@ -630,7 +697,7 @@ def move_file(self, source_dir: str, dest_dir: str): isinstance(source_dir, str) is False or isinstance(dest_dir, str) is False or source_dir is None - or dest_dir is False + or dest_dir is None ): return False return self._ws.send_request( @@ -644,7 +711,7 @@ def copy_file(self, source_dir: str, dest_dir: str): isinstance(source_dir, str) is False or isinstance(dest_dir, str) is False or source_dir is None - or dest_dir is False + or dest_dir is None ): return False return self._ws.send_request( @@ -757,7 +824,7 @@ def update_status(self, refresh: bool = False) -> bool: @QtCore.pyqtSlot(name="update-refresh") @QtCore.pyqtSlot(str, name="update-refresh") - def refresh_update_status(self, name: str = None) -> bool: + def refresh_update_status(self, name: str | None = None) -> bool: """Refresh packages state""" if isinstance(name, str): return self._ws.send_request( @@ -814,7 +881,7 @@ def rollback_update(self, name: str): if not isinstance(name, str) or not name: return False return self._ws.send_request( - method="machine,update.rollback", params={"name": name} + method="machine.update.rollback", params={"name": name} ) def history_list(self, limit, start, since, before, order): diff --git a/BlocksScreen/lib/utils/RepeatedTimer.py b/BlocksScreen/lib/utils/RepeatedTimer.py index 42b42aa0..6e059a74 100644 --- a/BlocksScreen/lib/utils/RepeatedTimer.py +++ b/BlocksScreen/lib/utils/RepeatedTimer.py @@ -10,6 +10,7 @@ def __init__( *args, **kwargs, ): + """Initialize a repeating timer that invokes callback every timeout seconds.""" super().__init__(daemon=True) self.name = name self._timeout = timeout @@ -17,6 +18,7 @@ def __init__( self._args = args self._kwargs = kwargs + self._lock = threading.Lock() self.running = False self.timeoutEvent = threading.Event() self.stopEvent = threading.Event() @@ -24,36 +26,42 @@ def __init__( self.startTimer() def _run(self): - self.running = False - self.startTimer() - self.stopEvent.wait() + """Invoke the callback and restart the timer loop, unless stopped.""" + with self._lock: + self.running = False + if self.stopEvent.is_set(): + return if callable(self._function): self._function(*self._args, **self._kwargs) + self.startTimer() def startTimer(self): """Start timer""" - if self.running is False: + with self._lock: + if self.running: + return + self.stopEvent.clear() try: - self._timer = threading.Timer(self._timeout, self._run) - self._timer.daemon = True - self._timer.start() - if not self.stopEvent.is_set(): - self.stopEvent.set() + timer = threading.Timer(self._timeout, self._run) + timer.daemon = True + self._timer = timer + self.running = True except Exception as e: + self.running = False raise Exception( f"RepeatedTimer {self.name} error while starting timer, error: {e}" - ) - finally: - self.running = False - self.running = True + ) from e + # Start outside the lock to avoid holding it during thread creation + timer.start() def stopTimer(self): """Stop timer""" - if self._timer is None: - return - if self.running: - self._timer.cancel() - self._timer.join() + with self._lock: + if self._timer is None or not self.running: + return + timer = self._timer self._timer = None - self.stopEvent.clear() self.running = False + self.stopEvent.set() + timer.cancel() + timer.join()