spoolman: add a websocket connect

Connect to the spoolman sevice via websocket to receive
spool events.  In addition, this gives Moonraker a persistent
connection to know when the service is available.

Signed-off-by:  Eric Callahan  <arksine.code@gmail.com>

moonraker/components/spoolman.py

@@ -8,29 +8,46 @@ from __future__ import annotations
 import asyncio
 import datetime
 import logging
-from typing import TYPE_CHECKING, Dict, Any
+import re
-from ..common import RequestType
+import tornado.websocket as tornado_ws
+from ..common import RequestType, Sentinel
+from ..utils import json_wrapper as jsonw
+from typing import (
+    TYPE_CHECKING,
+    Dict,
+    Any,
+    Optional,
+    Union
+)
 
 if TYPE_CHECKING:
-    from typing import Optional
+    from ..confighelper import ConfigHelper
     from ..common import WebRequest
     from .http_client import HttpClient, HttpResponse
     from .database import MoonrakerDatabase
     from .announcements import Announcements
     from .klippy_apis import KlippyAPI as APIComp
-    from confighelper import ConfigHelper
+    from tornado.websocket import WebSocketClientConnection
 
 DB_NAMESPACE = "moonraker"
 ACTIVE_SPOOL_KEY = "spoolman.spool_id"
-
+CONNECTION_ERROR_LOG_TIME = 60.
 
 class SpoolManager:
     def __init__(self, config: ConfigHelper):
         self.server = config.get_server()
+        self.eventloop = self.server.get_event_loop()
+        self._get_spoolman_urls(config)
         self.sync_rate_seconds = config.getint("sync_rate", default=5, minval=1)
         self.last_sync_time = datetime.datetime.now()
         self.extruded_lock = asyncio.Lock()
-        self.spoolman_url = f"{config.get('server').rstrip('/')}/api"
+        self.spoolman_ws: Optional[WebSocketClientConnection] = None
+        self.connection_task: Optional[asyncio.Task] = None
+        self.spool_check_task: Optional[asyncio.Task] = None
+        self.spool_lock = asyncio.Lock()
+        self.ws_connected: bool = False
+        self.reconnect_delay: float = 2.
+        self.is_closing: bool = False
         self.spool_id: Optional[int] = None
         self.extruded: float = 0
         self._error_logged: bool = False
@@ -47,6 +64,19 @@ class SpoolManager:
             "spoolman_set_active_spool", self.set_active_spool
         )
 
+    def _get_spoolman_urls(self, config: ConfigHelper) -> None:
+        orig_url = config.get('server')
+        url_match = re.match(r"(?i:(?P<scheme>https?)://)?(?P<host>.+)", orig_url)
+        if url_match is None:
+            raise config.error(
+                f"Section [spoolman], Option server: {orig_url}: Invalid URL format"
+            )
+        scheme = url_match["scheme"] or "http"
+        host = url_match["host"].rstrip("/")
+        ws_scheme = "wss" if scheme == "https" else "ws"
+        self.spoolman_url = f"{scheme}://{host}/api"
+        self.ws_url = f"{ws_scheme}://{host}/api/v1/spool"
+
     def _register_notifications(self):
         self.server.register_notification("spoolman:active_spool_set")
 
@@ -71,22 +101,101 @@ class SpoolManager:
         self.spool_id = await self.database.get_item(
             DB_NAMESPACE, ACTIVE_SPOOL_KEY, None
         )
+        self.connection_task = self.eventloop.create_task(self._connect_websocket())
+
+    async def _connect_websocket(self) -> None:
+        log_connect: bool = True
+        last_err: Exception = Exception()
+        while not self.is_closing:
+            if log_connect:
+                logging.info(f"Connecting To Spoolman: {self.ws_url}")
+                log_connect = False
+            try:
+                self.spoolman_ws = await tornado_ws.websocket_connect(
+                    self.ws_url,
+                    connect_timeout=5.,
+                    ping_interval=20.,
+                    ping_timeout=60.
+                )
+                setattr(self.spoolman_ws, "on_ping", self._on_ws_ping)
+                cur_time = self.eventloop.get_loop_time()
+                self._last_ping_received = cur_time
+            except asyncio.CancelledError:
+                raise
+            except Exception as e:
+                if type(last_err) is not type(e) or last_err.args != e.args:
+                    logging.exception("Failed to connect to Spoolman")
+                    last_err = e
+            else:
+                self.ws_connected = True
+                self._error_logged = False
+                logging.info("Connected to Spoolman Spool Manager")
+                if self.spool_id is not None:
+                    self._cancel_spool_check_task()
+                    coro = self._check_spool_deleted()
+                    self.spool_check_task = self.eventloop.create_task(coro)
+                await self._read_messages()
+                log_connect = True
+                last_err = Exception()
+            if not self.is_closing:
+                await asyncio.sleep(self.reconnect_delay)
+
+    async def _read_messages(self) -> None:
+        message: Union[str, bytes, None]
+        while self.spoolman_ws is not None:
+            message = await self.spoolman_ws.read_message()
+            if isinstance(message, str):
+                self._decode_message(message)
+            elif message is None:
+                self.ws_connected = False
+                cur_time = self.eventloop.get_loop_time()
+                ping_time: float = cur_time - self._last_ping_received
+                reason = code = None
+                if self.spoolman_ws is not None:
+                    reason = self.spoolman_ws.close_reason
+                    code = self.spoolman_ws.close_code
+                logging.info(
+                    f"Spoolman Disconnected - Code: {code}, Reason: {reason}, "
+                    f"Server Ping Time Elapsed: {ping_time}"
+                )
+                self.spoolman_ws = None
+                break
+
+    def _decode_message(self, message: str) -> None:
+        event: Dict[str, Any] = jsonw.loads(message)
+        if event.get("resource") != "spool":
+            return
+        if self.spool_id is not None and event.get("type") == "deleted":
+            payload: Dict[str, Any] = event.get("payload", {})
+            if payload.get("id") == self.spool_id:
+                self.eventloop.create_task(self.set_active_spool(Sentinel.MISSING))
+
+    def _cancel_spool_check_task(self) -> None:
+        if self.spool_check_task is None or self.spool_check_task.done():
+            return
+        self.spool_check_task.cancel()
+
+    async def _check_spool_deleted(self) -> None:
         if self.spool_id is not None:
             response = await self.http_client.get(
                 f"{self.spoolman_url}/v1/spool/{self.spool_id}",
-                connect_timeout=1., request_timeout=2.,
+                connect_timeout=1., request_timeout=2.
             )
             if response.status_code == 404:
                 logging.info(f"Spool ID {self.spool_id} not found, setting to None")
-                self._set_spool(None)
+                await self.set_active_spool(Sentinel.MISSING)
             elif response.has_error():
                 err_msg = self._get_response_error(response)
-                logging.info(
+                logging.info(f"Attempt to check spool status failed: {err_msg}")
-                    "Attempt to initialize Spoolman connection failed with the "
-                    f"following: {err_msg}"
-                )
             else:
                 logging.info(f"Found Spool ID {self.spool_id} on spoolman instance")
+        self.spool_check_task = None
+
+    def connected(self) -> bool:
+        return self.ws_connected
+
+    def _on_ws_ping(self, data: bytes = b"") -> None:
+        self._last_ping_received = self.eventloop.get_loop_time()
 
     async def _handle_klippy_ready(self):
         result = await self.klippy_apis.subscribe_objects(
@@ -130,27 +239,29 @@ class SpoolManager:
                 logging.debug("Sync period elapsed, tracking usage")
                 await self.track_filament_usage()
 
-    async def set_active_spool(self, spool_id: Optional[int]) -> None:
+    async def set_active_spool(self, spool_id: Union[int, Sentinel, None]) -> None:
-        if self.spool_id == spool_id:
+        async with self.spool_lock:
-            logging.info(f"Spool ID already set to: {spool_id}")
+            deleted_spool = False
-            return
+            if spool_id is Sentinel.MISSING:
-        # Store the current spool usage before switching
+                spool_id = None
-        if self.spool_id is not None:
+                deleted_spool = True
-            await self.track_filament_usage()
+            if self.spool_id == spool_id:
-        elif spool_id is not None:
+                logging.info(f"Spool ID already set to: {spool_id}")
-            async with self.extruded_lock:
+                return
-                self.extruded = 0
+            # Store the current spool usage before switching, unless it has been deleted
-        self._set_spool(spool_id)
+            if not deleted_spool:
-        logging.info(f"Setting active spool to: {spool_id}")
+                if self.spool_id is not None:
-
+                    await self.track_filament_usage()
-    def _set_spool(self, spool_id: Optional[int]) -> None:
+                elif spool_id is not None:
-        if spool_id == self.spool_id:
+                    # No need to track, just reset extrusion
-            return
+                    async with self.extruded_lock:
-        self.spool_id = spool_id
+                        self.extruded = 0
-        self.database.insert_item(DB_NAMESPACE, ACTIVE_SPOOL_KEY, spool_id)
+            self.spool_id = spool_id
-        self.server.send_event(
+            self.database.insert_item(DB_NAMESPACE, ACTIVE_SPOOL_KEY, spool_id)
-            "spoolman:active_spool_set", {"spool_id": spool_id}
+            self.server.send_event(
-        )
+                "spoolman:active_spool_set", {"spool_id": spool_id}
+            )
+            logging.info(f"Setting active spool to: {spool_id}")
 
     async def track_filament_usage(self):
         spool_id = self.spool_id
@@ -158,7 +269,7 @@ class SpoolManager:
             logging.debug("No active spool, skipping tracking")
             return
         async with self.extruded_lock:
-            if self.extruded > 0:
+            if self.extruded > 0 and self.ws_connected:
                 used_length = self.extruded
 
                 logging.debug(
@@ -176,19 +287,20 @@ class SpoolManager:
                 )
                 if response.has_error():
                     if response.status_code == 404:
+                        self._error_logged = False
                         logging.info(
                             f"Spool ID {self.spool_id} not found, setting to None"
                         )
-                        self._set_spool(None)
+                        coro = self.set_active_spool(Sentinel.MISSING)
-                    else:
+                        self.eventloop.create_task(coro)
-                        if not self._error_logged:
+                    elif not self._error_logged:
-                            error_msg = self._get_response_error(response)
+                        error_msg = self._get_response_error(response)
-                            self._error_logged = True
+                        self._error_logged = True
-                            logging.info(
+                        logging.info(
-                                f"Failed to update extrusion for spool id {spool_id}, "
+                            f"Failed to update extrusion for spool id {spool_id}, "
-                                f"received {error_msg}"
+                            f"received {error_msg}"
-                            )
+                        )
-                        return
+                    return
                 self._error_logged = False
                 self.extruded = 0
 
@@ -204,36 +316,38 @@ class SpoolManager:
         path = web_request.get_str("path")
         query = web_request.get_str("query", None)
         body = web_request.get("body", None)
-
         if method not in {"GET", "POST", "PUT", "PATCH", "DELETE"}:
             raise self.server.error(f"Invalid HTTP method: {method}")
-
         if body is not None and method == "GET":
             raise self.server.error("GET requests cannot have a body")
-
         if len(path) < 4 or path[:4] != "/v1/":
             raise self.server.error(
                 "Invalid path, must start with the API version, e.g. /v1"
             )
-
+        query = f"?{query}" if query is not None else ""
-        if query is not None:
-            query = f"?{query}"
-        else:
-            query = ""
-
         full_url = f"{self.spoolman_url}{path}{query}"
-
+        if not self.ws_connected:
+            raise self.server.error("Spoolman server not available", 503)
         logging.debug(f"Proxying {method} request to {full_url}")
-
         response = await self.http_client.request(
             method=method,
             url=full_url,
             body=body,
         )
         response.raise_for_status()
-
         return response.json()
 
+    async def close(self):
+        self.is_closing = True
+        if self.spoolman_ws is not None:
+            self.spoolman_ws.close(1001, "Moonraker Shutdown")
+        self._cancel_spool_check_task()
+        if self.connection_task is None or self.connection_task.done():
+            return
+        try:
+            await asyncio.wait_for(self.connection_task, 2.)
+        except asyncio.TimeoutError:
+            pass
 
 def load_component(config: ConfigHelper) -> SpoolManager:
     return SpoolManager(config)