Add an auto-reconnect websocket API
Tyler Goodlet
parent
bbd54e8f95
commit
add63734f1
|
|
@ -17,14 +17,22 @@
|
|||
"""
|
||||
Kraken backend.
|
||||
"""
|
||||
from contextlib import asynccontextmanager
|
||||
from contextlib import asynccontextmanager, AsyncExitStack
|
||||
from dataclasses import asdict, field
|
||||
from types import ModuleType
|
||||
from typing import List, Dict, Any, Tuple, Optional
|
||||
import json
|
||||
import time
|
||||
|
||||
import trio_websocket
|
||||
from trio_websocket._impl import ConnectionClosed, DisconnectionTimeout
|
||||
from trio_websocket._impl import (
|
||||
ConnectionClosed,
|
||||
DisconnectionTimeout,
|
||||
ConnectionRejected,
|
||||
HandshakeError,
|
||||
ConnectionTimeout,
|
||||
)
|
||||
|
||||
import arrow
|
||||
import asks
|
||||
import numpy as np
|
||||
|
|
@ -229,22 +237,27 @@ async def get_client() -> Client:
|
|||
yield Client()
|
||||
|
||||
|
||||
async def recv_msg(recv):
|
||||
async def stream_messages(ws):
|
||||
|
||||
too_slow_count = last_hb = 0
|
||||
|
||||
while True:
|
||||
with trio.move_on_after(1.5) as cs:
|
||||
msg = await recv()
|
||||
|
||||
# trigger reconnection logic if too slow
|
||||
with trio.move_on_after(5) as cs:
|
||||
msg = await ws.recv_msg()
|
||||
|
||||
# trigger reconnection if heartbeat is laggy
|
||||
if cs.cancelled_caught:
|
||||
|
||||
too_slow_count += 1
|
||||
if too_slow_count > 2:
|
||||
|
||||
if too_slow_count > 10:
|
||||
log.warning(
|
||||
"Heartbeat is to slow, "
|
||||
"resetting ws connection")
|
||||
raise trio_websocket._impl.ConnectionClosed(
|
||||
"Reset Connection")
|
||||
"Heartbeat is too slow, resetting ws connection")
|
||||
|
||||
await ws._connect()
|
||||
too_slow_count = 0
|
||||
continue
|
||||
|
||||
if isinstance(msg, dict):
|
||||
if msg.get('event') == 'heartbeat':
|
||||
|
|
@ -252,11 +265,11 @@ async def recv_msg(recv):
|
|||
now = time.time()
|
||||
delay = now - last_hb
|
||||
last_hb = now
|
||||
log.trace(f"Heartbeat after {delay}")
|
||||
|
||||
# TODO: hmm i guess we should use this
|
||||
# for determining when to do connection
|
||||
# resets eh?
|
||||
# XXX: why tf is this not printing without --tl flag?
|
||||
log.debug(f"Heartbeat after {delay}")
|
||||
# print(f"Heartbeat after {delay}")
|
||||
|
||||
continue
|
||||
|
||||
err = msg.get('errorMessage')
|
||||
|
|
@ -326,6 +339,95 @@ def make_sub(pairs: List[str], data: Dict[str, Any]) -> Dict[str, str]:
|
|||
}
|
||||
|
||||
|
||||
class AutoReconWs:
|
||||
"""Make ``trio_websocketw` sockets stay up no matter the bs.
|
||||
|
||||
"""
|
||||
recon_errors = (
|
||||
ConnectionClosed,
|
||||
DisconnectionTimeout,
|
||||
ConnectionRejected,
|
||||
HandshakeError,
|
||||
ConnectionTimeout,
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
url: str,
|
||||
stack: AsyncExitStack,
|
||||
serializer: ModuleType = json,
|
||||
):
|
||||
self.url = url
|
||||
self._stack = stack
|
||||
self._ws: 'WebSocketConnection' = None # noqa
|
||||
|
||||
async def _connect(
|
||||
self,
|
||||
tries: int = 10000,
|
||||
) -> None:
|
||||
try:
|
||||
await self._stack.aclose()
|
||||
except (DisconnectionTimeout, RuntimeError):
|
||||
await trio.sleep(1)
|
||||
|
||||
last_err = None
|
||||
for i in range(tries):
|
||||
try:
|
||||
self._ws = await self._stack.enter_async_context(
|
||||
trio_websocket.open_websocket_url(self.url)
|
||||
)
|
||||
log.info(f'Connection success: {self.url}')
|
||||
return
|
||||
except self.recon_errors as err:
|
||||
last_err = err
|
||||
log.error(
|
||||
f'{self} connection bail with '
|
||||
f'{type(err)}...retry attempt {i}'
|
||||
)
|
||||
await trio.sleep(1)
|
||||
continue
|
||||
else:
|
||||
log.exception('ws connection fail...')
|
||||
raise last_err
|
||||
|
||||
async def send_msg(
|
||||
self,
|
||||
data: Any,
|
||||
) -> None:
|
||||
while True:
|
||||
try:
|
||||
return await self._ws.send_message(json.dumps(data))
|
||||
except self.recon_errors:
|
||||
await self._connect()
|
||||
|
||||
async def recv_msg(
|
||||
self,
|
||||
) -> Any:
|
||||
while True:
|
||||
try:
|
||||
return json.loads(await self._ws.get_message())
|
||||
except self.recon_errors:
|
||||
await self._connect()
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def open_autorecon_ws(url):
|
||||
"""Apparently we can QoS for all sorts of reasons..so catch em.
|
||||
|
||||
"""
|
||||
async with AsyncExitStack() as stack:
|
||||
ws = AutoReconWs(url, stack)
|
||||
# async with trio_websocket.open_websocket_url(url) as ws:
|
||||
# await tractor.breakpoint()
|
||||
|
||||
await ws._connect()
|
||||
try:
|
||||
yield ws
|
||||
|
||||
finally:
|
||||
await stack.aclose()
|
||||
|
||||
|
||||
# @tractor.msg.pub
|
||||
async def stream_quotes(
|
||||
# get_topics: Callable,
|
||||
|
|
@ -393,11 +495,7 @@ async def stream_quotes(
|
|||
}
|
||||
yield init_msgs
|
||||
|
||||
while True:
|
||||
try:
|
||||
async with trio_websocket.open_websocket_url(
|
||||
'wss://ws.kraken.com/',
|
||||
) as ws:
|
||||
async with open_autorecon_ws('wss://ws.kraken.com/') as ws:
|
||||
|
||||
# XXX: setup subs
|
||||
# https://docs.kraken.com/websockets/#message-subscribe
|
||||
|
|
@ -412,21 +510,19 @@ async def stream_quotes(
|
|||
# be completely fine to request from a separate task
|
||||
# since internally the ws methods appear to be FIFO
|
||||
# locked.
|
||||
await ws.send_message(json.dumps(ohlc_sub))
|
||||
await ws.send_msg(ohlc_sub)
|
||||
|
||||
# trade data (aka L1)
|
||||
l1_sub = make_sub(
|
||||
list(ws_pairs.values()),
|
||||
{'name': 'spread'} # 'depth': 10}
|
||||
|
||||
)
|
||||
await ws.send_message(json.dumps(l1_sub))
|
||||
|
||||
async def recv():
|
||||
return json.loads(await ws.get_message())
|
||||
await ws.send_msg(l1_sub)
|
||||
|
||||
# pull a first quote and deliver
|
||||
msg_gen = recv_msg(recv)
|
||||
msg_gen = stream_messages(ws)
|
||||
|
||||
typ, ohlc_last = await msg_gen.__anext__()
|
||||
|
||||
topic, quote = normalize(ohlc_last)
|
||||
|
|
@ -510,6 +606,3 @@ async def stream_quotes(
|
|||
# XXX: format required by ``tractor.msg.pub``
|
||||
# requires a ``Dict[topic: str, quote: dict]``
|
||||
yield {topic: quote}
|
||||
|
||||
except (ConnectionClosed, DisconnectionTimeout):
|
||||
log.exception("Good job kraken...reconnecting")
|
||||
|
|
|
|||
Loading…
Reference in New Issue