first commit
This commit is contained in:
198
app/services/realtime.py
Normal file
198
app/services/realtime.py
Normal file
@@ -0,0 +1,198 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Callable
|
||||
|
||||
import websockets
|
||||
|
||||
from app.core.auth import token_manager
|
||||
from app.core.config import get_ws_url, settings
|
||||
from app.core.logger import setup_logger
|
||||
from app.core.rate_limiter import RateLimiter
|
||||
|
||||
logger = setup_logger("realtime")
|
||||
|
||||
_SUBSCRIBE_TR_IDS = {
|
||||
"price": "H0STCNT0",
|
||||
"orderbook": "H0STASP0",
|
||||
"trade": "H0STCNI0",
|
||||
}
|
||||
|
||||
|
||||
class RealtimeService:
|
||||
def __init__(self) -> None:
|
||||
self._ws: websockets.WebSocketClientProtocol | None = None
|
||||
self._running = False
|
||||
self._callbacks: dict[str, list[Callable]] = {}
|
||||
self._subscriptions: dict[str, list[str]] = {}
|
||||
self._rate_limiter = RateLimiter(requests_per_second=0.5)
|
||||
self._heartbeat_task: asyncio.Task | None = None
|
||||
self._receive_task: asyncio.Task | None = None
|
||||
|
||||
def on(self, event: str, callback: Callable) -> None:
|
||||
self._callbacks.setdefault(event, []).append(callback)
|
||||
|
||||
async def connect(self) -> None:
|
||||
approval_key = await token_manager.get_approval_key()
|
||||
url = get_ws_url()
|
||||
|
||||
try:
|
||||
self._ws = await websockets.connect(
|
||||
url,
|
||||
extra_headers={"approval_key": approval_key, "type": "Y"},
|
||||
)
|
||||
self._running = True
|
||||
self._receive_task = asyncio.create_task(self._receive_loop())
|
||||
self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
|
||||
logger.info("WebSocket 연결 성공: %s", url)
|
||||
except Exception as e:
|
||||
logger.error("WebSocket 연결 실패: %s", e)
|
||||
raise
|
||||
|
||||
async def subscribe(self, stock_code: str, data_type: str = "price") -> None:
|
||||
if not self._ws:
|
||||
await self.connect()
|
||||
|
||||
tr_id = _SUBSCRIBE_TR_IDS.get(data_type, "H0STCNT0")
|
||||
msg = {
|
||||
"header": {
|
||||
"approval_key": await token_manager.get_approval_key(),
|
||||
"custtype": "P",
|
||||
"tr_type": "1",
|
||||
"content-type": "utf-8",
|
||||
},
|
||||
"body": {
|
||||
"input": {
|
||||
"tr_id": tr_id,
|
||||
"tr_key": stock_code,
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
await self._ws.send(json.dumps(msg))
|
||||
self._subscriptions.setdefault(data_type, []).append(stock_code)
|
||||
logger.info("구독 요청: %s - %s", data_type, stock_code)
|
||||
|
||||
async def unsubscribe(self, stock_code: str, data_type: str = "price") -> None:
|
||||
if not self._ws:
|
||||
return
|
||||
|
||||
tr_id = _SUBSCRIBE_TR_IDS.get(data_type, "H0STCNT0")
|
||||
msg = {
|
||||
"header": {
|
||||
"approval_key": await token_manager.get_approval_key(),
|
||||
"custtype": "P",
|
||||
"tr_type": "2",
|
||||
"content-type": "utf-8",
|
||||
},
|
||||
"body": {
|
||||
"input": {
|
||||
"tr_id": tr_id,
|
||||
"tr_key": stock_code,
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
await self._ws.send(json.dumps(msg))
|
||||
subs = self._subscriptions.get(data_type, [])
|
||||
if stock_code in subs:
|
||||
subs.remove(stock_code)
|
||||
logger.info("구독 해제: %s - %s", data_type, stock_code)
|
||||
|
||||
async def _receive_loop(self) -> None:
|
||||
try:
|
||||
async for raw_msg in self._ws:
|
||||
await self._handle_message(raw_msg)
|
||||
except websockets.ConnectionClosed:
|
||||
logger.warning("WebSocket 연결 끊김. 5초 후 재연결...")
|
||||
if self._running:
|
||||
await asyncio.sleep(5)
|
||||
await self.connect()
|
||||
except Exception as e:
|
||||
logger.error("WebSocket 수신 오류: %s", e)
|
||||
|
||||
async def _handle_message(self, raw_msg: str) -> None:
|
||||
try:
|
||||
msg = json.loads(raw_msg)
|
||||
header = msg.get("header", {})
|
||||
body = msg.get("body", {})
|
||||
|
||||
tr_id = header.get("tr_id", "")
|
||||
|
||||
if tr_id == "PINGPONG":
|
||||
return
|
||||
|
||||
parsed = self._parse_data(tr_id, body.get("output", {}))
|
||||
if not parsed:
|
||||
return
|
||||
|
||||
event_type = "price"
|
||||
if "bid_prices" in parsed:
|
||||
event_type = "orderbook"
|
||||
elif "trade_price" in parsed:
|
||||
event_type = "trade"
|
||||
|
||||
for callback in self._callbacks.get(event_type, []):
|
||||
try:
|
||||
result = callback(parsed)
|
||||
if asyncio.iscoroutine(result):
|
||||
await result
|
||||
except Exception as e:
|
||||
logger.error("콜백 실행 오류: %s", e)
|
||||
|
||||
except json.JSONDecodeError:
|
||||
logger.warning("JSON 파싱 실패")
|
||||
|
||||
def _parse_data(self, tr_id: str, output: dict) -> dict | None:
|
||||
if not output:
|
||||
return None
|
||||
|
||||
if tr_id == "H0STCNT0":
|
||||
return {
|
||||
"type": "price",
|
||||
"stock_code": output.get("mksc_shrn_iscd", ""),
|
||||
"trade_price": int(output.get("stck_prpr", 0)),
|
||||
"change_price": int(output.get("prdy_vrss", 0)),
|
||||
"change_rate": float(output.get("prdy_ctrt", 0)),
|
||||
"open_price": int(output.get("stck_oprc", 0)),
|
||||
"high_price": int(output.get("stck_hgpr", 0)),
|
||||
"low_price": int(output.get("stck_lwpr", 0)),
|
||||
"volume": int(output.get("acml_vol", 0)),
|
||||
"trade_time": output.get("stck_cntg_hour", ""),
|
||||
}
|
||||
|
||||
if tr_id == "H0STASP0":
|
||||
return {
|
||||
"type": "orderbook",
|
||||
"stock_code": output.get("mksc_shrn_iscd", ""),
|
||||
"bid_prices": [int(output.get(f"phsc_kprc_{i}", 0)) for i in range(1, 6)],
|
||||
"ask_prices": [int(output.get(f"sats_kprc_{i}", 0)) for i in range(1, 6)],
|
||||
"bid_volumes": [int(output.get(f"phsc_vola_{i}", 0)) for i in range(1, 6)],
|
||||
"ask_volumes": [int(output.get(f"sats_ac_vola_{i}", 0)) for i in range(1, 6)],
|
||||
}
|
||||
|
||||
return None
|
||||
|
||||
async def _heartbeat_loop(self) -> None:
|
||||
while self._running:
|
||||
try:
|
||||
await asyncio.sleep(30)
|
||||
if self._ws and self._ws.open:
|
||||
await self._ws.send(json.dumps({"header": {"tr_id": "PINGPONG"}}))
|
||||
except Exception:
|
||||
break
|
||||
|
||||
async def disconnect(self) -> None:
|
||||
self._running = False
|
||||
if self._heartbeat_task:
|
||||
self._heartbeat_task.cancel()
|
||||
if self._receive_task:
|
||||
self._receive_task.cancel()
|
||||
if self._ws:
|
||||
await self._ws.close()
|
||||
self._ws = None
|
||||
logger.info("WebSocket 연결 종료")
|
||||
|
||||
|
||||
realtime_service = RealtimeService()
|
||||
Reference in New Issue
Block a user