websocket: refactor how the connection is handled

ks_includes/KlippyWebsocket.py

@@ -30,6 +30,7 @@ api = {
 class KlippyWebsocket(threading.Thread):
     _req_id = 0
     connected = False
+    connecting = True
     callback_table = {}
     reconnect_count = 0
     max_retries = 4
@@ -59,35 +60,23 @@ class KlippyWebsocket(threading.Thread):
 
     def retry(self):
         self.reconnect_count = 0
+        self.connecting = True
         self.initial_connect()
 
     def initial_connect(self):
         # Enable a timeout so that way if moonraker is not running, it will attempt to reconnect
         self.connect()
+        if self.connecting:
             GLib.timeout_add_seconds(6, self.reconnect)
 
     def connect(self):
-        def ws_on_close(ws, a=None, b=None):
+        if self.connected:
-            self.on_close(ws)
+            return
-
-        def ws_on_error(ws, msg):
-            self.on_error(ws, msg)
-
-        def ws_on_message(ws, msg):
-            self.on_message(ws, msg)
-
-        def ws_on_open(ws):
-            self.on_open(ws)
-
         self.reconnect_count += 1
         try:
             state = self._screen.apiclient.get_server_info()
             if state is False:
-                if self.reconnect_count > self.max_retries:
+                if self.reconnect_count > 2:
-                    self._screen.printer_initializing(
-                        _("Cannot connect to Moonraker")
-                        + f'\n\n{self._url}')
-                elif self.reconnect_count > 2:
                     self._screen.printer_initializing(
                         _("Cannot connect to Moonraker")
                         + f'\n\n{self._url}\n\n'
@@ -102,8 +91,9 @@ class KlippyWebsocket(threading.Thread):
 
         self.ws_url = f"{self.ws_proto}://{self._url}/websocket?token={token}"
         self.ws = websocket.WebSocketApp(
-            self.ws_url, on_close=ws_on_close, on_error=ws_on_error, on_message=ws_on_message, on_open=ws_on_open)
+            self.ws_url,
-
+            on_close=self.on_close, on_error=self.on_error, on_message=self.on_message, on_open=self.on_open
+        )
         self._wst = threading.Thread(target=self.ws.run_forever, daemon=True)
         try:
             logging.debug("Starting websocket thread")
@@ -114,15 +104,11 @@ class KlippyWebsocket(threading.Thread):
 
     def close(self):
         self.closing = True
+        self.connecting = False
         if self.ws is not None:
+            logging.debug("Closing websocket")
             self.ws.close()
 
-    def is_connected(self):
-        return self.connected
-
-    def is_connecting(self):
-        return self.reconnect_count > self.max_retries
-
     def on_message(self, ws, message):
         response = json.loads(message)
         if "id" in response and response['id'] in self.callback_table:
@@ -140,10 +126,10 @@ class KlippyWebsocket(threading.Thread):
         return
 
     def send_method(self, method, params=None, callback=None, *args):
+        if not self.connected:
+            return False
         if params is None:
             params = {}
-        if self.is_connected() is False:
-            return False
 
         self._req_id += 1
         if callback is not None:
@@ -160,36 +146,33 @@ class KlippyWebsocket(threading.Thread):
 
     def on_open(self, ws):
         logging.info("Moonraker Websocket Open")
-        logging.info(f"Self.connected = {self.is_connected()}")
         self.connected = True
         self.reconnect_count = 0
         if "on_connect" in self._callback:
             GLib.idle_add(self._callback['on_connect'])
 
-    def on_close(self, ws):
+    def on_close(self, ws, status, message):
-        if self.is_connected() is False:
+        if message is not None:
+            logging.info(f"{status} {message}")
+        if not self.connected:
             logging.debug("Connection already closed")
             return
 
-        if self.closing is True:
+        if "on_close" in self._callback:
-            logging.debug("Closing websocket")
+            GLib.idle_add(self._callback['on_close'], "Lost Connection to Moonraker")
-            self.ws.keep_running = False
-            self.close()
-            self.closing = False
-            return
 
         logging.info("Moonraker Websocket Closed")
         self.connected = False
 
-        if "on_close" in self._callback:
-            GLib.idle_add(self._callback['on_close'], "Lost Connection to Moonraker")
-        self.retry()
-
     def reconnect(self):
-        if self.is_connected():
+        if self.connected:
             return False
         if self.reconnect_count > self.max_retries:
             logging.debug("Stopping reconnections")
+            self.connecting = False
+            self._screen.printer_initializing(
+                _("Cannot connect to Moonraker")
+                + f'\n\n{self._url}')
             return False
         logging.debug("Attempting to reconnect")
         self.connect()

ks_includes/printer.py

@@ -37,9 +37,6 @@ class Printer:
         self.output_pin_count = None
         self.store_timeout = None
         self.tempstore = None
-        if self.store_timeout is not None:
-            GLib.source_remove(self.store_timeout)
-            self.store_timeout = None
 
     def reinit(self, printer_info, data):
         self.config = data['configfile']['config']
@@ -87,15 +84,6 @@ class Printer:
                     self.fancount += 1
             if x.startswith('output_pin ') and not x.split()[1].startswith("_"):
                 self.output_pin_count += 1
-            if x.startswith('bed_mesh '):
-                r = self.config[x]
-                r['x_count'] = int(r['x_count'])
-                r['y_count'] = int(r['y_count'])
-                r['max_x'] = float(r['max_x'])
-                r['min_x'] = float(r['min_x'])
-                r['max_y'] = float(r['max_y'])
-                r['min_y'] = float(r['min_y'])
-                r['points'] = [[float(j.strip()) for j in i.split(",")] for i in r['points'].strip().split("\n")]
         self.process_update(data)
 
         logging.info(f"Klipper version: {printer_info['software_version']}")
@@ -320,6 +308,8 @@ class Printer:
         self.devices[dev][stat] = value
 
     def _update_temp_store(self):
+        if self.tempstore is None:
+            return False
         for device in self.tempstore:
             for x in self.tempstore[device]:
                 if len(self.tempstore[device][x]) >= 1200:

panels/menu.py

@@ -32,7 +32,7 @@ class MenuPanel(ScreenPanel):
     def activate(self):
         self.j2_data = self._printer.get_printer_status_data()
         self.j2_data.update({
-            'moonraker_connected': self._screen._ws.is_connected()
+            'moonraker_connected': self._screen._ws.connected
         })
         if self._screen.vertical_mode:
             self.arrangeMenuItems(self.items, 3)
@@ -96,8 +96,8 @@ class MenuPanel(ScreenPanel):
 
     def evaluate_enable(self, enable):
         if enable == "{{ moonraker_connected }}":
-            logging.info("moonraker is_connected %s", self._screen._ws.is_connected())
+            logging.info(f"moonraker connected {self._screen._ws.connected}")
-            return self._screen._ws.is_connected()
+            return self._screen._ws.connected
 
         self.j2_data = self._printer.get_printer_status_data()
         try:

panels/printer_select.py

@@ -47,4 +47,5 @@ class PrinterSelect(ScreenPanel):
 
     def activate(self):
         self._screen.base_panel.action_bar.hide()
+        self._screen._ws.connecting = False
         GLib.timeout_add(100, self._screen.base_panel.action_bar.hide)

panels/splash_screen.py

@@ -72,7 +72,7 @@ class SplashScreenPanel(ScreenPanel):
 
         self.clear_action_bar()
         printer = self._screen.connected_printer
-        if printer is not None and self._screen._ws.is_connected():
+        if printer is not None and self._screen._ws.connected:
             printer_cfg = self._config.get_printer_config(printer)
             if printer_cfg is not None:
                 power_devices = printer_cfg.get("power_devices", "")
@@ -81,14 +81,14 @@ class SplashScreenPanel(ScreenPanel):
                     logging.info(f"Associated power devices: {power_devices}")
                     self.add_power_button(self._screen.search_power_devices(power_devices))
 
-        if self._screen._ws and self._screen._ws.is_connected():
+        if self._screen._ws and self._screen._ws.connected:
             self.labels['actions'].add(self.labels['restart'])
             self.labels['actions'].add(self.labels['firmware_restart'])
         else:
             self.labels['actions'].add(self.labels['restart_system'])
             self.labels['actions'].add(self.labels['shutdown'])
         self.labels['actions'].add(self.labels['menu'])
-        if self._screen._ws and self._screen._ws.is_connecting():
+        if self._screen._ws and not self._screen._ws.connecting:
             self.labels['actions'].add(self.labels['retry'])
         self.labels['actions'].show_all()
 
@@ -123,8 +123,7 @@ class SplashScreenPanel(ScreenPanel):
         self._screen._ws.klippy.restart()
 
     def shutdown(self, widget):
-
+        if self._screen._ws.connected:
-        if self._screen._ws.is_connected():
             self._screen._confirm_send_action(widget,
                                               _("Are you sure you wish to shutdown the system?"),
                                               "machine.shutdown")
@@ -134,7 +133,7 @@ class SplashScreenPanel(ScreenPanel):
 
     def restart_system(self, widget):
 
-        if self._screen._ws.is_connected():
+        if self._screen._ws.connected:
             self._screen._confirm_send_action(widget,
                                               _("Are you sure you wish to reboot the system?"),
                                               "machine.reboot")

screen.py

@@ -148,16 +148,22 @@ class KlipperScreen(Gtk.Window):
             pname = list(printers[0])[0]
             self.connect_printer(pname)
         else:
+            self.base_panel.show_printer_select(True)
             self.show_printer_select()
 
     def connect_printer(self, name):
+        self.connecting_to_printer = name
+        if self._ws is not None and self._ws.connected:
+            self._ws.close()
+            return
+
         data = {
             "moonraker_host": "127.0.0.1",
             "moonraker_port": "7125",
             "moonraker_api_key": False
         }
 
-        self.connecting_to_printer = name
+        logging.info(f"Connecting to printer: {name}")
         for printer in self._config.get_printers():
             pname = list(printer)[0]
             if pname != name:
@@ -165,19 +171,7 @@ class KlipperScreen(Gtk.Window):
             data = printer[pname]
             break
 
-        if self._ws is not None:
-            self._ws.close()
-        self.connecting = True
-
-        logging.info(f"Connecting to printer: {name}")
         self.apiclient = KlippyRest(data["moonraker_host"], data["moonraker_port"], data["moonraker_api_key"])
-        # Cleanup
-        if self.files is not None:
-            self.files.reset()
-            self.files = None
-        if self.printer is not None:
-            self.printer.reset()
-            self.printer = None
         self.printer = Printer(self.state_execute)
         self.printer.state_callbacks = {
             "disconnected": self.state_disconnected,
@@ -189,14 +183,13 @@ class KlipperScreen(Gtk.Window):
             "shutdown": self.state_shutdown
         }
         self._remove_all_panels()
-        self.base_panel.show_printer_select(True)
         self.printer_initializing(_("Connecting to %s") % name)
 
         self._ws = KlippyWebsocket(self,
                                    {
                                        "on_connect": self.init_printer,
                                        "on_message": self._websocket_callback,
-                                       "on_close": self.state_disconnected
+                                       "on_close": self.websocket_disconnected
                                    },
                                    data["moonraker_host"],
                                    data["moonraker_port"]
@@ -204,10 +197,6 @@ class KlipperScreen(Gtk.Window):
 
         self.files = KlippyFiles(self)
         self._ws.initial_connect()
-        self.connecting = False
-        self.connected_printer = name
-        self.base_panel.set_ks_printer_cfg(name)
-        logging.debug(f"Connected to printer: {name}")
 
     def ws_subscribe(self):
         requested_updates = {
@@ -662,12 +651,20 @@ class KlipperScreen(Gtk.Window):
             self.init_printer()
             callback()
 
-    def state_disconnected(self, msg=None):
+    def websocket_disconnected(self, msg):
-        logging.debug("### Going to disconnected")
-        self.printer.state = "disconnected"
-        self.close_screensaver()
-        msg = _("Klipper has disconnected") if msg is None else msg
         self.printer_initializing(msg)
+        self.connecting = True
+        self.connected_printer = None
+        self.files.reset()
+        self.files = None
+        self.printer.reset()
+        self.printer = None
+        self.connect_printer(self.connecting_to_printer)
+
+    def state_disconnected(self):
+        logging.debug("### Going to disconnected")
+        self.close_screensaver()
+        self.printer_initializing(_("Klipper has disconnected"))
 
     def state_error(self):
         self.close_screensaver()
@@ -803,7 +800,7 @@ class KlipperScreen(Gtk.Window):
 
     def printer_initializing(self, msg):
         self.close_popup_message()
-        if 'splash_screen' not in self._cur_panels:
+        if 'splash_screen' not in self.panels or 'printer_select' not in self._cur_panels:
             self.show_panel('splash_screen', "splash_screen", None, 2)
         self.panels['splash_screen'].update_text(msg)
 
@@ -839,6 +836,10 @@ class KlipperScreen(Gtk.Window):
         if state is False:
             logging.info("Moonraker not connected")
             return
+        self.connecting = not self._ws.connected
+        self.connected_printer = self.connecting_to_printer
+        self.base_panel.set_ks_printer_cfg(self.connected_printer)
+
         # Moonraker is ready, set a loop to init the printer
         self.reinit_count += 1