From 217eec50b6b5b17e49332edefee1b7e2a379988c Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Sat, 24 Oct 2026 00:22:08 -0600 Subject: [PATCH] feat: add opt-in WebSocket reconnect options connect() reconnects at once after a dropped connection, returns when the server closes the connection normally, retries a handshake error such as HTTP 404 forever, and counts every disconnect over the client's lifetime toward max_reconnects. A caller cannot change any of this. Add keyword arguments to WebSocketClient. Each defaults to today's behavior: - reconnect_backoff, max_reconnect_backoff: wait before reconnecting after an open connection closes. The wait doubles with each failure in a row, up to the cap, plus up to 20% jitter; at the cap it is drawn from [cap / 1.2, cap], never below reconnect_backoff. The constructor raises ValueError if reconnect_backoff is negative or greater than max_reconnect_backoff. - reconnect_on_close: also reconnect when the server closes normally. A close() by the caller while connected still ends connect(). - max_consecutive_failures, stable_uptime: limit failures in a row. A failure is an attempt that fails, or a connection that closes before it has been open for stable_uptime seconds. max_reconnects keeps its meaning. - raise_fatal_errors: use websockets' default process_exception, which raises errors that retrying will not fix. The reconnect loop is still async for ... in websockets' connect(). connect() now also accepts process_exception, which used to raise TypeError because the client passed its own. The new attribute max_connections_reached tells whether the server sent a max_connections status during the last handshake; the processor receives the same messages as before. --- massive/websocket/__init__.py | 112 ++++++++++++++++---- test_websocket/mock_server.py | 8 +- test_websocket/test_reconnect_options.py | 125 +++++++++++++++++++++++ 3 files changed, 225 insertions(+), 20 deletions(-) create mode 100644 test_websocket/test_reconnect_options.py diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py index 0f2d9610..5d2ebdbe 100644 --- a/massive/websocket/__init__.py +++ b/massive/websocket/__init__.py @@ -5,10 +5,11 @@ import logging import json import asyncio +import random import ssl import time from .models import * -from websockets.asyncio.client import connect, ClientConnection +from websockets.asyncio.client import connect, ClientConnection, process_exception from websockets.exceptions import ConnectionClosedOK, ConnectionClosedError from ..certs import ca_bundle_path from ..logging import get_logger @@ -36,6 +37,12 @@ def __init__( secure: bool = True, custom_json: Optional[Any] = None, ca_certs: Optional[str] = None, + reconnect_backoff: float = 0.0, + max_reconnect_backoff: float = 120.0, + reconnect_on_close: bool = False, + max_consecutive_failures: Optional[int] = None, + stable_uptime: float = 60.0, + raise_fatal_errors: bool = False, **kwargs, ): """ @@ -49,8 +56,20 @@ def __init__( :param max_reconnects: How many times to reconnect on network outage before ending .connect event loop. :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 an open connection closes. The wait doubles with each consecutive failure, up to max_reconnect_backoff, with random jitter. 0 (the default) reconnects at once. + :param max_reconnect_backoff: The longest wait before reconnecting, in seconds. + :param reconnect_on_close: Also reconnect when the server closes the connection normally (e.g. code 1000), counting toward max_reconnects. By default, .connect returns then. + :param max_consecutive_failures: How many failures in a row .connect allows before it raises the last error. A failure is a connection attempt that fails, or a connection that closes before it has been open for stable_uptime seconds. None (the default) allows any number. max_reconnects applies too. + :param stable_uptime: Seconds a connection must stay open to reset the failure count and the backoff. + :param raise_fatal_errors: Raise connection errors that retrying will not fix, such as an HTTP 404 during the handshake, instead of retrying them forever. Uses websockets' default process_exception. + :raises ValueError: If reconnect_backoff is negative or greater than max_reconnect_backoff. :return: A client. """ + if not 0 <= reconnect_backoff <= max_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" @@ -73,6 +92,15 @@ def __init__( self.subscribed = False self.subs: Set[str] = set() self.max_reconnects = max_reconnects + self.reconnect_backoff = reconnect_backoff + self.max_reconnect_backoff = max_reconnect_backoff + self.reconnect_on_close = reconnect_on_close + self.max_consecutive_failures = max_consecutive_failures + self.stable_uptime = stable_uptime + self.raise_fatal_errors = raise_fatal_errors + # whether the server sent a max_connections status during the last + # handshake: the account has too many connections open + self.max_connections_reached = False self.websocket: Optional[ClientConnection] = None if subscriptions is None: subscriptions = [] @@ -101,28 +129,45 @@ async def connect( :param processor: The callback to process messages. :param close_timeout: How long to wait for handshake when calling .close. + :param kwargs: Passed to websockets' connect(), e.g. process_exception. :raises AuthError: If invalid API key is supplied. """ # a new connection has no subscriptions, whatever an earlier call left self.subs = set() self.schedule_resub = True reconnects = 0 + # failed attempts and short-lived connections in a row + failures = 0 logger.debug("connect: %s", self.url) ssl_context = self._ssl_context() - last_exc = None + retry_or_raise = kwargs.pop( + "process_exception", + process_exception if self.raise_fatal_errors else lambda exc: None, + ) + + def on_connect_error(exc: Exception) -> Optional[Exception]: + nonlocal failures + failures += 1 + if self._too_many(failures, self.max_consecutive_failures): + return exc + return retry_or_raise(exc) + + last_exc: Optional[Exception] = None async for s in connect( self.url, close_timeout=close_timeout, ssl=ssl_context, - process_exception=lambda exc: None, + process_exception=on_connect_error, **kwargs, ): self.websocket = s + connected_at = None 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}) @@ -132,6 +177,11 @@ 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 = any( + m.get("status") == "max_connections" + for m in msg_parsed + auth_msg_parsed + ) + connected_at = time.monotonic() # e.g. max_connections, which the server sends before closing self._log_server_statuses(auth_msg_parsed) self._flush_server_errors(force=True) @@ -172,27 +222,51 @@ async def connect( logger.debug("connection closed (OK): %s", e) # the server sends a status (e.g. slow_consumer) before closing self._flush_server_errors(force=True) - return + # rcvd_then_sent: the server closed it, rather than close() + if not (self.reconnect_on_close and e.rcvd_then_sent): + return + last_exc = e except ConnectionClosedError as e: last_exc = e logger.debug("connection closed (ERR): %s", e) self._flush_server_errors(force=True) - reconnects += 1 - # keep scheduled_subs, which may hold changes not sent yet, - # and subscribe to all of it on the next connection - 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 - ): + reconnects += 1 + # keep scheduled_subs, which may hold changes not sent yet, + # and subscribe to all of it on the next connection + self.subs = set() + self.schedule_resub = True + if ( + connected_at is not None + and time.monotonic() - connected_at >= self.stable_uptime + ): + failures = 0 + else: + failures += 1 + if self._too_many(reconnects, self.max_reconnects) or self._too_many( + failures, self.max_consecutive_failures + ): + break + if self.reconnect_backoff > 0: + delay = self._reconnect_delay(failures) + logger.debug("reconnecting in %.1f s", delay) + await asyncio.sleep(delay) + + if last_exc is not None: raise last_exc + @staticmethod + def _too_many(count: int, limit: Optional[int]) -> bool: + return limit is not None and count > limit + + def _reconnect_delay(self, failures: int) -> float: + # doubles with each failure in a row, up to the cap; jitter adds up + # to 20%, and at the cap draws from [max(base, cap / 1.2), cap] instead + base, cap = self.reconnect_backoff, self.max_reconnect_backoff + delay = min(base * 2.0 ** min(max(failures - 1, 0), 64), cap) + if delay < cap: + return random.uniform(delay, min(delay * 1.2, cap)) + return random.uniform(max(base, cap / 1.2), cap) + def _ssl_context(self) -> Optional[ssl.SSLContext]: # darwin needs some extra <3 if not self.url.startswith("wss://"): diff --git a/test_websocket/mock_server.py b/test_websocket/mock_server.py index db00e84a..882bb4ec 100644 --- a/test_websocket/mock_server.py +++ b/test_websocket/mock_server.py @@ -137,8 +137,14 @@ async def mock_server(websocket): ) +def process_request(connection, request): + if request.path.startswith("/404/"): + # emulate the server refusing the handshake + return connection.respond(404, "Not Found\n") + + async def run_mock_server(): - async with serve(mock_server, "localhost", port): + async with serve(mock_server, "localhost", port, process_request=process_request): await asyncio.Future() # run forever diff --git a/test_websocket/test_reconnect_options.py b/test_websocket/test_reconnect_options.py new file mode 100644 index 00000000..96c03b8b --- /dev/null +++ b/test_websocket/test_reconnect_options.py @@ -0,0 +1,125 @@ +from massive import WebSocketClient +from base_ws import BaseTest +from mock_server import inbox, port +from unittest import mock +from websockets.exceptions import ( + ConnectionClosedError, + ConnectionClosedOK, + InvalidStatus, +) +import asyncio + + +class RecordSleeps: + # asyncio, except that sleep() records the delay and returns at once + def __init__(self): + self.delays = [] + + def __getattr__(self, name): + return getattr(asyncio, name) + + async def sleep(self, delay): + self.delays.append(delay) + + +async def ignore(msgs): + pass + + +class ReconnectOptionsTest(BaseTest): + def client(self, *subscriptions, **kwargs): + return WebSocketClient( + api_key=kwargs.pop("api_key", ""), + feed=kwargs.pop("feed", f"localhost:{port}"), + secure=False, + subscriptions=list(subscriptions), + **kwargs, + ) + + async def drop_until_raised(self, c, exc=ConnectionClosedError): + # DROP.X makes the server drop every connection right away + sleeps = RecordSleeps() + with mock.patch("massive.websocket.asyncio", sleeps), mock.patch( + "massive.websocket.random" + ) as rnd: + rnd.uniform.side_effect = lambda a, b: a + with self.assertRaises(exc): + await asyncio.wait_for(c.connect(ignore), 5) + return sleeps.delays, rnd.uniform.call_args_list + + async def test_reconnects_at_once_by_default(self): + c = self.client("DROP.X", max_reconnects=2) + delays, _ = await self.drop_until_raised(c) + self.assertEqual(delays, []) + self.assertEqual(len(inbox), 3) + + async def test_backoff_doubles_up_to_the_cap(self): + c = self.client( + "DROP.X", max_reconnects=4, reconnect_backoff=1, max_reconnect_backoff=4 + ) + delays, uniform = await self.drop_until_raised(c) + at_cap = mock.call(4 / 1.2, 4) + self.assertEqual( + uniform, [mock.call(1, 1.2), mock.call(2, 2.4), at_cap, at_cap] + ) + self.assertEqual(delays, [1, 2, 4 / 1.2, 4 / 1.2]) + + async def test_consecutive_failures_are_limited(self): + c = self.client("DROP.X", max_reconnects=None, max_consecutive_failures=2) + await self.drop_until_raised(c) + self.assertEqual(len(inbox), 3) + + async def test_stable_connection_resets_failures(self): + c = self.client( + "DROP.X", max_reconnects=3, max_consecutive_failures=0, stable_uptime=0 + ) + await self.drop_until_raised(c) + self.assertEqual(len(inbox), 4) + + async def test_reconnect_on_close(self): + # SLOW.X makes the server close every connection normally + c = self.client("SLOW.X", max_reconnects=2, reconnect_on_close=True) + with self.assertLogs("WebSocketClient", level="WARNING"): + await self.drop_until_raised(c, ConnectionClosedOK) + self.assertEqual(len(inbox), 3) + + async def test_close_still_returns_with_reconnect_on_close(self): + c = self.client("T.AAPL", reconnect_on_close=True) + task = asyncio.create_task(c.connect(ignore)) + await self.until(lambda: inbox and len(inbox[0]) == 2) + await c.close() + await asyncio.wait_for(task, 5) + self.assertEqual(len(inbox), 1) + + async def test_raise_fatal_errors(self): + c = self.client(feed=f"localhost:{port}/404", raise_fatal_errors=True) + with self.assertRaises(InvalidStatus): + await asyncio.wait_for(c.connect(ignore), 5) + + async def test_handshake_failures_count(self): + c = self.client(feed=f"localhost:{port}/404", max_consecutive_failures=0) + with self.assertRaises(InvalidStatus): + await asyncio.wait_for(c.connect(ignore), 5) + + async def test_process_exception_is_passed_to_websockets(self): + c = self.client(feed=f"localhost:{port}/404") + with self.assertRaises(InvalidStatus): + await asyncio.wait_for( + c.connect(ignore, process_exception=lambda exc: exc), 5 + ) + + async def test_max_connections_reached(self): + processor = mock.AsyncMock() + c = self.client(api_key="MAX-CONNECTIONS", raw=True, max_reconnects=0) + self.assertFalse(c.max_connections_reached) + with self.assertLogs("WebSocketClient", level="WARNING"): + with self.assertRaises(ConnectionClosedError): + await asyncio.wait_for(c.connect(processor), 5) + self.assertTrue(c.max_connections_reached) + processor.assert_not_called() + + async def test_backoff_is_checked(self): + with self.assertRaises(ValueError): + self.client(reconnect_backoff=-1) + with self.assertRaises(ValueError): + self.client(reconnect_backoff=10, max_reconnect_backoff=5)