feat: Enhance tank control service with connection management and reset functionality

Author: 02loveslollipop

control_broker/app.py

@@ -6,7 +6,7 @@
 from datetime import timedelta
 from typing import List
 
-from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
 from fastapi.middleware.cors import CORSMiddleware
 import redis.asyncio as redis
 from redis import exceptions as redis_exceptions
@@ -25,7 +25,10 @@
 config = get_config()
 
 # Initialize connection manager
-manager = ConnectionManager()
+manager = ConnectionManager(
+    stale_timeout_seconds=config.tank_stale_timeout_seconds,
+    prune_interval_seconds=config.tank_prune_interval_seconds,
+)
 radar_broker = RadarBroker()
 redis_client_lock = asyncio.Lock()
 
@@ -88,6 +91,7 @@ async def get_redis_client() -> redis.Redis:
 async def on_startup() -> None:
     """Initialize Redis connection and start command listener."""
     await reset_redis_client()
+    await manager.start()
 
     # Start Redis command stream listener
     listener = RedisCommandListener(get_redis_client, reset_redis_client, config, manager)
@@ -108,6 +112,9 @@ async def on_shutdown() -> None:
         await redis_client.close()
     app.state.redis = None
 
+    await manager.stop()
+    await manager.close_all()
+
 
 # ========================================
 # Middleware
@@ -141,6 +148,15 @@ async def list_tanks() -> List[dict]:
     return await manager.snapshot()
 
 
+@app.post("/tanks/{tank_id}/reset")
+async def reset_tank(tank_id: str) -> dict:
+    """Forcefully reset a tank connection and clear its state."""
+    reset = await manager.force_reset(tank_id)
+    if not reset:
+        raise HTTPException(status_code=404, detail="Tank not found")
+    return {"status": "reset", "tankId": tank_id, "timestamp": utcnow().isoformat()}
+
+
 @app.get("/radars")
 async def list_radars() -> List[dict]:
     """List all connected radar sources."""

control_broker/core/config.py

@@ -15,6 +15,8 @@ class Config:
     redis_status_maxlen: int
     redis_radar_stream: str
     redis_radar_maxlen: int
+    tank_stale_timeout_seconds: int
+    tank_prune_interval_seconds: int
 
 
 def get_config() -> Config:
@@ -27,4 +29,6 @@ def get_config() -> Config:
         redis_status_maxlen=int(os.getenv("REDIS_STATUS_MAXLEN", "500")),
         redis_radar_stream=os.getenv("REDIS_RADAR_STREAM", "tank_radar"),
         redis_radar_maxlen=int(os.getenv("REDIS_RADAR_MAXLEN", "1000")),
+        tank_stale_timeout_seconds=int(os.getenv("TANK_STALE_TIMEOUT_SECONDS", "600")),
+        tank_prune_interval_seconds=int(os.getenv("TANK_PRUNE_INTERVAL_SECONDS", "30")),
     )

control_broker/services/connection_manager.py

@@ -4,7 +4,7 @@
 import json
 from contextlib import suppress
 from datetime import timedelta
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple
 from fastapi import WebSocket
 
 from models import TankInfo
@@ -14,25 +14,57 @@
 class ConnectionManager:
     """Manages WebSocket connections to tanks."""
 
-    def __init__(self) -> None:
+    def __init__(
+        self,
+        *,
+        stale_timeout_seconds: int = 600,
+        prune_interval_seconds: int = 30,
+    ) -> None:
         self._tanks: Dict[str, TankInfo] = {}
         self._lock = asyncio.Lock()
-        self._stale_timeout = timedelta(minutes=10)
+        self._stale_timeout = timedelta(seconds=max(1, stale_timeout_seconds))
+        self._prune_interval = timedelta(seconds=max(5, prune_interval_seconds))
+        self._maintenance_task: Optional[asyncio.Task] = None
 
-    async def _prune_stale(self) -> None:
+    async def start(self) -> None:
+        """Start background maintenance tasks."""
+        if self._maintenance_task and not self._maintenance_task.done():
+            return
+        self._maintenance_task = asyncio.create_task(self._run_auto_prune())
+
+    async def stop(self) -> None:
+        """Stop maintenance tasks and wait for completion."""
+        task = self._maintenance_task
+        if task:
+            task.cancel()
+            with suppress(asyncio.CancelledError):
+                await task
+        self._maintenance_task = None
+
+    async def _run_auto_prune(self) -> None:
+        """Periodically prune stale tank connections."""
+        try:
+            while True:
+                await asyncio.sleep(self._prune_interval.total_seconds())
+                await self._prune_stale(reason="auto-prune")
+        except asyncio.CancelledError:
+            pass
+
+    async def _prune_stale(self, *, reason: str = "stale") -> None:
         """Remove tanks that have been inactive for longer than the timeout."""
         now = utcnow()
-        to_close: List[WebSocket] = []
+        to_close: List[Tuple[str, WebSocket]] = []
         async with self._lock:
             for tank_id, info in list(self._tanks.items()):
                 if (now - info.last_seen) > self._stale_timeout:
                     websocket = info.websocket
                     if websocket is not None:
-                        to_close.append(websocket)
+                        to_close.append((tank_id, websocket))
                     self._tanks.pop(tank_id, None)
-        for websocket in to_close:
+        for tank_id, websocket in to_close:
             with suppress(Exception):
                 await websocket.close(code=1011)
+            print(f"[MANAGER] Pruned tank '{tank_id}' due to {reason}")
 
     async def register_tank(self, tank_id: str, websocket: WebSocket) -> TankInfo:
         """Register a new tank connection."""
@@ -110,3 +142,28 @@ async def update_last_seen(self, tank_id: str, payload: Optional[dict]) -> None:
             info.last_seen = utcnow()
             if payload is not None:
                 info.last_payload = payload
+
+    async def force_reset(self, tank_id: str) -> bool:
+        """Forcefully terminate and remove a tank connection."""
+        async with self._lock:
+            info = self._tanks.pop(tank_id, None)
+        if not info:
+            return False
+        websocket = info.websocket
+        if websocket:
+            with suppress(Exception):
+                await websocket.close(code=1012)
+        print(f"[MANAGER] Forced reset for tank '{tank_id}'")
+        return True
+
+    async def close_all(self) -> None:
+        """Close all tracked tank connections."""
+        async with self._lock:
+            entries = list(self._tanks.items())
+            self._tanks.clear()
+        for tank_id, info in entries:
+            websocket = info.websocket
+            if websocket:
+                with suppress(Exception):
+                    await websocket.close(code=1001)
+            print(f"[MANAGER] Closed connection for tank '{tank_id}' during shutdown")