save state to pickle file, restart timers on restart

hbd/server/ws.py

@@ -65,11 +65,7 @@ async def _handler(websocket, path=None):
         logger.exception("WebSocket handler exception from %s: %s", remote_address, e)
     finally:
         logger.debug("Removing WebSocket connection from %s", remote_address)
-        try:
-            _connections.remove(websocket)
-        except KeyError:
-            pass
-        await websocket.wait_closed()
+        _connections.discard(websocket)
 
 
 async def start(
@@ -93,33 +89,40 @@ async def start(
     _verbose = config.get("verbose", False),
     _debug = config.get("debug", 0),
 
-    servers = []
-    # plain WebSocket
-    websockets_logger = logging.getLogger("websockets.server")
-    #if _debug > 2:
-    #    websockets_logger.setLevel(logging.DEBUG)
-    #else:   
-    #    websockets_logger.setLevel(logging.INFO)
-    # regular WebSocket
-    ws_server = websockets.serve(_handler, host, ws_port)  # , subprotocols=["hbd"])
-    servers.append(ws_server)
-    # secure WebSocket (optional)
+    # Start servers and keep the server objects for clean shutdown
+    running_servers = []
+    ws_server = await websockets.serve(_handler, host, ws_port)
+    running_servers.append(ws_server)
     if wss_port and ssl_context:
-        wss_server = websockets.serve(
-            _handler, host, wss_port, ssl=ssl_context
-        )  # , subprotocols=["hbd"])
-        servers.append(wss_server)
-
-    # await starting of all servers
-    for srv in servers:
-        await srv
+        wss_server = await websockets.serve(_handler, host, wss_port, ssl=ssl_context)
+        running_servers.append(wss_server)
 
     logger.info(
         "WebSocket server(s) started on port %s (wss %s)", ws_port, wss_port
     )
 
-    # block forever (until loop is stopped or cancelled)
-    await asyncio.Future()
+    try:
+        # Block until cancelled
+        await asyncio.Future()
+    except asyncio.CancelledError:
+        pass
+    finally:
+        # Close all active browser connections so their handler coroutines exit
+        active = list(_connections)
+        if active:
+            logger.info("Closing %d active WebSocket connection(s)...", len(active))
+            await asyncio.gather(
+                *[ws.close() for ws in active],
+                return_exceptions=True,
+            )
+        # Stop the listening servers and wait for all handlers to finish
+        for srv in running_servers:
+            srv.close()
+        await asyncio.gather(
+            *[srv.wait_closed() for srv in running_servers],
+            return_exceptions=True,
+        )
+        logger.info("WebSocket server(s) stopped")
 
 
 def broadcast(typ: str, data) -> bool: