Coverage for adrs/io/stream.py: 39%
206 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-29 05:07 +0800
1import json
2import zlib
3import time
4import asyncio
5import logging
6import websockets
8from datetime import datetime, timezone
9from typing import AsyncIterator, AsyncGenerator, Awaitable, Callable, TypeVar, cast
11from grpc import StatusCode
12from grpc.aio import AioRpcError
13from nats_client import NATSClient, Msg
14from websockets.exceptions import ConnectionClosed # type: ignore[import-untyped]
16from adrs.types import (
17 Topic,
18 Message,
19 CollectedData,
20 Data,
21 SubscriptionResponse,
22 BoundedSet,
23)
25logger = logging.getLogger(__name__)
27_T = TypeVar("_T")
29# gRPC codes worth retrying. JetStream rate-limiting (429) surfaces as
30# UNAVAILABLE with "Received http2 header with status: 429" in the details;
31# RESOURCE_EXHAUSTED also maps to 429 in some gRPC implementations.
32_RETRYABLE_CODES = frozenset(
33 {
34 StatusCode.UNAVAILABLE,
35 StatusCode.RESOURCE_EXHAUSTED,
36 StatusCode.DEADLINE_EXCEEDED,
37 }
38)
41async def _retry_grpc(
42 op: Callable[[], Awaitable[_T]],
43 *,
44 desc: str,
45 max_retries: int,
46 retry_backoff: float,
47) -> _T:
48 """Run `op`, retrying transient gRPC failures with exponential backoff.
50 Rate-limited JetStream ops are "rejected-not-processed", so retrying does
51 not risk duplicate side effects. Non-retryable codes propagate immediately.
52 """
53 for attempt in range(max_retries + 1):
54 try:
55 return await op()
56 except AioRpcError as e:
57 if e.code() not in _RETRYABLE_CODES or attempt == max_retries:
58 raise
59 delay = retry_backoff * (2**attempt)
60 logger.warning(
61 "%s transient gRPC failure (%s); retry %d/%d in %.2fs",
62 desc,
63 e.code(),
64 attempt + 1,
65 max_retries,
66 delay,
67 )
68 await asyncio.sleep(delay)
69 raise AssertionError("unreachable") # loop either returns or raises
72_BINANCE_WS = {
73 "binance-spot": "wss://stream.binance.com:9443/ws",
74 "binance-futures": "wss://fstream.binance.com/ws",
75}
77_BYBIT_WS = {
78 "bybit-linear": "wss://stream.bybit.com/v5/public/linear",
79 "bybit-spot": "wss://stream.bybit.com/v5/public/spot",
80 "bybit-inverse": "wss://stream.bybit.com/v5/public/inverse",
81}
83_BYBIT_INTERVAL = {
84 "1m": "1",
85 "3m": "3",
86 "5m": "5",
87 "15m": "15",
88 "30m": "30",
89 "1h": "60",
90 "2h": "120",
91 "4h": "240",
92 "6h": "360",
93 "12h": "720",
94 "1d": "D",
95 "1w": "W",
96 "1M": "M",
97}
98_BYBIT_INTERVAL_INV = {v: k for k, v in _BYBIT_INTERVAL.items()}
101def _ms(ms: int) -> datetime:
102 return datetime.fromtimestamp(ms / 1000, tz=timezone.utc)
105def _now_ms() -> int:
106 return int(time.time() * 1000)
109async def _binance_worker(url: str, topics: list[Topic], queue: asyncio.Queue) -> None:
110 streams = [
111 f"{t.params['symbol'].lower()}@kline_{t.params['interval']}" for t in topics
112 ]
113 backoff = 1.0
114 while True:
115 try:
116 async with websockets.connect(url, ping_interval=20) as ws:
117 await ws.send(
118 json.dumps({"method": "SUBSCRIBE", "params": streams, "id": 1})
119 )
120 await queue.put(
121 SubscriptionResponse(
122 conn_id=url, success=True, message=f"subscribed {streams}"
123 )
124 )
125 backoff = 1.0
126 async for raw in ws:
127 msg = json.loads(raw)
128 if msg.get("e") != "kline":
129 continue
130 k = msg["k"]
131 if not k.get("x"):
132 continue
133 sym, iv = msg["s"], k["i"]
134 topic_str = next(
135 (
136 str(t)
137 for t in topics
138 if t.params["symbol"] == sym and t.params["interval"] == iv
139 ),
140 None,
141 )
142 if topic_str is None:
143 continue
144 await queue.put(
145 CollectedData(
146 topic=topic_str,
147 data=cast(
148 list[Data],
149 [
150 {
151 "start_time": _ms(k["t"]),
152 "open": float(k["o"]),
153 "high": float(k["h"]),
154 "low": float(k["l"]),
155 "close": float(k["c"]),
156 "volume": float(k["v"]),
157 }
158 ],
159 ),
160 local_timestamp_ms=_now_ms(),
161 type="delta",
162 )
163 )
164 except (ConnectionClosed, OSError, Exception) as e:
165 logger.warning("[binance] %s — reconnect in %.0fs", e, backoff)
166 await asyncio.sleep(backoff)
167 backoff = min(backoff * 2, 60.0)
170async def _bybit_worker(url: str, topics: list[Topic], queue: asyncio.Queue) -> None:
171 args = [
172 f"kline.{_BYBIT_INTERVAL.get(t.params['interval'], t.params['interval'])}.{t.params['symbol']}"
173 for t in topics
174 ]
175 backoff = 1.0
176 while True:
177 try:
178 async with websockets.connect(url, ping_interval=20) as ws:
179 await ws.send(json.dumps({"op": "subscribe", "args": args}))
180 await queue.put(
181 SubscriptionResponse(
182 conn_id=url, success=True, message=f"subscribed {args}"
183 )
184 )
185 backoff = 1.0
186 async for raw in ws:
187 msg = json.loads(raw)
188 topic_field = msg.get("topic", "")
189 if not topic_field.startswith("kline."):
190 continue
191 _, bybit_iv, sym = topic_field.split(".", 2)
192 iv = _BYBIT_INTERVAL_INV.get(bybit_iv, bybit_iv)
193 topic_str = next(
194 (
195 str(t)
196 for t in topics
197 if t.params["symbol"] == sym and t.params["interval"] == iv
198 ),
199 None,
200 )
201 if topic_str is None:
202 continue
203 for item in msg.get("data", []):
204 if not item.get("confirm"):
205 continue
206 await queue.put(
207 CollectedData(
208 topic=topic_str,
209 data=cast(
210 list[Data],
211 [
212 {
213 "start_time": _ms(item["start"]),
214 "open": float(item["open"]),
215 "high": float(item["high"]),
216 "low": float(item["low"]),
217 "close": float(item["close"]),
218 "volume": float(item["volume"]),
219 }
220 ],
221 ),
222 local_timestamp_ms=_now_ms(),
223 type="delta",
224 )
225 )
226 except (ConnectionClosed, OSError, Exception) as e:
227 logger.warning("[bybit] %s — reconnect in %.0fs", e, backoff)
228 await asyncio.sleep(backoff)
229 backoff = min(backoff * 2, 60.0)
232async def _coinbase_worker(url: str, topics: list[Topic], queue: asyncio.Queue) -> None:
233 product_ids = list({t.params["symbol"] for t in topics})
234 backoff = 1.0
235 while True:
236 try:
237 async with websockets.connect(url, ping_interval=30) as ws:
238 await ws.send(
239 json.dumps(
240 {
241 "type": "subscribe",
242 "product_ids": product_ids,
243 "channel": "candles",
244 }
245 )
246 )
247 await queue.put(
248 SubscriptionResponse(
249 conn_id=url, success=True, message=f"subscribed {product_ids}"
250 )
251 )
252 backoff = 1.0
253 async for raw in ws:
254 msg = json.loads(raw)
255 if msg.get("channel") != "candles":
256 continue
257 for event in msg.get("events", []):
258 msg_type = (
259 "snapshot" if event.get("type") == "snapshot" else "delta"
260 )
261 for candle in event.get("candles", []):
262 pid = candle["product_id"]
263 topic_str = next(
264 (str(t) for t in topics if t.params["symbol"] == pid),
265 None,
266 )
267 if topic_str is None:
268 continue
269 await queue.put(
270 CollectedData(
271 topic=topic_str,
272 data=cast(
273 list[Data],
274 [
275 {
276 "start_time": datetime.fromtimestamp(
277 int(candle["start"]),
278 tz=timezone.utc,
279 ),
280 "open": float(candle["open"]),
281 "high": float(candle["high"]),
282 "low": float(candle["low"]),
283 "close": float(candle["close"]),
284 "volume": float(candle["volume"]),
285 }
286 ],
287 ),
288 local_timestamp_ms=_now_ms(),
289 type=msg_type,
290 )
291 )
292 except (ConnectionClosed, OSError, Exception) as e:
293 logger.warning("[coinbase] %s — reconnect in %.0fs", e, backoff)
294 await asyncio.sleep(backoff)
295 backoff = min(backoff * 2, 60.0)
298class PublicDatasourceStream:
299 """DatasourceStream backed by public exchange WebSocket APIs.
301 Supports: binance-spot, binance-futures, bybit-linear, bybit-spot,
302 bybit-inverse, coinbase.
303 """
305 async def connect(self, topics: list[Topic]) -> AsyncIterator[Message] | None:
306 by_provider: dict[str, list[Topic]] = {}
307 for t in topics:
308 by_provider.setdefault(t.provider, []).append(t)
310 queue: asyncio.Queue[Message] = asyncio.Queue()
311 tasks: list[asyncio.Task] = []
313 for provider, pts in by_provider.items():
314 if provider in _BINANCE_WS:
315 tasks.append(
316 asyncio.create_task(
317 _binance_worker(_BINANCE_WS[provider], pts, queue),
318 name=f"ws_{provider}",
319 )
320 )
321 elif provider in _BYBIT_WS:
322 tasks.append(
323 asyncio.create_task(
324 _bybit_worker(_BYBIT_WS[provider], pts, queue),
325 name=f"ws_{provider}",
326 )
327 )
328 elif provider == "coinbase":
329 tasks.append(
330 asyncio.create_task(
331 _coinbase_worker(
332 "wss://advanced-trade-ws.coinbase.com", pts, queue
333 ),
334 name=f"ws_{provider}",
335 )
336 )
337 else:
338 logger.warning(
339 "[PublicDatasourceStream] unsupported provider: %s", provider
340 )
342 if not tasks:
343 return None
345 async def _iter() -> AsyncGenerator[Message, None]:
346 try:
347 while True:
348 yield await queue.get()
349 finally:
350 for task in tasks:
351 task.cancel()
353 return _iter()
356def _transform_data_from_timestamp_ms(data: list[Data]) -> list[Data]:
357 return list(
358 map(
359 lambda d: {
360 **d,
361 "start_time": datetime.fromtimestamp(
362 cast(int, d["start_time"]) / 1000, tz=timezone.utc
363 ),
364 },
365 data,
366 )
367 )
370class PublicNatsDatasourceStream:
371 def __init__(
372 self,
373 flow_nats: NATSClient,
374 timeout_secs: int = 5,
375 max_timeout_secs: int = 60,
376 multiplier: int = 2,
377 max_retries: int = 5,
378 retry_backoff: float = 0.5,
379 log_title: str = "connect",
380 ):
381 self.flow_nats = flow_nats
382 self.timeout_secs = timeout_secs
383 self.max_timeout_secs = max_timeout_secs
384 self.multiplier = multiplier
385 self.max_retries = max_retries
386 self.retry_backoff = retry_backoff
387 self.log_title = log_title
389 async def _stream(self, topics: list[Topic]):
390 queue: asyncio.Queue[Message] = asyncio.Queue()
392 def make_handler(topic: Topic):
393 dedup = BoundedSet(10)
395 async def handler(msg):
396 try:
397 message = json.loads(zlib.decompress(msg.data))
398 message["data"] = _transform_data_from_timestamp_ms(message["data"])
400 if len(message["data"]) == 0:
401 await queue.put(message)
402 return
404 key = (message["topic"], message["data"][0]["start_time"])
405 if dedup.add(key):
406 await queue.put(message)
407 else:
408 logging.warning(
409 f"[stream] received duplicated message on {topic}: {message}"
410 )
411 except Exception as e:
412 logging.error(f"[stream] failed to process message on {topic}: {e}")
414 return handler
416 for topic in topics:
417 handler = make_handler(topic)
418 await _retry_grpc(
419 lambda t=str(topic), cb=handler: self.flow_nats.subscribe(t, cb=cb),
420 desc=f"[{self.log_title}] subscribe {topic}",
421 max_retries=self.max_retries,
422 retry_backoff=self.retry_backoff,
423 )
424 logging.info(f"[stream] subscribed to {topic}")
426 async def generator() -> AsyncIterator[Message]:
427 while True:
428 message = await queue.get()
429 yield message
431 return generator()
433 async def connect(self, topics: list[Topic]) -> AsyncIterator[Message] | None:
434 attempts = 0
435 while attempts < self.max_retries and self.timeout_secs < self.max_timeout_secs:
436 try:
437 async with asyncio.timeout(self.timeout_secs):
438 return await self._stream(topics)
439 except asyncio.TimeoutError:
440 logger.warning(
441 f"[{self.log_title}] connect operation timed out after {attempts + 1} attempts, retrying..."
442 )
443 attempts += 1
444 self.timeout_secs *= self.multiplier
445 return None
448class PublicMetricStream:
449 """Thin NATS transport: publishes/subscribes the subject verbatim.
451 It used to prepend an `insert_prefix` to JetStream publishes and a separate
452 `publish_prefix` to core publishes. That blanket `publish_prefix` mangled
453 routing subjects (an alpha signal meant for `alpha_signal.<id>` went out on
454 `portfolio_signal.alpha_signal.<id>`, which the PortfolioExecutor's
455 `alpha_signal.*` subscription never matched). Subject construction now lives
456 with the caller: MetricBuilder owns the dashboard-insert namespace, the
457 executors own their routing roots."""
459 def __init__(
460 self,
461 nats: NATSClient,
462 max_retries: int = 5,
463 retry_backoff: float = 0.5,
464 ):
465 self.nats = nats
466 self.max_retries = max_retries
467 self.retry_backoff = retry_backoff
469 async def init(self):
470 await self.nats.jetstream()
472 async def subscribe(
473 self, subject: str, callback: Callable[[Msg], Awaitable[None]] | None
474 ):
475 async def default_cb(_: Msg) -> None:
476 pass
478 if callback is None:
479 callback = default_cb
481 await self.nats.subscribe(subject, cb=callback)
483 async def publish(self, subject: str, payload: bytes, **kwargs) -> None:
484 use_jetstream = kwargs.get("use_jetstream", False)
485 headers = kwargs.get("headers")
487 async def _op() -> None:
488 if use_jetstream:
489 # headers may carry Nats-Msg-Id: JetStream server-side dedup,
490 # so a duplicate publish (e.g. two OMS replicas overlapping
491 # during a rolling deploy) collapses into one message instead
492 # of landing twice downstream.
493 await self.nats.js_publish(
494 subject=subject, payload=payload, headers=headers
495 )
496 else:
497 await self.nats.publish(subject, payload)
499 await _retry_grpc(
500 _op,
501 desc=f"publish {subject}",
502 max_retries=self.max_retries,
503 retry_backoff=self.retry_backoff,
504 )