replay system
This commit is contained in:
Dobromir Popov
306
COBY/api/replay_api.py
Normal file
306
COBY/api/replay_api.py
Normal file
@ -0,0 +1,306 @@
|
||||
"""
|
||||
REST API endpoints for historical data replay functionality.
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Query, Path
|
||||
from typing import Optional, List, Dict, Any
|
||||
from datetime import datetime
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from ..replay.replay_manager import HistoricalReplayManager
|
||||
from ..models.core import ReplayStatus
|
||||
from ..utils.logging import get_logger, set_correlation_id
|
||||
from ..utils.exceptions import ReplayError, ValidationError
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class CreateReplayRequest(BaseModel):
|
||||
"""Request model for creating replay session"""
|
||||
start_time: datetime = Field(..., description="Replay start time")
|
||||
end_time: datetime = Field(..., description="Replay end time")
|
||||
speed: float = Field(1.0, gt=0, le=100, description="Playback speed multiplier")
|
||||
symbols: Optional[List[str]] = Field(None, description="Symbols to replay")
|
||||
exchanges: Optional[List[str]] = Field(None, description="Exchanges to replay")
|
||||
|
||||
|
||||
class ReplayControlRequest(BaseModel):
|
||||
"""Request model for replay control operations"""
|
||||
action: str = Field(..., description="Control action: start, pause, resume, stop")
|
||||
|
||||
|
||||
class SeekRequest(BaseModel):
|
||||
"""Request model for seeking in replay"""
|
||||
timestamp: datetime = Field(..., description="Target timestamp")
|
||||
|
||||
|
||||
class SpeedRequest(BaseModel):
|
||||
"""Request model for changing replay speed"""
|
||||
speed: float = Field(..., gt=0, le=100, description="New playback speed")
|
||||
|
||||
|
||||
def create_replay_router(replay_manager: HistoricalReplayManager) -> APIRouter:
|
||||
"""Create replay API router with endpoints"""
|
||||
|
||||
router = APIRouter(prefix="/replay", tags=["replay"])
|
||||
|
||||
@router.post("/sessions", response_model=Dict[str, str])
|
||||
async def create_replay_session(request: CreateReplayRequest):
|
||||
"""Create a new replay session"""
|
||||
try:
|
||||
set_correlation_id()
|
||||
|
||||
session_id = replay_manager.create_replay_session(
|
||||
start_time=request.start_time,
|
||||
end_time=request.end_time,
|
||||
speed=request.speed,
|
||||
symbols=request.symbols,
|
||||
exchanges=request.exchanges
|
||||
)
|
||||
|
||||
logger.info(f"Created replay session {session_id}")
|
||||
|
||||
return {
|
||||
"session_id": session_id,
|
||||
"status": "created",
|
||||
"message": "Replay session created successfully"
|
||||
}
|
||||
|
||||
except ValidationError as e:
|
||||
logger.warning(f"Invalid replay request: {e}")
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except ReplayError as e:
|
||||
logger.error(f"Replay creation failed: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error creating replay session: {e}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
@router.get("/sessions", response_model=List[Dict[str, Any]])
|
||||
async def list_replay_sessions():
|
||||
"""List all replay sessions"""
|
||||
try:
|
||||
sessions = replay_manager.list_replay_sessions()
|
||||
|
||||
return [
|
||||
{
|
||||
"session_id": session.session_id,
|
||||
"start_time": session.start_time.isoformat(),
|
||||
"end_time": session.end_time.isoformat(),
|
||||
"current_time": session.current_time.isoformat(),
|
||||
"speed": session.speed,
|
||||
"status": session.status.value,
|
||||
"symbols": session.symbols,
|
||||
"exchanges": session.exchanges,
|
||||
"progress": session.progress,
|
||||
"events_replayed": session.events_replayed,
|
||||
"total_events": session.total_events,
|
||||
"created_at": session.created_at.isoformat(),
|
||||
"started_at": session.started_at.isoformat() if session.started_at else None,
|
||||
"error_message": getattr(session, 'error_message', None)
|
||||
}
|
||||
for session in sessions
|
||||
]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing replay sessions: {e}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
@router.get("/sessions/{session_id}", response_model=Dict[str, Any])
|
||||
async def get_replay_session(session_id: str = Path(..., description="Session ID")):
|
||||
"""Get replay session details"""
|
||||
try:
|
||||
session = replay_manager.get_replay_status(session_id)
|
||||
|
||||
if not session:
|
||||
raise HTTPException(status_code=404, detail="Session not found")
|
||||
|
||||
return {
|
||||
"session_id": session.session_id,
|
||||
"start_time": session.start_time.isoformat(),
|
||||
"end_time": session.end_time.isoformat(),
|
||||
"current_time": session.current_time.isoformat(),
|
||||
"speed": session.speed,
|
||||
"status": session.status.value,
|
||||
"symbols": session.symbols,
|
||||
"exchanges": session.exchanges,
|
||||
"progress": session.progress,
|
||||
"events_replayed": session.events_replayed,
|
||||
"total_events": session.total_events,
|
||||
"created_at": session.created_at.isoformat(),
|
||||
"started_at": session.started_at.isoformat() if session.started_at else None,
|
||||
"paused_at": session.paused_at.isoformat() if session.paused_at else None,
|
||||
"stopped_at": session.stopped_at.isoformat() if session.stopped_at else None,
|
||||
"completed_at": session.completed_at.isoformat() if session.completed_at else None,
|
||||
"error_message": getattr(session, 'error_message', None)
|
||||
}
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting replay session {session_id}: {e}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
@router.post("/sessions/{session_id}/control", response_model=Dict[str, str])
|
||||
async def control_replay_session(
|
||||
session_id: str = Path(..., description="Session ID"),
|
||||
request: ReplayControlRequest = None
|
||||
):
|
||||
"""Control replay session (start, pause, resume, stop)"""
|
||||
try:
|
||||
set_correlation_id()
|
||||
|
||||
if not request:
|
||||
raise HTTPException(status_code=400, detail="Control action required")
|
||||
|
||||
action = request.action.lower()
|
||||
|
||||
if action == "start":
|
||||
await replay_manager.start_replay(session_id)
|
||||
message = "Replay started"
|
||||
elif action == "pause":
|
||||
await replay_manager.pause_replay(session_id)
|
||||
message = "Replay paused"
|
||||
elif action == "resume":
|
||||
await replay_manager.resume_replay(session_id)
|
||||
message = "Replay resumed"
|
||||
elif action == "stop":
|
||||
await replay_manager.stop_replay(session_id)
|
||||
message = "Replay stopped"
|
||||
else:
|
||||
raise HTTPException(status_code=400, detail="Invalid action")
|
||||
|
||||
logger.info(f"Replay session {session_id} action: {action}")
|
||||
|
||||
return {
|
||||
"session_id": session_id,
|
||||
"action": action,
|
||||
"message": message
|
||||
}
|
||||
|
||||
except ReplayError as e:
|
||||
logger.error(f"Replay control failed for {session_id}: {e}")
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error controlling replay {session_id}: {e}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
@router.post("/sessions/{session_id}/seek", response_model=Dict[str, str])
|
||||
async def seek_replay_session(
|
||||
session_id: str = Path(..., description="Session ID"),
|
||||
request: SeekRequest = None
|
||||
):
|
||||
"""Seek to specific timestamp in replay"""
|
||||
try:
|
||||
if not request:
|
||||
raise HTTPException(status_code=400, detail="Timestamp required")
|
||||
|
||||
success = replay_manager.seek_replay(session_id, request.timestamp)
|
||||
|
||||
if not success:
|
||||
raise HTTPException(status_code=400, detail="Seek failed")
|
||||
|
||||
logger.info(f"Seeked replay session {session_id} to {request.timestamp}")
|
||||
|
||||
return {
|
||||
"session_id": session_id,
|
||||
"timestamp": request.timestamp.isoformat(),
|
||||
"message": "Seek successful"
|
||||
}
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error seeking replay session {session_id}: {e}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
@router.post("/sessions/{session_id}/speed", response_model=Dict[str, Any])
|
||||
async def set_replay_speed(
|
||||
session_id: str = Path(..., description="Session ID"),
|
||||
request: SpeedRequest = None
|
||||
):
|
||||
"""Change replay speed"""
|
||||
try:
|
||||
if not request:
|
||||
raise HTTPException(status_code=400, detail="Speed required")
|
||||
|
||||
success = replay_manager.set_replay_speed(session_id, request.speed)
|
||||
|
||||
if not success:
|
||||
raise HTTPException(status_code=400, detail="Speed change failed")
|
||||
|
||||
logger.info(f"Set replay speed to {request.speed}x for session {session_id}")
|
||||
|
||||
return {
|
||||
"session_id": session_id,
|
||||
"speed": request.speed,
|
||||
"message": "Speed changed successfully"
|
||||
}
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error setting replay speed for {session_id}: {e}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
@router.delete("/sessions/{session_id}", response_model=Dict[str, str])
|
||||
async def delete_replay_session(session_id: str = Path(..., description="Session ID")):
|
||||
"""Delete replay session"""
|
||||
try:
|
||||
success = replay_manager.delete_replay_session(session_id)
|
||||
|
||||
if not success:
|
||||
raise HTTPException(status_code=404, detail="Session not found")
|
||||
|
||||
logger.info(f"Deleted replay session {session_id}")
|
||||
|
||||
return {
|
||||
"session_id": session_id,
|
||||
"message": "Session deleted successfully"
|
||||
}
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting replay session {session_id}: {e}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
@router.get("/data-range/{symbol}", response_model=Dict[str, Any])
|
||||
async def get_data_range(
|
||||
symbol: str = Path(..., description="Trading symbol"),
|
||||
exchange: Optional[str] = Query(None, description="Exchange name")
|
||||
):
|
||||
"""Get available data time range for a symbol"""
|
||||
try:
|
||||
data_range = await replay_manager.get_available_data_range(symbol, exchange)
|
||||
|
||||
if not data_range:
|
||||
raise HTTPException(status_code=404, detail="No data available for symbol")
|
||||
|
||||
return {
|
||||
"symbol": symbol,
|
||||
"exchange": exchange,
|
||||
"start_time": data_range['start'].isoformat(),
|
||||
"end_time": data_range['end'].isoformat(),
|
||||
"duration_days": (data_range['end'] - data_range['start']).days
|
||||
}
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting data range for {symbol}: {e}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
@router.get("/stats", response_model=Dict[str, Any])
|
||||
async def get_replay_stats():
|
||||
"""Get replay system statistics"""
|
||||
try:
|
||||
return replay_manager.get_stats()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting replay stats: {e}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
return router
|
||||
435
COBY/api/replay_websocket.py
Normal file
435
COBY/api/replay_websocket.py
Normal file
@ -0,0 +1,435 @@
|
||||
"""
|
||||
WebSocket server for real-time replay data streaming.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from typing import Dict, Set, Optional, Any
|
||||
from fastapi import WebSocket, WebSocketDisconnect
|
||||
from datetime import datetime
|
||||
|
||||
from ..replay.replay_manager import HistoricalReplayManager
|
||||
from ..models.core import OrderBookSnapshot, TradeEvent, ReplayStatus
|
||||
from ..utils.logging import get_logger, set_correlation_id
|
||||
from ..utils.exceptions import ReplayError
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class ReplayWebSocketManager:
|
||||
"""
|
||||
WebSocket manager for replay data streaming.
|
||||
|
||||
Provides:
|
||||
- Real-time replay data streaming
|
||||
- Session-based connections
|
||||
- Automatic cleanup on disconnect
|
||||
- Status updates
|
||||
"""
|
||||
|
||||
def __init__(self, replay_manager: HistoricalReplayManager):
|
||||
"""
|
||||
Initialize WebSocket manager.
|
||||
|
||||
Args:
|
||||
replay_manager: Replay manager instance
|
||||
"""
|
||||
self.replay_manager = replay_manager
|
||||
|
||||
# Connection management
|
||||
self.connections: Dict[str, Set[WebSocket]] = {} # session_id -> websockets
|
||||
self.websocket_sessions: Dict[WebSocket, str] = {} # websocket -> session_id
|
||||
|
||||
# Statistics
|
||||
self.stats = {
|
||||
'active_connections': 0,
|
||||
'total_connections': 0,
|
||||
'messages_sent': 0,
|
||||
'connection_errors': 0
|
||||
}
|
||||
|
||||
logger.info("Replay WebSocket manager initialized")
|
||||
|
||||
async def connect_to_session(self, websocket: WebSocket, session_id: str) -> bool:
|
||||
"""
|
||||
Connect WebSocket to a replay session.
|
||||
|
||||
Args:
|
||||
websocket: WebSocket connection
|
||||
session_id: Replay session ID
|
||||
|
||||
Returns:
|
||||
bool: True if connected successfully, False otherwise
|
||||
"""
|
||||
try:
|
||||
set_correlation_id()
|
||||
|
||||
# Check if session exists
|
||||
session = self.replay_manager.get_replay_status(session_id)
|
||||
if not session:
|
||||
await websocket.send_json({
|
||||
"type": "error",
|
||||
"message": f"Session {session_id} not found"
|
||||
})
|
||||
return False
|
||||
|
||||
# Accept WebSocket connection
|
||||
await websocket.accept()
|
||||
|
||||
# Add to connection tracking
|
||||
if session_id not in self.connections:
|
||||
self.connections[session_id] = set()
|
||||
|
||||
self.connections[session_id].add(websocket)
|
||||
self.websocket_sessions[websocket] = session_id
|
||||
|
||||
# Update statistics
|
||||
self.stats['active_connections'] += 1
|
||||
self.stats['total_connections'] += 1
|
||||
|
||||
# Add callbacks to replay session
|
||||
self.replay_manager.add_data_callback(session_id, self._data_callback)
|
||||
self.replay_manager.add_status_callback(session_id, self._status_callback)
|
||||
|
||||
# Send initial session status
|
||||
await self._send_session_status(websocket, session)
|
||||
|
||||
logger.info(f"WebSocket connected to replay session {session_id}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to connect WebSocket to session {session_id}: {e}")
|
||||
self.stats['connection_errors'] += 1
|
||||
return False
|
||||
|
||||
async def disconnect(self, websocket: WebSocket) -> None:
|
||||
"""
|
||||
Disconnect WebSocket and cleanup.
|
||||
|
||||
Args:
|
||||
websocket: WebSocket connection to disconnect
|
||||
"""
|
||||
try:
|
||||
session_id = self.websocket_sessions.get(websocket)
|
||||
|
||||
if session_id:
|
||||
# Remove from connection tracking
|
||||
if session_id in self.connections:
|
||||
self.connections[session_id].discard(websocket)
|
||||
|
||||
# Clean up empty session connections
|
||||
if not self.connections[session_id]:
|
||||
del self.connections[session_id]
|
||||
|
||||
del self.websocket_sessions[websocket]
|
||||
|
||||
# Update statistics
|
||||
self.stats['active_connections'] -= 1
|
||||
|
||||
logger.info(f"WebSocket disconnected from replay session {session_id}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error during WebSocket disconnect: {e}")
|
||||
|
||||
async def handle_websocket_messages(self, websocket: WebSocket) -> None:
|
||||
"""
|
||||
Handle incoming WebSocket messages.
|
||||
|
||||
Args:
|
||||
websocket: WebSocket connection
|
||||
"""
|
||||
try:
|
||||
while True:
|
||||
# Receive message
|
||||
message = await websocket.receive_json()
|
||||
|
||||
# Process message
|
||||
await self._process_websocket_message(websocket, message)
|
||||
|
||||
except WebSocketDisconnect:
|
||||
logger.info("WebSocket disconnected")
|
||||
except Exception as e:
|
||||
logger.error(f"WebSocket message handling error: {e}")
|
||||
await websocket.send_json({
|
||||
"type": "error",
|
||||
"message": "Message processing error"
|
||||
})
|
||||
|
||||
async def _process_websocket_message(self, websocket: WebSocket, message: Dict[str, Any]) -> None:
|
||||
"""
|
||||
Process incoming WebSocket message.
|
||||
|
||||
Args:
|
||||
websocket: WebSocket connection
|
||||
message: Received message
|
||||
"""
|
||||
try:
|
||||
message_type = message.get('type')
|
||||
session_id = self.websocket_sessions.get(websocket)
|
||||
|
||||
if not session_id:
|
||||
await websocket.send_json({
|
||||
"type": "error",
|
||||
"message": "Not connected to any session"
|
||||
})
|
||||
return
|
||||
|
||||
if message_type == "control":
|
||||
await self._handle_control_message(websocket, session_id, message)
|
||||
elif message_type == "seek":
|
||||
await self._handle_seek_message(websocket, session_id, message)
|
||||
elif message_type == "speed":
|
||||
await self._handle_speed_message(websocket, session_id, message)
|
||||
elif message_type == "status":
|
||||
await self._handle_status_request(websocket, session_id)
|
||||
else:
|
||||
await websocket.send_json({
|
||||
"type": "error",
|
||||
"message": f"Unknown message type: {message_type}"
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing WebSocket message: {e}")
|
||||
await websocket.send_json({
|
||||
"type": "error",
|
||||
"message": "Message processing failed"
|
||||
})
|
||||
|
||||
async def _handle_control_message(self, websocket: WebSocket, session_id: str,
|
||||
message: Dict[str, Any]) -> None:
|
||||
"""Handle replay control messages."""
|
||||
try:
|
||||
action = message.get('action')
|
||||
|
||||
if action == "start":
|
||||
await self.replay_manager.start_replay(session_id)
|
||||
elif action == "pause":
|
||||
await self.replay_manager.pause_replay(session_id)
|
||||
elif action == "resume":
|
||||
await self.replay_manager.resume_replay(session_id)
|
||||
elif action == "stop":
|
||||
await self.replay_manager.stop_replay(session_id)
|
||||
else:
|
||||
await websocket.send_json({
|
||||
"type": "error",
|
||||
"message": f"Invalid control action: {action}"
|
||||
})
|
||||
return
|
||||
|
||||
await websocket.send_json({
|
||||
"type": "control_response",
|
||||
"action": action,
|
||||
"status": "success"
|
||||
})
|
||||
|
||||
except ReplayError as e:
|
||||
await websocket.send_json({
|
||||
"type": "error",
|
||||
"message": str(e)
|
||||
})
|
||||
except Exception as e:
|
||||
logger.error(f"Control message error: {e}")
|
||||
await websocket.send_json({
|
||||
"type": "error",
|
||||
"message": "Control action failed"
|
||||
})
|
||||
|
||||
async def _handle_seek_message(self, websocket: WebSocket, session_id: str,
|
||||
message: Dict[str, Any]) -> None:
|
||||
"""Handle seek messages."""
|
||||
try:
|
||||
timestamp_str = message.get('timestamp')
|
||||
if not timestamp_str:
|
||||
await websocket.send_json({
|
||||
"type": "error",
|
||||
"message": "Timestamp required for seek"
|
||||
})
|
||||
return
|
||||
|
||||
timestamp = datetime.fromisoformat(timestamp_str.replace('Z', '+00:00'))
|
||||
success = self.replay_manager.seek_replay(session_id, timestamp)
|
||||
|
||||
await websocket.send_json({
|
||||
"type": "seek_response",
|
||||
"timestamp": timestamp_str,
|
||||
"status": "success" if success else "failed"
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Seek message error: {e}")
|
||||
await websocket.send_json({
|
||||
"type": "error",
|
||||
"message": "Seek failed"
|
||||
})
|
||||
|
||||
async def _handle_speed_message(self, websocket: WebSocket, session_id: str,
|
||||
message: Dict[str, Any]) -> None:
|
||||
"""Handle speed change messages."""
|
||||
try:
|
||||
speed = message.get('speed')
|
||||
if not speed or speed <= 0:
|
||||
await websocket.send_json({
|
||||
"type": "error",
|
||||
"message": "Valid speed required"
|
||||
})
|
||||
return
|
||||
|
||||
success = self.replay_manager.set_replay_speed(session_id, speed)
|
||||
|
||||
await websocket.send_json({
|
||||
"type": "speed_response",
|
||||
"speed": speed,
|
||||
"status": "success" if success else "failed"
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Speed message error: {e}")
|
||||
await websocket.send_json({
|
||||
"type": "error",
|
||||
"message": "Speed change failed"
|
||||
})
|
||||
|
||||
async def _handle_status_request(self, websocket: WebSocket, session_id: str) -> None:
|
||||
"""Handle status request messages."""
|
||||
try:
|
||||
session = self.replay_manager.get_replay_status(session_id)
|
||||
if session:
|
||||
await self._send_session_status(websocket, session)
|
||||
else:
|
||||
await websocket.send_json({
|
||||
"type": "error",
|
||||
"message": "Session not found"
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Status request error: {e}")
|
||||
await websocket.send_json({
|
||||
"type": "error",
|
||||
"message": "Status request failed"
|
||||
})
|
||||
|
||||
async def _data_callback(self, data) -> None:
|
||||
"""Callback for replay data - broadcasts to all connected WebSockets."""
|
||||
try:
|
||||
# Determine which session this data belongs to
|
||||
# This is a simplified approach - in practice, you'd need to track
|
||||
# which session generated this callback
|
||||
|
||||
# Serialize data
|
||||
if isinstance(data, OrderBookSnapshot):
|
||||
message = {
|
||||
"type": "orderbook",
|
||||
"data": {
|
||||
"symbol": data.symbol,
|
||||
"exchange": data.exchange,
|
||||
"timestamp": data.timestamp.isoformat(),
|
||||
"bids": [{"price": b.price, "size": b.size} for b in data.bids[:10]],
|
||||
"asks": [{"price": a.price, "size": a.size} for a in data.asks[:10]],
|
||||
"sequence_id": data.sequence_id
|
||||
}
|
||||
}
|
||||
elif isinstance(data, TradeEvent):
|
||||
message = {
|
||||
"type": "trade",
|
||||
"data": {
|
||||
"symbol": data.symbol,
|
||||
"exchange": data.exchange,
|
||||
"timestamp": data.timestamp.isoformat(),
|
||||
"price": data.price,
|
||||
"size": data.size,
|
||||
"side": data.side,
|
||||
"trade_id": data.trade_id
|
||||
}
|
||||
}
|
||||
else:
|
||||
return
|
||||
|
||||
# Broadcast to all connections
|
||||
await self._broadcast_message(message)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Data callback error: {e}")
|
||||
|
||||
async def _status_callback(self, session_id: str, status: ReplayStatus) -> None:
|
||||
"""Callback for replay status changes."""
|
||||
try:
|
||||
message = {
|
||||
"type": "status",
|
||||
"session_id": session_id,
|
||||
"status": status.value,
|
||||
"timestamp": datetime.utcnow().isoformat()
|
||||
}
|
||||
|
||||
# Send to connections for this session
|
||||
if session_id in self.connections:
|
||||
await self._broadcast_to_session(session_id, message)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Status callback error: {e}")
|
||||
|
||||
async def _send_session_status(self, websocket: WebSocket, session) -> None:
|
||||
"""Send session status to WebSocket."""
|
||||
try:
|
||||
message = {
|
||||
"type": "session_status",
|
||||
"data": {
|
||||
"session_id": session.session_id,
|
||||
"status": session.status.value,
|
||||
"progress": session.progress,
|
||||
"current_time": session.current_time.isoformat(),
|
||||
"speed": session.speed,
|
||||
"events_replayed": session.events_replayed,
|
||||
"total_events": session.total_events
|
||||
}
|
||||
}
|
||||
|
||||
await websocket.send_json(message)
|
||||
self.stats['messages_sent'] += 1
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error sending session status: {e}")
|
||||
|
||||
async def _broadcast_message(self, message: Dict[str, Any]) -> None:
|
||||
"""Broadcast message to all connected WebSockets."""
|
||||
disconnected = []
|
||||
|
||||
for session_id, websockets in self.connections.items():
|
||||
for websocket in websockets.copy():
|
||||
try:
|
||||
await websocket.send_json(message)
|
||||
self.stats['messages_sent'] += 1
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to send message to WebSocket: {e}")
|
||||
disconnected.append((session_id, websocket))
|
||||
|
||||
# Clean up disconnected WebSockets
|
||||
for session_id, websocket in disconnected:
|
||||
await self.disconnect(websocket)
|
||||
|
||||
async def _broadcast_to_session(self, session_id: str, message: Dict[str, Any]) -> None:
|
||||
"""Broadcast message to WebSockets connected to a specific session."""
|
||||
if session_id not in self.connections:
|
||||
return
|
||||
|
||||
disconnected = []
|
||||
|
||||
for websocket in self.connections[session_id].copy():
|
||||
try:
|
||||
await websocket.send_json(message)
|
||||
self.stats['messages_sent'] += 1
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to send message to WebSocket: {e}")
|
||||
disconnected.append(websocket)
|
||||
|
||||
# Clean up disconnected WebSockets
|
||||
for websocket in disconnected:
|
||||
await self.disconnect(websocket)
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""Get WebSocket manager statistics."""
|
||||
return {
|
||||
**self.stats,
|
||||
'sessions_with_connections': len(self.connections),
|
||||
'total_websockets': sum(len(ws_set) for ws_set in self.connections.values())
|
||||
}
|
||||
Reference in New Issue
Block a user