Add an auto-reconnect websocket API

piker/brokers/kraken.py

@@ -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, "
+                    "Heartbeat is too slow, resetting ws connection")
-                    "resetting ws connection")
+
-                raise trio_websocket._impl.ConnectionClosed(
+                await ws._connect()
-                    "Reset Connection")
+                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
+                # XXX: why tf is this not printing without --tl flag?
-                # for determining when to do connection
+                log.debug(f"Heartbeat after {delay}")
-                # resets eh?
+                # 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:
+        async with open_autorecon_ws('wss://ws.kraken.com/') as ws:
-            try:
-                async with trio_websocket.open_websocket_url(
-                    '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():
+            await ws.send_msg(l1_sub)
-                        return json.loads(await ws.get_message())
 
             # 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")