134 lines
5.6 KiB
Python
134 lines
5.6 KiB
Python
from typing import Dict, Set
|
|
from fastapi import WebSocket
|
|
from app.core.json_helpers import _dumps
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class WebSocketManager:
|
|
def __init__(self):
|
|
self.active_connections: Dict[str, WebSocket] = {}
|
|
# 클라이언트별 구독 종목 세트
|
|
self._subscriptions: Dict[str, Set[str]] = {}
|
|
|
|
# ------------------------------------------------------------------
|
|
# 연결 관리
|
|
# ------------------------------------------------------------------
|
|
async def connect(self, websocket: WebSocket, client_id: str):
|
|
"""클라이언트 연결"""
|
|
await websocket.accept()
|
|
self.active_connections[client_id] = websocket
|
|
self._subscriptions.setdefault(client_id, set())
|
|
logger.info(f"Client {client_id} connected. Total: {len(self.active_connections)}")
|
|
|
|
def disconnect(self, client_id: str):
|
|
"""클라이언트 연결 해제"""
|
|
self.active_connections.pop(client_id, None)
|
|
self._subscriptions.pop(client_id, None)
|
|
logger.info(f"Client {client_id} disconnected. Total: {len(self.active_connections)}")
|
|
|
|
# ------------------------------------------------------------------
|
|
# 구독 관리
|
|
# ------------------------------------------------------------------
|
|
def subscribe(self, client_id: str, stock_codes: list[str]):
|
|
"""클라이언트에 종목 추가 구독."""
|
|
bucket = self._subscriptions.setdefault(client_id, set())
|
|
bucket.update(stock_codes)
|
|
logger.debug(f"[sub+] {client_id}: {stock_codes} → 현재 {sorted(bucket)}")
|
|
|
|
def unsubscribe(self, client_id: str, stock_codes: list[str]):
|
|
"""클라이언트의 특정 종목 구독 해제."""
|
|
bucket = self._subscriptions.get(client_id, set())
|
|
bucket.difference_update(stock_codes)
|
|
logger.debug(f"[sub-] {client_id}: {stock_codes} → 현재 {sorted(bucket)}")
|
|
|
|
def set_subscriptions(self, client_id: str, stock_codes: list[str]):
|
|
"""클라이언트의 구독 목록을 전체 교체."""
|
|
self._subscriptions[client_id] = set(stock_codes)
|
|
logger.debug(f"[sub=] {client_id}: 구독 목록 재설정 → {sorted(stock_codes)}")
|
|
|
|
def get_subscriptions(self, client_id: str) -> list[str]:
|
|
"""클라이언트의 현재 구독 목록 반환."""
|
|
return sorted(self._subscriptions.get(client_id, set()))
|
|
|
|
def get_all_subscribed_codes(self) -> Set[str]:
|
|
"""모든 활성 클라이언트에 걸쳐 구독 중인 종목 합집합 반환.
|
|
KisRealtimeStream이 이 목록을 KIS 측에 구독한다.
|
|
"""
|
|
result: Set[str] = set()
|
|
for client_id in self.active_connections:
|
|
result.update(self._subscriptions.get(client_id, set()))
|
|
return result
|
|
|
|
# ------------------------------------------------------------------
|
|
# 메시지 전송
|
|
# ------------------------------------------------------------------
|
|
async def send_personal_message(self, message: str, client_id: str):
|
|
"""특정 클라이언트에게 텍스트 메시지 전송"""
|
|
ws = self.active_connections.get(client_id)
|
|
if ws:
|
|
try:
|
|
await ws.send_text(message)
|
|
except Exception as e:
|
|
logger.error(f"Error sending to {client_id}: {e}")
|
|
self.disconnect(client_id)
|
|
|
|
async def broadcast(self, message: str):
|
|
"""모든 연결된 클라이언트에게 메시지 브로드캐스트"""
|
|
disconnected = []
|
|
for client_id, ws in list(self.active_connections.items()):
|
|
try:
|
|
await ws.send_text(message)
|
|
except Exception as e:
|
|
logger.error(f"Error broadcasting to {client_id}: {e}")
|
|
disconnected.append(client_id)
|
|
for client_id in disconnected:
|
|
self.disconnect(client_id)
|
|
|
|
async def broadcast_json(self, data: dict):
|
|
"""JSON 데이터를 모든 클라이언트에게 브로드캐스트"""
|
|
await self.broadcast(_dumps(data))
|
|
|
|
async def send_market_update(self, stock_code: str, data: dict):
|
|
"""시장 데이터를 해당 종목을 구독한 클라이언트에게만 전송.
|
|
구독자가 없으면 전체 브로드캐스트(기존 동작 유지).
|
|
"""
|
|
message = _dumps({
|
|
"type": "market_update",
|
|
"stock_code": stock_code,
|
|
"data": data,
|
|
})
|
|
# 해당 종목 구독 클라이언트 필터링
|
|
targets = [
|
|
cid for cid, codes in self._subscriptions.items()
|
|
if stock_code in codes and cid in self.active_connections
|
|
]
|
|
if not targets:
|
|
# 구독자 없음 — 전체 브로드캐스트(fallback)
|
|
await self.broadcast(message)
|
|
return
|
|
disconnected = []
|
|
for cid in targets:
|
|
ws = self.active_connections.get(cid)
|
|
if not ws:
|
|
continue
|
|
try:
|
|
await ws.send_text(message)
|
|
except Exception as e:
|
|
logger.error(f"send_market_update error → {cid}: {e}")
|
|
disconnected.append(cid)
|
|
for cid in disconnected:
|
|
self.disconnect(cid)
|
|
|
|
async def send_trade_notification(self, trade_data: dict):
|
|
"""거래 알림 전송"""
|
|
await self.broadcast_json({"type": "trade_notification", "data": trade_data})
|
|
|
|
async def send_mcp_status(self, status: dict):
|
|
"""MCP 상태 업데이트 전송"""
|
|
await self.broadcast_json({"type": "mcp_status", "data": status})
|
|
|
|
|
|
websocket_manager = WebSocketManager()
|