Merge fix/ws-reconnect into deephaven-fixes

Author: chipkent

docs/source/WebSocket.rst

@@ -49,6 +49,11 @@ Close
 ============================
 .. automethod:: massive.WebSocketClient.close
 
+============================
+Request close
+============================
+.. automethod:: massive.WebSocketClient.request_close
+
 .. _Stocks getting started: https://massive.com/docs/stocks/ws_getting-started
 .. _Options getting started: https://massive.com/docs/options/ws_getting-started
 .. _Forex getting started: https://massive.com/docs/forex/ws_getting-started

massive/websocket/__init__.py

@@ -15,13 +15,20 @@
 import logging
 import json
 import asyncio
+import random
 import ssl
 import sys
 import threading
 import time
 from .models import *
 from websockets.asyncio.client import connect, ClientConnection
-from websockets.exceptions import ConnectionClosedOK, ConnectionClosedError
+from websockets.exceptions import (
+    ConnectionClosed,
+    ConnectionClosedOK,
+    ConnectionClosedError,
+    InvalidMessage,
+    InvalidStatus,
+)
 from ..certs import ca_bundle_path
 from ..logging import get_logger
 import logging
@@ -44,10 +51,12 @@ def __init__(
         raw: bool = False,
         verbose: bool = False,
         subscriptions: Optional[List[str]] = None,
-        max_reconnects: Optional[int] = 5,
+        max_reconnects: Optional[int] = None,
         secure: bool = True,
         custom_json: Optional[Any] = None,
         ca_certs: Optional[str] = None,
+        reconnect_backoff: float = 15.0,
+        max_reconnect_backoff: float = 120.0,
         **kwargs,
     ):
         """
@@ -58,11 +67,19 @@ def __init__(
         :param raw: Whether to pass raw Union[str, bytes] to user callback.
         :param verbose: Whether to log client and server status messages.
         :param subscriptions: List of subscription parameters.
-        :param max_reconnects: How many times to reconnect on network outage before ending .connect event loop.
+        :param max_reconnects: How many consecutive reconnect attempts may fail (the connection is refused, or ends before or after authenticating) before .connect raises the last error. The count starts over whenever a connection authenticates. None, the default, retries forever. Only transient failures are retried and the wait between attempts grows, while other errors (e.g. HTTP 4xx) are raised at once, so with the default .run is safe in a long-lived process.
         :param custom_json: Optional module exposing loads/dumps functions (similar to Python's json module) to be used for JSON conversions.
         :param ca_certs: Path to a CA bundle to verify TLS with. Defaults to the SSL_CERT_FILE or REQUESTS_CA_BUNDLE environment variable, then certifi.
+        :param reconnect_backoff: Seconds to wait before reconnecting after the connection ends or an attempt to connect fails. The wait doubles with each consecutive failure and gets up to 20% random jitter, but never leaves [reconnect_backoff, max_reconnect_backoff]. Defaults to 15, the minimum Massive asks clients to wait; 0 disables the wait and is meant only for tests.
+        :param max_reconnect_backoff: The longest wait between reconnect attempts, jitter included.
+        :raises ValueError: If reconnect_backoff is negative or greater than max_reconnect_backoff.
         :return: A client.
         """
+        if reconnect_backoff < 0 or max_reconnect_backoff < reconnect_backoff:
+            raise ValueError(
+                "need 0 <= reconnect_backoff <= max_reconnect_backoff, got "
+                f"{reconnect_backoff} and {max_reconnect_backoff}"
+            )
         if api_key is None:
             raise AuthError(
                 f"Must specify env var {env_key} or pass api_key in constructor"
@@ -97,7 +114,7 @@ def __init__(
         # topics to subscribe to again even if already subscribed
         self._resend: Set[str] = set()
         # wake-up of connect() when subscriptions change: the loop it runs
-        # on, whether a wake-up is queued there, whether one arrived since the
+        # on (set only while connect() runs), whether a wake-up is queued there, whether one arrived since the
         # last reconciliation, and how to interrupt the recv() in progress
         self._loop: Optional[asyncio.AbstractEventLoop] = None
         self._wake_pending = False
@@ -113,6 +130,14 @@ def __init__(
             self.json = custom_json
         else:
             self.json = json
+        # close() and request_close() stop connect()
+        self._close_requested = threading.Event()
+        self._close_event: Optional[asyncio.Event] = None
+        self._close_task: Optional[asyncio.Future] = None
+        # whether the server said max_connections during the last handshake
+        self.max_connections_reached = False
+        self.reconnect_backoff = reconnect_backoff
+        self.max_reconnect_backoff = max_reconnect_backoff
 
     # https://websockets.readthedocs.io/en/stable/reference/client.html#opening-a-connection
     async def connect(
@@ -127,31 +152,88 @@ async def connect(
         """
         Connect to websocket server and run `processor(msg)` on every new `msg`.
 
+        When the connection ends, whether the server closed it cleanly or not,
+        or an attempt to connect fails for a reason that may pass (the server
+        is unreachable, times out, or answers with HTTP 5xx or 429), the
+        client waits (see `reconnect_backoff`), reconnects, and subscribes
+        again to every scheduled subscription. Other errors, e.g. HTTP 4xx or
+        a TLS certificate that cannot be verified, are raised at once. It
+        returns once `close()` or `request_close()` is called.
+
+        If the server sends a `max_connections` status during the handshake
+        (the account has too many connections open), `max_connections_reached`
+        is set until a later handshake succeeds, and in raw mode the message is
+        also passed to `processor`. The server then closes the connection, and
+        the client reconnects after the backoff.
+
         :param processor: The callback to process messages.
         :param close_timeout: How long to wait for handshake when calling .close.
         :raises AuthError: If invalid API key is supplied.
+        :raises Exception: The last error, if more than `max_reconnects`
+            consecutive attempts fail; ConnectionClosedError if the server
+            closed the last connection, even cleanly.
         """
-        reconnects = 0
-        logger.debug("connect: %s", self.url)
+        # reset the wake-up state before publishing the loop to other threads
         self._woken = False
         with self._subs_lock:
             self._wake_pending = False
         self._loop = asyncio.get_running_loop()
+        self._close_event = asyncio.Event()
+        try:
+            await self._connect(processor, close_timeout, **kwargs)
+            # it returned because close was requested, which is now done
+            self._close_requested.clear()
+        finally:
+            self._loop = None
+            self._close_event = None
+            closing, self._close_task = self._close_task, None
+            if closing is not None:
+                await asyncio.gather(closing, return_exceptions=True)
+
+    async def _connect(
+        self,
+        processor: Union[
+            Callable[[List[WebSocketMessage]], Awaitable],
+            Callable[[Union[str, bytes]], Awaitable],
+        ],
+        close_timeout: int,
+        **kwargs,
+    ):
+        # consecutive failed attempts, which set the backoff and are limited
+        # by max_reconnects
+        failures = 0
+        logger.debug("connect: %s", self.url)
         ssl_context = self._ssl_context()
 
-        last_exc = None
-        async for s in connect(
-            self.url,
-            close_timeout=close_timeout,
-            ssl=ssl_context,
-            process_exception=lambda exc: None,
-            **kwargs,
-        ):
+        # invalid arguments raise here; each attempt awaits it again
+        connector = connect(
+            self.url, close_timeout=close_timeout, ssl=ssl_context, **kwargs
+        )
+        last_exc: Optional[Exception] = None
+        delay = 0.0
+        while not self._close_requested.is_set():
+            if failures > 0 and await self._wait_for_close(delay):
+                return
+            try:
+                s = await self._open(connector)
+            except Exception as e:
+                if not self._is_transient(e):
+                    raise
+                last_exc = e
+                failures += 1
+                if self._too_many_failures(failures):
+                    break
+                delay = self._reconnect_delay(failures - 1)
+                logger.warning("connect failed; reconnecting in %.1f s: %s", delay, e)
+                continue
+            if s is None:
+                return
             self.websocket = s
             try:
                 msg = await s.recv()
                 logger.debug("connected: %s", msg)
-                self._log_server_statuses(self.json.loads(msg))
+                msg_parsed = self.json.loads(msg)
+                self._log_server_statuses(msg_parsed)
                 logger.debug("authing...")
                 await s.send(
                     self.json.dumps({"action": "auth", "params": self.api_key})
@@ -161,6 +243,14 @@ async def connect(
                 logger.debug("authed: %s", auth_msg)
                 if auth_msg_parsed[0]["status"] == "auth_failed":
                     raise AuthError(auth_msg_parsed[0]["message"])
+                self.max_connections_reached = False
+                for frame, parsed in ((msg, msg_parsed), (auth_msg, auth_msg_parsed)):
+                    if any(m.get("status") == "max_connections" for m in parsed):
+                        self.max_connections_reached = True
+                        if self.raw:
+                            await processor(frame)  # type: ignore
+                if not self.max_connections_reached:
+                    failures = 0
                 # e.g. max_connections, which the server sends before closing
                 self._log_server_statuses(auth_msg_parsed)
                 self._flush_server_errors(force=True)
@@ -186,31 +276,103 @@ async def connect(
 
                     if len(cmsg) > 0:
                         await processor(cmsg)  # type: ignore
-            except ConnectionClosedOK as e:
-                logger.debug("connection closed (OK): %s", e)
+            except ConnectionClosed as e:
+                ok = isinstance(e, ConnectionClosedOK)
+                logger.debug("connection closed (%s): %s", "OK" if ok else "ERR", e)
                 # the server sends a status (e.g. slow_consumer) before closing
                 self._flush_server_errors(force=True)
-                return
-            except ConnectionClosedError as e:
+                if self._close_requested.is_set():
+                    return
+                # the server can close cleanly too, e.g. after max_connections
                 last_exc = e
-                logger.debug("connection closed (ERR): %s", e)
-                self._flush_server_errors(force=True)
-                reconnects += 1
+                failures += 1
+                if self._too_many_failures(failures):
+                    break
+                delay = self._reconnect_delay(failures - 1)
+                logger.warning(
+                    "connection closed; reconnecting in %.1f s: %s", delay, e
+                )
+            finally:
+                self.websocket = None
+                # scheduled_subs is the desired set, including changes not sent
+                # yet: keep it, and subscribe all of it on the next connection
                 with self._subs_lock:
-                    self.scheduled_subs = set(self.subs)
                     self.subs = set()
                     self.schedule_resub = True
-                if self.max_reconnects is not None and reconnects > self.max_reconnects:
-                    break
-                continue
-
-        if (
-            last_exc
-            and self.max_reconnects is not None
-            and reconnects > self.max_reconnects
-        ):
+                await s.close()
+
+        if last_exc is not None and self._too_many_failures(failures):
+            if isinstance(last_exc, ConnectionClosedOK):
+                # a clean close is still a failure once connect() gives up
+                raise ConnectionClosedError(
+                    last_exc.rcvd, last_exc.sent, last_exc.rcvd_then_sent
+                ) from last_exc
             raise last_exc
 
+    def _too_many_failures(self, failures: int) -> bool:
+        return self.max_reconnects is not None and failures > self.max_reconnects
+
+    @staticmethod
+    def _is_transient(exc: Exception) -> bool:
+        """Whether a failed attempt to connect is worth retrying."""
+        if isinstance(exc, InvalidStatus):
+            status = exc.response.status_code
+            return status >= 500 or status == 429
+        if isinstance(exc, InvalidMessage):
+            # e.g. the server closed the connection during the handshake
+            return isinstance(exc.__cause__, EOFError)
+        if isinstance(exc, ssl.SSLCertVerificationError):
+            return False
+        # e.g. refused, unreachable, or timed out
+        return isinstance(exc, (OSError, asyncio.TimeoutError, ConnectionClosed))
+
+    def _reconnect_delay(self, attempt: int) -> float:
+        """
+        The wait before reconnect attempt `attempt`, counting from 0.
+
+        The base doubles from `reconnect_backoff` up to `max_reconnect_backoff`;
+        jitter only lengthens it, by up to 20%, and is clamped to the cap.
+        """
+        # bound the exponent so a long outage cannot overflow the float
+        base = min(
+            self.reconnect_backoff * 2.0 ** min(attempt, 1000),
+            self.max_reconnect_backoff,
+        )
+        return random.uniform(base, min(1.2 * base, self.max_reconnect_backoff))
+
+    async def _open(
+        self, connector: Awaitable[ClientConnection]
+    ) -> Optional[ClientConnection]:
+        """Open a connection, or return None if close is requested first."""
+        assert self._close_event is not None
+        attempt = asyncio.ensure_future(connector)
+        closing = asyncio.ensure_future(self._close_event.wait())
+        try:
+            await asyncio.wait({attempt, closing}, return_when=asyncio.FIRST_COMPLETED)
+        except BaseException:
+            attempt.cancel()
+            closing.cancel()
+            raise
+        stop = closing.done() or self._close_requested.is_set()
+        closing.cancel()
+        if not stop:
+            return attempt.result()
+        # close was requested: abandon the attempt, or close what it opened
+        attempt.cancel()
+        await asyncio.wait({attempt})
+        if not attempt.cancelled() and attempt.exception() is None:
+            await attempt.result().close()
+        return None
+
+    async def _wait_for_close(self, timeout: float) -> bool:
+        """Wait up to `timeout` seconds; return whether close was requested."""
+        assert self._close_event is not None
+        try:
+            await asyncio.wait_for(self._close_event.wait(), timeout)
+        except asyncio.TimeoutError:
+            pass
+        return self._close_requested.is_set()
+
     def _ssl_context(self) -> Optional[ssl.SSLContext]:
         # darwin needs some extra <3
         if not self.url.startswith("wss://"):
@@ -516,15 +678,51 @@ def unsubscribe_all(self):
             self.schedule_resub = True
         self._notify()
 
+    def request_close(self):
+        """
+        Close the websocket connection and make `connect()` return.
+
+        Unlike `close()`, this can be called from any thread, and it does not
+        wait: `connect()` returns soon after, including while it is
+        connecting or waiting to reconnect. If no `connect()` is running, the
+        next one returns without connecting; the request stays pending until a
+        `connect()` returns because of it.
+        """
+        logger.debug("close requested")
+        self._close_requested.set()
+        loop = self._loop
+        if loop is None:
+            return
+        try:
+            loop.call_soon_threadsafe(self._on_close_requested)
+        except RuntimeError:
+            # the loop is closed, so connect() has returned
+            pass
+
+    def _on_close_requested(self):
+        # runs on the event loop of connect(); the request may be stale, e.g.
+        # queued for an earlier connect() on the same loop
+        if not self._close_requested.is_set():
+            return
+        if self._close_event is not None:
+            self._close_event.set()
+        if self.websocket is not None:
+            self._close_task = asyncio.ensure_future(self.websocket.close())
+
     async def close(self):
         """
-        Close the websocket connection.
+        Close the websocket connection and stop `connect()` from reconnecting.
         """
         logger.debug("closing")
+        if self._loop is not None:
+            self._close_requested.set()
+            if self._close_event is not None:
+                self._close_event.set()
         self._flush_server_errors(force=True)
 
         if self.websocket:
             await self.websocket.close()
             self.websocket = None
-        else:
+        elif self._loop is None:
+            # (while connect() runs, it may be between connections)
             logger.warning("no websocket open to close")

test_websocket/test_fake_server.py

@@ -66,7 +66,7 @@ async def test_auth_failed_raises(self) -> None:
         self.assertEqual(self.server.inbox, [{"action": "auth", "params": "bad"}])
 
     async def test_resubscribes_after_reconnect(self) -> None:
-        c = self.client(subscriptions=["T.AAPL", "Q.AAPL"])
+        c = self.client(subscriptions=["T.AAPL", "Q.AAPL"], reconnect_backoff=0.01)
         task = asyncio.create_task(c.connect(lambda msgs: asyncio.sleep(0)))
         await self.server.wait_for_subscriptions({"T.AAPL", "Q.AAPL"})