feat: implement graceful shutdown with SIGINT/SIGTERM support
- Add WebSocket shutdown() with client notification and graceful close - Enhance download service stop() with pending state persistence - Expand FastAPI lifespan shutdown with proper cleanup sequence - Add SQLite WAL checkpoint before database close - Update stop_server.sh to use SIGTERM with timeout fallback - Configure uvicorn timeout_graceful_shutdown=30s - Update ARCHITECTURE.md with shutdown documentation
This commit is contained in:
@@ -322,6 +322,85 @@ class ConnectionManager:
|
||||
connection_id=connection_id,
|
||||
)
|
||||
|
||||
async def shutdown(self, timeout: float = 5.0) -> None:
|
||||
"""Gracefully shutdown all WebSocket connections.
|
||||
|
||||
Broadcasts a shutdown notification to all clients, then closes
|
||||
each connection with proper close codes.
|
||||
|
||||
Args:
|
||||
timeout: Maximum time (seconds) to wait for all closes to complete
|
||||
"""
|
||||
logger.info(
|
||||
"Initiating WebSocket shutdown, connections=%d",
|
||||
len(self._active_connections)
|
||||
)
|
||||
|
||||
# Broadcast shutdown notification to all clients
|
||||
shutdown_message = {
|
||||
"type": "server_shutdown",
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"data": {
|
||||
"message": "Server is shutting down",
|
||||
"reason": "graceful_shutdown",
|
||||
},
|
||||
}
|
||||
|
||||
try:
|
||||
await self.broadcast(shutdown_message)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to broadcast shutdown message: %s", e)
|
||||
|
||||
# Close all connections gracefully
|
||||
async with self._lock:
|
||||
connection_ids = list(self._active_connections.keys())
|
||||
|
||||
close_tasks = []
|
||||
for connection_id in connection_ids:
|
||||
websocket = self._active_connections.get(connection_id)
|
||||
if websocket:
|
||||
close_tasks.append(
|
||||
self._close_connection_gracefully(connection_id, websocket)
|
||||
)
|
||||
|
||||
if close_tasks:
|
||||
# Wait for all closes with timeout
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
asyncio.gather(*close_tasks, return_exceptions=True),
|
||||
timeout=timeout
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning(
|
||||
"WebSocket shutdown timed out after %.1f seconds", timeout
|
||||
)
|
||||
|
||||
# Clear all data structures
|
||||
async with self._lock:
|
||||
self._active_connections.clear()
|
||||
self._rooms.clear()
|
||||
self._connection_metadata.clear()
|
||||
|
||||
logger.info("WebSocket shutdown complete")
|
||||
|
||||
async def _close_connection_gracefully(
|
||||
self, connection_id: str, websocket: WebSocket
|
||||
) -> None:
|
||||
"""Close a single WebSocket connection gracefully.
|
||||
|
||||
Args:
|
||||
connection_id: The connection identifier
|
||||
websocket: The WebSocket connection to close
|
||||
"""
|
||||
try:
|
||||
# Code 1001 = Going Away (server shutdown)
|
||||
await websocket.close(code=1001, reason="Server shutdown")
|
||||
logger.debug("Closed WebSocket connection: %s", connection_id)
|
||||
except Exception as e:
|
||||
logger.debug(
|
||||
"Error closing WebSocket %s: %s", connection_id, str(e)
|
||||
)
|
||||
|
||||
|
||||
class WebSocketService:
|
||||
"""High-level WebSocket service for application-wide messaging.
|
||||
@@ -579,6 +658,18 @@ class WebSocketService:
|
||||
elapsed_seconds=round(elapsed_seconds, 2),
|
||||
)
|
||||
|
||||
async def shutdown(self, timeout: float = 5.0) -> None:
|
||||
"""Gracefully shutdown the WebSocket service.
|
||||
|
||||
Broadcasts shutdown notification and closes all connections.
|
||||
|
||||
Args:
|
||||
timeout: Maximum time (seconds) to wait for shutdown
|
||||
"""
|
||||
logger.info("Shutting down WebSocket service...")
|
||||
await self._manager.shutdown(timeout=timeout)
|
||||
logger.info("WebSocket service shutdown complete")
|
||||
|
||||
|
||||
# Singleton instance for application-wide access
|
||||
_websocket_service: Optional[WebSocketService] = None
|
||||
|
||||
Reference in New Issue
Block a user