143 lines
5.8 KiB
Python
143 lines
5.8 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from typing import Dict, Set, Tuple, Any
|
|
|
|
import websockets
|
|
|
|
from app.streaming.live_store import live_store
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class StreamKey:
|
|
symbol: str
|
|
timeframe: str # only '1m' supported in hub
|
|
|
|
|
|
class BinanceStreamHub:
|
|
"""
|
|
Maintains a single upstream websocket per (symbol,timeframe) and fans out
|
|
kline events to multiple subscribers via asyncio.Queues.
|
|
"""
|
|
|
|
def __init__(self, base_ws: str | None = None) -> None:
|
|
self.base_ws = (base_ws or os.getenv("BINANCE_WS_URL", "wss://stream.binance.com:9443/ws")).rstrip("/")
|
|
self._subs: Dict[StreamKey, Set[asyncio.Queue]] = {}
|
|
self._tasks: Dict[StreamKey, asyncio.Task] = {}
|
|
self._lock = asyncio.Lock()
|
|
|
|
def get_status(self) -> list[dict]:
|
|
"""Return status snapshot of active streams."""
|
|
out: list[dict] = []
|
|
for key, subs in self._subs.items():
|
|
hist = live_store.get_history(key.symbol, key.timeframe)
|
|
last_ts = hist[-1]["time"] if hist else None
|
|
last_iso = None
|
|
if isinstance(last_ts, (int, float)):
|
|
try:
|
|
last_iso = datetime.utcfromtimestamp(int(last_ts)).isoformat() + "Z"
|
|
except Exception:
|
|
last_iso = None
|
|
out.append({
|
|
"symbol": key.symbol,
|
|
"timeframe": key.timeframe,
|
|
"subscribers": len(subs),
|
|
"last_event_time": last_iso,
|
|
})
|
|
return out
|
|
|
|
async def subscribe(self, symbol: str, timeframe: str = "1m") -> Tuple[asyncio.Queue, Any]:
|
|
"""Subscribe to a stream. Returns (queue, unsubscribe_cb)."""
|
|
if timeframe != "1m":
|
|
raise ValueError("BinanceStreamHub currently supports timeframe '1m' only")
|
|
key = StreamKey(symbol=symbol.upper().replace("/", ""), timeframe=timeframe)
|
|
q: asyncio.Queue = asyncio.Queue(maxsize=1000)
|
|
async with self._lock:
|
|
subs = self._subs.get(key)
|
|
if not subs:
|
|
subs = set()
|
|
self._subs[key] = subs
|
|
subs.add(q)
|
|
if key not in self._tasks:
|
|
self._tasks[key] = asyncio.create_task(self._run_stream(key))
|
|
|
|
async def _unsubscribe() -> None:
|
|
async with self._lock:
|
|
s = self._subs.get(key)
|
|
if s and q in s:
|
|
s.remove(q)
|
|
# Close queue to unblock listeners
|
|
try:
|
|
q.put_nowait(None)
|
|
except Exception:
|
|
pass
|
|
if s is not None and len(s) == 0:
|
|
# cancel task and cleanup
|
|
t = self._tasks.pop(key, None)
|
|
if t:
|
|
t.cancel()
|
|
self._subs.pop(key, None)
|
|
return q, _unsubscribe
|
|
|
|
async def _run_stream(self, key: StreamKey) -> None:
|
|
symbol = key.symbol
|
|
stream = f"{symbol.lower()}@kline_{key.timeframe}"
|
|
url = self.base_ws.replace("/ws", "/stream") + f"?streams={stream}"
|
|
# Reconnect loop
|
|
while True:
|
|
try:
|
|
async with websockets.connect(url, ping_interval=20, ping_timeout=20) as ws:
|
|
async for message in ws:
|
|
try:
|
|
data = json.loads(message)
|
|
k = (data.get("data") or {}).get("k") or {}
|
|
if not k:
|
|
continue
|
|
# Normalize event
|
|
evt = {
|
|
"symbol": symbol,
|
|
"timeframe": key.timeframe,
|
|
"open_time": datetime.fromtimestamp(k["t"] / 1000.0).isoformat(),
|
|
"close_time": datetime.fromtimestamp(k["T"] / 1000.0).isoformat(),
|
|
"open": float(k["o"]),
|
|
"high": float(k["h"]),
|
|
"low": float(k["l"]),
|
|
"close": float(k["c"]),
|
|
"volume": float(k.get("v", 0.0)),
|
|
"is_closed": bool(k.get("x", False)),
|
|
"source": "binance",
|
|
}
|
|
# Update live store (1m bar)
|
|
try:
|
|
tsec = int(k["T"] // 1000)
|
|
live_store.ingest_bar(symbol=symbol, timeframe=key.timeframe, bar={
|
|
"time": tsec, "open": evt["open"], "high": evt["high"], "low": evt["low"], "close": evt["close"], "volume": evt["volume"],
|
|
})
|
|
except Exception:
|
|
pass
|
|
# Fan-out to subscribers
|
|
subs = self._subs.get(key) or set()
|
|
for q in list(subs):
|
|
try:
|
|
if q.full():
|
|
q.get_nowait()
|
|
q.put_nowait(evt)
|
|
except Exception:
|
|
# Drop failed subscriber
|
|
try:
|
|
subs.remove(q)
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
continue
|
|
except Exception:
|
|
await asyncio.sleep(1.5)
|
|
|
|
|
|
# Singleton hub instance
|
|
hub = BinanceStreamHub()
|