""" WebSocket 연결 관리 모듈 - 클라이언트 연결/해제 및 인증 관리 - 종목별 구독(subscribe/unsubscribe/set_subscriptions) 관리 - 개별 클라이언트, 전체, 종목 구독자 타겟 메시지 전송 - 싱글톤 websocket_manager 인스턴스 제공 (main.py 에서 import) """ from typing import Dict, Set import asyncio from fastapi import WebSocket from app.core.json_helpers import _dumps import logging logger = logging.getLogger(__name__) # [L-6] 실시간 시세 전송 1건당 클라이언트별 타임아웃(초). 초과 = 정체된 클라이언트 → 연결 해제. _MARKET_UPDATE_SEND_TIMEOUT = 5.0 class WebSocketManager: def __init__(self): self.active_connections: Dict[str, WebSocket] = {} # 클라이언트별 구독 종목 세트 self._subscriptions: Dict[str, Set[str]] = {} # 클라이언트별 "선택 종목"(차트를 보고 있는 종목) — KIS 등록 우선순위·호가 구독 대상(H-4) self._selected: Dict[str, Set[str]] = {} # ------------------------------------------------------------------ # 연결 관리 # ------------------------------------------------------------------ async def connect(self, websocket: WebSocket, client_id: str): """클라이언트 연결 (서브프로토콜 없음 — connect_v2 위임).""" await self.connect_v2(websocket, client_id) async def connect_v2(self, websocket: WebSocket, client_id: str, subprotocol: str | None = None): """클라이언트 연결. subprotocol 이 있으면 accept(subprotocol=...) 로 응답 (브라우저는 요청한 서브프로토콜이 응답에 없으면 연결을 끊는다). 동일 client_id로 재접속 시 active_connections를 교체. 이전 핸들러는 websocket_endpoint의 루프에서 교체 감지 후 자연 종료. 구독 목록은 유지 (새 연결이 즉시 실시간 데이터를 수신할 수 있도록). """ if client_id in self.active_connections: logger.info(f"Client {client_id}: 기존 연결 교체 (이전 핸들러 자연 종료 대기)") await websocket.accept(subprotocol=subprotocol) 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) self._selected.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 set_selected_codes_v2(self, client_id: str, stock_codes: list[str]): """클라이언트의 선택 종목(차트 표시 종목)을 전체 교체.""" self._selected[client_id] = {str(c) for c in stock_codes if c} def get_selected_codes_v2(self) -> Set[str]: """활성 클라이언트의 선택 종목 합집합 (해당 클라이언트가 실제 구독 중인 종목만).""" result: Set[str] = set() for client_id in self.active_connections: subs = self._subscriptions.get(client_id, set()) result.update(c for c in self._selected.get(client_id, set()) if c in subs) return result 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): """시장 데이터를 해당 종목을 구독한 클라이언트에게만 전송. [L-5] 구독자가 없으면 아무에게도 보내지 않는다(과거: 전체 브로드캐스트 폴백 — 구독하지 않은 종목 틱이 모든 화면에 쏟아졌다). [L-6] 클라이언트별 전송을 병렬 + 타임아웃으로 — 느린 클라이언트 1개가 KIS 분배를 막지 않게. """ targets = [ cid for cid, codes in self._subscriptions.items() if stock_code in codes and cid in self.active_connections ] if not targets: return message = _dumps({ "type": "market_update", "stock_code": stock_code, "data": data, }) async def _send_one(cid: str): ws = self.active_connections.get(cid) if not ws: return None try: await asyncio.wait_for(ws.send_text(message), timeout=_MARKET_UPDATE_SEND_TIMEOUT) return None except asyncio.TimeoutError: logger.warning(f"send_market_update 타임아웃({_MARKET_UPDATE_SEND_TIMEOUT}s) → {cid} 연결 해제") return (cid, ws) except Exception as e: logger.error(f"send_market_update error → {cid}: {e}") return (cid, ws) results = await asyncio.gather(*(_send_one(cid) for cid in targets)) for item in results: if not item: continue cid, ws = item # 전송 중 같은 client_id 로 재접속했으면 새 연결은 건드리지 않는다 if self.active_connections.get(cid) is ws: 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()