69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
|
|
|
from app.core.logger import setup_logger
|
|
from app.services.realtime import realtime_service
|
|
|
|
router = APIRouter(tags=["websocket"])
|
|
logger = setup_logger("ws_router")
|
|
|
|
_connected_clients: set[WebSocket] = set()
|
|
|
|
|
|
async def _broadcast(data: dict) -> None:
|
|
message = json.dumps(data, ensure_ascii=False)
|
|
disconnected: list[WebSocket] = []
|
|
for client in _connected_clients:
|
|
try:
|
|
await client.send_text(message)
|
|
except Exception:
|
|
disconnected.append(client)
|
|
for client in disconnected:
|
|
_connected_clients.discard(client)
|
|
|
|
|
|
@router.websocket("/ws/realtime")
|
|
async def websocket_endpoint(websocket: WebSocket) -> None:
|
|
await websocket.accept()
|
|
_connected_clients.add(websocket)
|
|
logger.info("WebSocket 클라이언트 연결 (총 %d)", len(_connected_clients))
|
|
|
|
realtime_service.on("price", _broadcast)
|
|
realtime_service.on("orderbook", _broadcast)
|
|
|
|
try:
|
|
while True:
|
|
raw = await websocket.receive_text()
|
|
try:
|
|
msg = json.loads(raw)
|
|
action = msg.get("action", "")
|
|
|
|
if action == "subscribe":
|
|
stock_code = msg.get("stock_code", "")
|
|
data_type = msg.get("data_type", "price")
|
|
await realtime_service.subscribe(stock_code, data_type)
|
|
await websocket.send_text(
|
|
json.dumps({"status": "subscribed", "stock_code": stock_code, "data_type": data_type})
|
|
)
|
|
|
|
elif action == "unsubscribe":
|
|
stock_code = msg.get("stock_code", "")
|
|
data_type = msg.get("data_type", "price")
|
|
await realtime_service.unsubscribe(stock_code, data_type)
|
|
await websocket.send_text(
|
|
json.dumps({"status": "unsubscribed", "stock_code": stock_code})
|
|
)
|
|
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
except WebSocketDisconnect:
|
|
pass
|
|
finally:
|
|
_connected_clients.discard(websocket)
|
|
logger.info("WebSocket 클라이언트 해제 (총 %d)", len(_connected_clients))
|