mirror of
https://github.com/ruvnet/RuView
synced 2026-07-27 18:11:43 +00:00
feat: Complete Rust port of WiFi-DensePose with modular crates
Major changes: - Organized Python v1 implementation into v1/ subdirectory - Created Rust workspace with 9 modular crates: - wifi-densepose-core: Core types, traits, errors - wifi-densepose-signal: CSI processing, phase sanitization, FFT - wifi-densepose-nn: Neural network inference (ONNX/Candle/tch) - wifi-densepose-api: Axum-based REST/WebSocket API - wifi-densepose-db: SQLx database layer - wifi-densepose-config: Configuration management - wifi-densepose-hardware: Hardware abstraction - wifi-densepose-wasm: WebAssembly bindings - wifi-densepose-cli: Command-line interface Documentation: - ADR-001: Workspace structure - ADR-002: Signal processing library selection - ADR-003: Neural network inference strategy - DDD domain model with bounded contexts Testing: - 69 tests passing across all crates - Signal processing: 45 tests - Neural networks: 21 tests - Core: 3 doc tests Performance targets: - 10x faster CSI processing (~0.5ms vs ~5ms) - 5x lower memory usage (~100MB vs ~500MB) - WASM support for browser deployment
This commit is contained in:
Claude
@@ -0,0 +1,8 @@
|
||||
"""
|
||||
WebSocket handlers package
|
||||
"""
|
||||
|
||||
from .connection_manager import ConnectionManager
|
||||
from .pose_stream import PoseStreamHandler
|
||||
|
||||
__all__ = ["ConnectionManager", "PoseStreamHandler"]
|
||||
@@ -0,0 +1,461 @@
|
||||
"""
|
||||
WebSocket connection manager for WiFi-DensePose API
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from typing import Dict, List, Optional, Any, Set
|
||||
from datetime import datetime, timedelta
|
||||
from collections import defaultdict
|
||||
|
||||
from fastapi import WebSocket, WebSocketDisconnect
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class WebSocketConnection:
|
||||
"""Represents a WebSocket connection with metadata."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
websocket: WebSocket,
|
||||
client_id: str,
|
||||
stream_type: str,
|
||||
zone_ids: Optional[List[str]] = None,
|
||||
**config
|
||||
):
|
||||
self.websocket = websocket
|
||||
self.client_id = client_id
|
||||
self.stream_type = stream_type
|
||||
self.zone_ids = zone_ids or []
|
||||
self.config = config
|
||||
self.connected_at = datetime.utcnow()
|
||||
self.last_ping = datetime.utcnow()
|
||||
self.message_count = 0
|
||||
self.is_active = True
|
||||
|
||||
async def send_json(self, data: Dict[str, Any]):
|
||||
"""Send JSON data to client."""
|
||||
try:
|
||||
await self.websocket.send_json(data)
|
||||
self.message_count += 1
|
||||
except Exception as e:
|
||||
logger.error(f"Error sending to client {self.client_id}: {e}")
|
||||
self.is_active = False
|
||||
raise
|
||||
|
||||
async def send_text(self, message: str):
|
||||
"""Send text message to client."""
|
||||
try:
|
||||
await self.websocket.send_text(message)
|
||||
self.message_count += 1
|
||||
except Exception as e:
|
||||
logger.error(f"Error sending text to client {self.client_id}: {e}")
|
||||
self.is_active = False
|
||||
raise
|
||||
|
||||
def update_config(self, config: Dict[str, Any]):
|
||||
"""Update connection configuration."""
|
||||
self.config.update(config)
|
||||
|
||||
# Update zone IDs if provided
|
||||
if "zone_ids" in config:
|
||||
self.zone_ids = config["zone_ids"] or []
|
||||
|
||||
def matches_filter(
|
||||
self,
|
||||
stream_type: Optional[str] = None,
|
||||
zone_ids: Optional[List[str]] = None,
|
||||
**filters
|
||||
) -> bool:
|
||||
"""Check if connection matches given filters."""
|
||||
# Check stream type
|
||||
if stream_type and self.stream_type != stream_type:
|
||||
return False
|
||||
|
||||
# Check zone IDs
|
||||
if zone_ids:
|
||||
if not self.zone_ids: # Connection listens to all zones
|
||||
return True
|
||||
# Check if any requested zone is in connection's zones
|
||||
if not any(zone in self.zone_ids for zone in zone_ids):
|
||||
return False
|
||||
|
||||
# Check additional filters
|
||||
for key, value in filters.items():
|
||||
if key in self.config and self.config[key] != value:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def get_info(self) -> Dict[str, Any]:
|
||||
"""Get connection information."""
|
||||
return {
|
||||
"client_id": self.client_id,
|
||||
"stream_type": self.stream_type,
|
||||
"zone_ids": self.zone_ids,
|
||||
"config": self.config,
|
||||
"connected_at": self.connected_at.isoformat(),
|
||||
"last_ping": self.last_ping.isoformat(),
|
||||
"message_count": self.message_count,
|
||||
"is_active": self.is_active,
|
||||
"uptime_seconds": (datetime.utcnow() - self.connected_at).total_seconds()
|
||||
}
|
||||
|
||||
|
||||
class ConnectionManager:
|
||||
"""Manages WebSocket connections for real-time streaming."""
|
||||
|
||||
def __init__(self):
|
||||
self.connections: Dict[str, WebSocketConnection] = {}
|
||||
self.connections_by_type: Dict[str, Set[str]] = defaultdict(set)
|
||||
self.connections_by_zone: Dict[str, Set[str]] = defaultdict(set)
|
||||
self.metrics = {
|
||||
"total_connections": 0,
|
||||
"active_connections": 0,
|
||||
"messages_sent": 0,
|
||||
"errors": 0,
|
||||
"start_time": datetime.utcnow()
|
||||
}
|
||||
self._cleanup_task = None
|
||||
self._started = False
|
||||
|
||||
async def connect(
|
||||
self,
|
||||
websocket: WebSocket,
|
||||
stream_type: str,
|
||||
zone_ids: Optional[List[str]] = None,
|
||||
**config
|
||||
) -> str:
|
||||
"""Register a new WebSocket connection."""
|
||||
client_id = str(uuid.uuid4())
|
||||
|
||||
try:
|
||||
# Create connection object
|
||||
connection = WebSocketConnection(
|
||||
websocket=websocket,
|
||||
client_id=client_id,
|
||||
stream_type=stream_type,
|
||||
zone_ids=zone_ids,
|
||||
**config
|
||||
)
|
||||
|
||||
# Store connection
|
||||
self.connections[client_id] = connection
|
||||
self.connections_by_type[stream_type].add(client_id)
|
||||
|
||||
# Index by zones
|
||||
if zone_ids:
|
||||
for zone_id in zone_ids:
|
||||
self.connections_by_zone[zone_id].add(client_id)
|
||||
|
||||
# Update metrics
|
||||
self.metrics["total_connections"] += 1
|
||||
self.metrics["active_connections"] = len(self.connections)
|
||||
|
||||
logger.info(f"WebSocket client {client_id} connected for {stream_type}")
|
||||
|
||||
return client_id
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error connecting WebSocket client: {e}")
|
||||
raise
|
||||
|
||||
async def disconnect(self, client_id: str) -> bool:
|
||||
"""Disconnect a WebSocket client."""
|
||||
if client_id not in self.connections:
|
||||
return False
|
||||
|
||||
try:
|
||||
connection = self.connections[client_id]
|
||||
|
||||
# Remove from indexes
|
||||
self.connections_by_type[connection.stream_type].discard(client_id)
|
||||
|
||||
for zone_id in connection.zone_ids:
|
||||
self.connections_by_zone[zone_id].discard(client_id)
|
||||
|
||||
# Close WebSocket if still active
|
||||
if connection.is_active:
|
||||
try:
|
||||
await connection.websocket.close()
|
||||
except:
|
||||
pass # Connection might already be closed
|
||||
|
||||
# Remove connection
|
||||
del self.connections[client_id]
|
||||
|
||||
# Update metrics
|
||||
self.metrics["active_connections"] = len(self.connections)
|
||||
|
||||
logger.info(f"WebSocket client {client_id} disconnected")
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error disconnecting client {client_id}: {e}")
|
||||
return False
|
||||
|
||||
async def disconnect_all(self):
|
||||
"""Disconnect all WebSocket clients."""
|
||||
client_ids = list(self.connections.keys())
|
||||
|
||||
for client_id in client_ids:
|
||||
await self.disconnect(client_id)
|
||||
|
||||
logger.info("All WebSocket clients disconnected")
|
||||
|
||||
async def send_to_client(self, client_id: str, data: Dict[str, Any]) -> bool:
|
||||
"""Send data to a specific client."""
|
||||
if client_id not in self.connections:
|
||||
return False
|
||||
|
||||
connection = self.connections[client_id]
|
||||
|
||||
try:
|
||||
await connection.send_json(data)
|
||||
self.metrics["messages_sent"] += 1
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error sending to client {client_id}: {e}")
|
||||
self.metrics["errors"] += 1
|
||||
|
||||
# Mark connection as inactive and schedule for cleanup
|
||||
connection.is_active = False
|
||||
return False
|
||||
|
||||
async def broadcast(
|
||||
self,
|
||||
data: Dict[str, Any],
|
||||
stream_type: Optional[str] = None,
|
||||
zone_ids: Optional[List[str]] = None,
|
||||
**filters
|
||||
) -> int:
|
||||
"""Broadcast data to matching clients."""
|
||||
sent_count = 0
|
||||
failed_clients = []
|
||||
|
||||
# Get matching connections
|
||||
matching_clients = self._get_matching_clients(
|
||||
stream_type=stream_type,
|
||||
zone_ids=zone_ids,
|
||||
**filters
|
||||
)
|
||||
|
||||
# Send to all matching clients
|
||||
for client_id in matching_clients:
|
||||
try:
|
||||
success = await self.send_to_client(client_id, data)
|
||||
if success:
|
||||
sent_count += 1
|
||||
else:
|
||||
failed_clients.append(client_id)
|
||||
except Exception as e:
|
||||
logger.error(f"Error broadcasting to client {client_id}: {e}")
|
||||
failed_clients.append(client_id)
|
||||
|
||||
# Clean up failed connections
|
||||
for client_id in failed_clients:
|
||||
await self.disconnect(client_id)
|
||||
|
||||
return sent_count
|
||||
|
||||
async def update_client_config(self, client_id: str, config: Dict[str, Any]) -> bool:
|
||||
"""Update client configuration."""
|
||||
if client_id not in self.connections:
|
||||
return False
|
||||
|
||||
connection = self.connections[client_id]
|
||||
old_zones = set(connection.zone_ids)
|
||||
|
||||
# Update configuration
|
||||
connection.update_config(config)
|
||||
|
||||
# Update zone indexes if zones changed
|
||||
new_zones = set(connection.zone_ids)
|
||||
|
||||
# Remove from old zones
|
||||
for zone_id in old_zones - new_zones:
|
||||
self.connections_by_zone[zone_id].discard(client_id)
|
||||
|
||||
# Add to new zones
|
||||
for zone_id in new_zones - old_zones:
|
||||
self.connections_by_zone[zone_id].add(client_id)
|
||||
|
||||
return True
|
||||
|
||||
async def get_client_status(self, client_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""Get status of a specific client."""
|
||||
if client_id not in self.connections:
|
||||
return None
|
||||
|
||||
return self.connections[client_id].get_info()
|
||||
|
||||
async def get_connected_clients(self) -> List[Dict[str, Any]]:
|
||||
"""Get list of all connected clients."""
|
||||
return [conn.get_info() for conn in self.connections.values()]
|
||||
|
||||
async def get_connection_stats(self) -> Dict[str, Any]:
|
||||
"""Get connection statistics."""
|
||||
stats = {
|
||||
"total_clients": len(self.connections),
|
||||
"clients_by_type": {
|
||||
stream_type: len(clients)
|
||||
for stream_type, clients in self.connections_by_type.items()
|
||||
},
|
||||
"clients_by_zone": {
|
||||
zone_id: len(clients)
|
||||
for zone_id, clients in self.connections_by_zone.items()
|
||||
if clients # Only include zones with active clients
|
||||
},
|
||||
"active_clients": sum(1 for conn in self.connections.values() if conn.is_active),
|
||||
"inactive_clients": sum(1 for conn in self.connections.values() if not conn.is_active)
|
||||
}
|
||||
|
||||
return stats
|
||||
|
||||
async def get_metrics(self) -> Dict[str, Any]:
|
||||
"""Get detailed metrics."""
|
||||
uptime = (datetime.utcnow() - self.metrics["start_time"]).total_seconds()
|
||||
|
||||
return {
|
||||
**self.metrics,
|
||||
"active_connections": len(self.connections),
|
||||
"uptime_seconds": uptime,
|
||||
"messages_per_second": self.metrics["messages_sent"] / max(uptime, 1),
|
||||
"error_rate": self.metrics["errors"] / max(self.metrics["messages_sent"], 1)
|
||||
}
|
||||
|
||||
def _get_matching_clients(
|
||||
self,
|
||||
stream_type: Optional[str] = None,
|
||||
zone_ids: Optional[List[str]] = None,
|
||||
**filters
|
||||
) -> List[str]:
|
||||
"""Get client IDs that match the given filters."""
|
||||
candidates = set(self.connections.keys())
|
||||
|
||||
# Filter by stream type
|
||||
if stream_type:
|
||||
type_clients = self.connections_by_type.get(stream_type, set())
|
||||
candidates &= type_clients
|
||||
|
||||
# Filter by zones
|
||||
if zone_ids:
|
||||
zone_clients = set()
|
||||
for zone_id in zone_ids:
|
||||
zone_clients.update(self.connections_by_zone.get(zone_id, set()))
|
||||
|
||||
# Also include clients listening to all zones (empty zone list)
|
||||
all_zone_clients = {
|
||||
client_id for client_id, conn in self.connections.items()
|
||||
if not conn.zone_ids
|
||||
}
|
||||
zone_clients.update(all_zone_clients)
|
||||
|
||||
candidates &= zone_clients
|
||||
|
||||
# Apply additional filters
|
||||
matching_clients = []
|
||||
for client_id in candidates:
|
||||
connection = self.connections[client_id]
|
||||
if connection.is_active and connection.matches_filter(**filters):
|
||||
matching_clients.append(client_id)
|
||||
|
||||
return matching_clients
|
||||
|
||||
async def ping_clients(self):
|
||||
"""Send ping to all connected clients."""
|
||||
ping_data = {
|
||||
"type": "ping",
|
||||
"timestamp": datetime.utcnow().isoformat()
|
||||
}
|
||||
|
||||
failed_clients = []
|
||||
|
||||
for client_id, connection in self.connections.items():
|
||||
try:
|
||||
await connection.send_json(ping_data)
|
||||
connection.last_ping = datetime.utcnow()
|
||||
except Exception as e:
|
||||
logger.warning(f"Ping failed for client {client_id}: {e}")
|
||||
failed_clients.append(client_id)
|
||||
|
||||
# Clean up failed connections
|
||||
for client_id in failed_clients:
|
||||
await self.disconnect(client_id)
|
||||
|
||||
async def cleanup_inactive_connections(self):
|
||||
"""Clean up inactive or stale connections."""
|
||||
now = datetime.utcnow()
|
||||
stale_threshold = timedelta(minutes=5) # 5 minutes without ping
|
||||
|
||||
stale_clients = []
|
||||
|
||||
for client_id, connection in self.connections.items():
|
||||
# Check if connection is inactive
|
||||
if not connection.is_active:
|
||||
stale_clients.append(client_id)
|
||||
continue
|
||||
|
||||
# Check if connection is stale (no ping response)
|
||||
if now - connection.last_ping > stale_threshold:
|
||||
logger.warning(f"Client {client_id} appears stale, disconnecting")
|
||||
stale_clients.append(client_id)
|
||||
|
||||
# Clean up stale connections
|
||||
for client_id in stale_clients:
|
||||
await self.disconnect(client_id)
|
||||
|
||||
if stale_clients:
|
||||
logger.info(f"Cleaned up {len(stale_clients)} stale connections")
|
||||
|
||||
async def start(self):
|
||||
"""Start the connection manager."""
|
||||
if not self._started:
|
||||
self._start_cleanup_task()
|
||||
self._started = True
|
||||
logger.info("Connection manager started")
|
||||
|
||||
def _start_cleanup_task(self):
|
||||
"""Start background cleanup task."""
|
||||
async def cleanup_loop():
|
||||
while True:
|
||||
try:
|
||||
await asyncio.sleep(60) # Run every minute
|
||||
await self.cleanup_inactive_connections()
|
||||
|
||||
# Send periodic ping every 2 minutes
|
||||
if datetime.utcnow().minute % 2 == 0:
|
||||
await self.ping_clients()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in cleanup task: {e}")
|
||||
|
||||
try:
|
||||
self._cleanup_task = asyncio.create_task(cleanup_loop())
|
||||
except RuntimeError:
|
||||
# No event loop running, will start later
|
||||
logger.debug("No event loop running, cleanup task will start later")
|
||||
|
||||
async def shutdown(self):
|
||||
"""Shutdown connection manager."""
|
||||
# Cancel cleanup task
|
||||
if self._cleanup_task:
|
||||
self._cleanup_task.cancel()
|
||||
try:
|
||||
await self._cleanup_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# Disconnect all clients
|
||||
await self.disconnect_all()
|
||||
|
||||
logger.info("Connection manager shutdown complete")
|
||||
|
||||
|
||||
# Global connection manager instance
|
||||
connection_manager = ConnectionManager()
|
||||
@@ -0,0 +1,384 @@
|
||||
"""
|
||||
Pose streaming WebSocket handler
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from typing import Dict, List, Optional, Any
|
||||
from datetime import datetime
|
||||
|
||||
from fastapi import WebSocket
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from src.api.websocket.connection_manager import ConnectionManager
|
||||
from src.services.pose_service import PoseService
|
||||
from src.services.stream_service import StreamService
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PoseStreamData(BaseModel):
|
||||
"""Pose stream data model."""
|
||||
|
||||
timestamp: datetime = Field(..., description="Data timestamp")
|
||||
zone_id: str = Field(..., description="Zone identifier")
|
||||
pose_data: Dict[str, Any] = Field(..., description="Pose estimation data")
|
||||
confidence: float = Field(..., ge=0.0, le=1.0, description="Confidence score")
|
||||
activity: Optional[str] = Field(default=None, description="Detected activity")
|
||||
metadata: Optional[Dict[str, Any]] = Field(default=None, description="Additional metadata")
|
||||
|
||||
|
||||
class PoseStreamHandler:
|
||||
"""Handles pose data streaming to WebSocket clients."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
connection_manager: ConnectionManager,
|
||||
pose_service: PoseService,
|
||||
stream_service: StreamService
|
||||
):
|
||||
self.connection_manager = connection_manager
|
||||
self.pose_service = pose_service
|
||||
self.stream_service = stream_service
|
||||
self.is_streaming = False
|
||||
self.stream_task = None
|
||||
self.subscribers = {}
|
||||
self.stream_config = {
|
||||
"fps": 30,
|
||||
"min_confidence": 0.5,
|
||||
"include_metadata": True,
|
||||
"buffer_size": 100
|
||||
}
|
||||
|
||||
async def start_streaming(self):
|
||||
"""Start pose data streaming."""
|
||||
if self.is_streaming:
|
||||
logger.warning("Pose streaming already active")
|
||||
return
|
||||
|
||||
self.is_streaming = True
|
||||
self.stream_task = asyncio.create_task(self._stream_loop())
|
||||
logger.info("Pose streaming started")
|
||||
|
||||
async def stop_streaming(self):
|
||||
"""Stop pose data streaming."""
|
||||
if not self.is_streaming:
|
||||
return
|
||||
|
||||
self.is_streaming = False
|
||||
|
||||
if self.stream_task:
|
||||
self.stream_task.cancel()
|
||||
try:
|
||||
await self.stream_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
logger.info("Pose streaming stopped")
|
||||
|
||||
async def _stream_loop(self):
|
||||
"""Main streaming loop."""
|
||||
try:
|
||||
logger.info("🚀 Starting pose streaming loop")
|
||||
while self.is_streaming:
|
||||
try:
|
||||
# Get current pose data from all zones
|
||||
logger.debug("📡 Getting current pose data...")
|
||||
pose_data = await self.pose_service.get_current_pose_data()
|
||||
logger.debug(f"📊 Received pose data: {pose_data}")
|
||||
|
||||
if pose_data:
|
||||
logger.debug("📤 Broadcasting pose data...")
|
||||
await self._process_and_broadcast_pose_data(pose_data)
|
||||
else:
|
||||
logger.debug("⚠️ No pose data received")
|
||||
|
||||
# Control streaming rate
|
||||
await asyncio.sleep(1.0 / self.stream_config["fps"])
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in pose streaming loop: {e}")
|
||||
await asyncio.sleep(1.0) # Brief pause on error
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.info("Pose streaming loop cancelled")
|
||||
except Exception as e:
|
||||
logger.error(f"Fatal error in pose streaming loop: {e}")
|
||||
finally:
|
||||
logger.info("🛑 Pose streaming loop stopped")
|
||||
self.is_streaming = False
|
||||
|
||||
async def _process_and_broadcast_pose_data(self, raw_pose_data: Dict[str, Any]):
|
||||
"""Process and broadcast pose data to subscribers."""
|
||||
try:
|
||||
# Process data for each zone
|
||||
for zone_id, zone_data in raw_pose_data.items():
|
||||
if not zone_data:
|
||||
continue
|
||||
|
||||
# Create structured pose data
|
||||
pose_stream_data = PoseStreamData(
|
||||
timestamp=datetime.utcnow(),
|
||||
zone_id=zone_id,
|
||||
pose_data=zone_data.get("pose", {}),
|
||||
confidence=zone_data.get("confidence", 0.0),
|
||||
activity=zone_data.get("activity"),
|
||||
metadata=zone_data.get("metadata") if self.stream_config["include_metadata"] else None
|
||||
)
|
||||
|
||||
# Filter by minimum confidence
|
||||
if pose_stream_data.confidence < self.stream_config["min_confidence"]:
|
||||
continue
|
||||
|
||||
# Broadcast to subscribers
|
||||
await self._broadcast_pose_data(pose_stream_data)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing pose data: {e}")
|
||||
|
||||
async def _broadcast_pose_data(self, pose_data: PoseStreamData):
|
||||
"""Broadcast pose data to matching WebSocket clients."""
|
||||
try:
|
||||
logger.debug(f"📡 Preparing to broadcast pose data for zone {pose_data.zone_id}")
|
||||
|
||||
# Prepare broadcast data
|
||||
broadcast_data = {
|
||||
"type": "pose_data",
|
||||
"timestamp": pose_data.timestamp.isoformat(),
|
||||
"zone_id": pose_data.zone_id,
|
||||
"data": {
|
||||
"pose": pose_data.pose_data,
|
||||
"confidence": pose_data.confidence,
|
||||
"activity": pose_data.activity
|
||||
}
|
||||
}
|
||||
|
||||
# Add metadata if enabled
|
||||
if pose_data.metadata and self.stream_config["include_metadata"]:
|
||||
broadcast_data["metadata"] = pose_data.metadata
|
||||
|
||||
logger.debug(f"📤 Broadcasting data: {broadcast_data}")
|
||||
|
||||
# Broadcast to pose stream subscribers
|
||||
sent_count = await self.connection_manager.broadcast(
|
||||
data=broadcast_data,
|
||||
stream_type="pose",
|
||||
zone_ids=[pose_data.zone_id]
|
||||
)
|
||||
|
||||
logger.info(f"✅ Broadcasted pose data for zone {pose_data.zone_id} to {sent_count} clients")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error broadcasting pose data: {e}")
|
||||
|
||||
async def handle_client_subscription(
|
||||
self,
|
||||
client_id: str,
|
||||
subscription_config: Dict[str, Any]
|
||||
):
|
||||
"""Handle client subscription configuration."""
|
||||
try:
|
||||
# Store client subscription config
|
||||
self.subscribers[client_id] = {
|
||||
"zone_ids": subscription_config.get("zone_ids", []),
|
||||
"min_confidence": subscription_config.get("min_confidence", 0.5),
|
||||
"max_fps": subscription_config.get("max_fps", 30),
|
||||
"include_metadata": subscription_config.get("include_metadata", True),
|
||||
"stream_types": subscription_config.get("stream_types", ["pose_data"]),
|
||||
"subscribed_at": datetime.utcnow()
|
||||
}
|
||||
|
||||
logger.info(f"Updated subscription for client {client_id}")
|
||||
|
||||
# Send confirmation
|
||||
confirmation = {
|
||||
"type": "subscription_updated",
|
||||
"client_id": client_id,
|
||||
"config": self.subscribers[client_id],
|
||||
"timestamp": datetime.utcnow().isoformat()
|
||||
}
|
||||
|
||||
await self.connection_manager.send_to_client(client_id, confirmation)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error handling client subscription: {e}")
|
||||
|
||||
async def handle_client_disconnect(self, client_id: str):
|
||||
"""Handle client disconnection."""
|
||||
if client_id in self.subscribers:
|
||||
del self.subscribers[client_id]
|
||||
logger.info(f"Removed subscription for disconnected client {client_id}")
|
||||
|
||||
async def send_historical_data(
|
||||
self,
|
||||
client_id: str,
|
||||
zone_id: str,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
limit: int = 100
|
||||
):
|
||||
"""Send historical pose data to client."""
|
||||
try:
|
||||
# Get historical data from pose service
|
||||
historical_data = await self.pose_service.get_historical_data(
|
||||
zone_id=zone_id,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
limit=limit
|
||||
)
|
||||
|
||||
# Send data in chunks to avoid overwhelming the client
|
||||
chunk_size = 10
|
||||
for i in range(0, len(historical_data), chunk_size):
|
||||
chunk = historical_data[i:i + chunk_size]
|
||||
|
||||
message = {
|
||||
"type": "historical_data",
|
||||
"zone_id": zone_id,
|
||||
"chunk_index": i // chunk_size,
|
||||
"total_chunks": (len(historical_data) + chunk_size - 1) // chunk_size,
|
||||
"data": chunk,
|
||||
"timestamp": datetime.utcnow().isoformat()
|
||||
}
|
||||
|
||||
await self.connection_manager.send_to_client(client_id, message)
|
||||
|
||||
# Small delay between chunks
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Send completion message
|
||||
completion_message = {
|
||||
"type": "historical_data_complete",
|
||||
"zone_id": zone_id,
|
||||
"total_records": len(historical_data),
|
||||
"timestamp": datetime.utcnow().isoformat()
|
||||
}
|
||||
|
||||
await self.connection_manager.send_to_client(client_id, completion_message)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error sending historical data: {e}")
|
||||
|
||||
# Send error message to client
|
||||
error_message = {
|
||||
"type": "error",
|
||||
"message": f"Failed to retrieve historical data: {str(e)}",
|
||||
"timestamp": datetime.utcnow().isoformat()
|
||||
}
|
||||
|
||||
await self.connection_manager.send_to_client(client_id, error_message)
|
||||
|
||||
async def send_zone_statistics(self, client_id: str, zone_id: str):
|
||||
"""Send zone statistics to client."""
|
||||
try:
|
||||
# Get zone statistics
|
||||
stats = await self.pose_service.get_zone_statistics(zone_id)
|
||||
|
||||
message = {
|
||||
"type": "zone_statistics",
|
||||
"zone_id": zone_id,
|
||||
"statistics": stats,
|
||||
"timestamp": datetime.utcnow().isoformat()
|
||||
}
|
||||
|
||||
await self.connection_manager.send_to_client(client_id, message)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error sending zone statistics: {e}")
|
||||
|
||||
async def broadcast_system_event(self, event_type: str, event_data: Dict[str, Any]):
|
||||
"""Broadcast system events to all connected clients."""
|
||||
try:
|
||||
message = {
|
||||
"type": "system_event",
|
||||
"event_type": event_type,
|
||||
"data": event_data,
|
||||
"timestamp": datetime.utcnow().isoformat()
|
||||
}
|
||||
|
||||
# Broadcast to all pose stream clients
|
||||
sent_count = await self.connection_manager.broadcast(
|
||||
data=message,
|
||||
stream_type="pose"
|
||||
)
|
||||
|
||||
logger.info(f"Broadcasted system event '{event_type}' to {sent_count} clients")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error broadcasting system event: {e}")
|
||||
|
||||
async def update_stream_config(self, config: Dict[str, Any]):
|
||||
"""Update streaming configuration."""
|
||||
try:
|
||||
# Validate and update configuration
|
||||
if "fps" in config:
|
||||
fps = max(1, min(60, config["fps"]))
|
||||
self.stream_config["fps"] = fps
|
||||
|
||||
if "min_confidence" in config:
|
||||
confidence = max(0.0, min(1.0, config["min_confidence"]))
|
||||
self.stream_config["min_confidence"] = confidence
|
||||
|
||||
if "include_metadata" in config:
|
||||
self.stream_config["include_metadata"] = bool(config["include_metadata"])
|
||||
|
||||
if "buffer_size" in config:
|
||||
buffer_size = max(10, min(1000, config["buffer_size"]))
|
||||
self.stream_config["buffer_size"] = buffer_size
|
||||
|
||||
logger.info(f"Updated stream configuration: {self.stream_config}")
|
||||
|
||||
# Broadcast configuration update to clients
|
||||
await self.broadcast_system_event("stream_config_updated", {
|
||||
"new_config": self.stream_config
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error updating stream configuration: {e}")
|
||||
|
||||
def get_stream_status(self) -> Dict[str, Any]:
|
||||
"""Get current streaming status."""
|
||||
return {
|
||||
"is_streaming": self.is_streaming,
|
||||
"config": self.stream_config,
|
||||
"subscriber_count": len(self.subscribers),
|
||||
"subscribers": {
|
||||
client_id: {
|
||||
"zone_ids": sub["zone_ids"],
|
||||
"min_confidence": sub["min_confidence"],
|
||||
"subscribed_at": sub["subscribed_at"].isoformat()
|
||||
}
|
||||
for client_id, sub in self.subscribers.items()
|
||||
}
|
||||
}
|
||||
|
||||
async def get_performance_metrics(self) -> Dict[str, Any]:
|
||||
"""Get streaming performance metrics."""
|
||||
try:
|
||||
# Get connection manager metrics
|
||||
conn_metrics = await self.connection_manager.get_metrics()
|
||||
|
||||
# Get pose service metrics
|
||||
pose_metrics = await self.pose_service.get_performance_metrics()
|
||||
|
||||
return {
|
||||
"streaming": {
|
||||
"is_active": self.is_streaming,
|
||||
"fps": self.stream_config["fps"],
|
||||
"subscriber_count": len(self.subscribers)
|
||||
},
|
||||
"connections": conn_metrics,
|
||||
"pose_service": pose_metrics,
|
||||
"timestamp": datetime.utcnow().isoformat()
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting performance metrics: {e}")
|
||||
return {}
|
||||
|
||||
async def shutdown(self):
|
||||
"""Shutdown pose stream handler."""
|
||||
await self.stop_streaming()
|
||||
self.subscribers.clear()
|
||||
logger.info("Pose stream handler shutdown complete")
|
||||
Reference in New Issue
Block a user