diff --git a/src/schematic/client.py b/src/schematic/client.py index d15bf11..adf07b9 100644 --- a/src/schematic/client.py +++ b/src/schematic/client.py @@ -109,6 +109,8 @@ class DataStreamConfig: replicator_mode: bool = False replicator_health_url: Optional[str] = None replicator_health_check: Optional[int] = None + # Largest WebSocket message in bytes we accept; leave unset for the default. + max_message_size: Optional[int] = None @dataclass @@ -507,6 +509,8 @@ def __init__(self, api_key: str, config: Optional[AsyncSchematicConfig] = None): ds_opts.replicator_health_url = ds.replicator_health_url if ds.replicator_health_check is not None: ds_opts.replicator_health_check = ds.replicator_health_check + if ds.max_message_size is not None: + ds_opts.max_message_size = ds.max_message_size self._datastream_client = DataStreamClient(ds_opts) diff --git a/src/schematic/datastream/datastream_client.py b/src/schematic/datastream/datastream_client.py index ee0a8e6..90faa8e 100644 --- a/src/schematic/datastream/datastream_client.py +++ b/src/schematic/datastream/datastream_client.py @@ -17,7 +17,7 @@ from .merge import partial_company, partial_user from .rules_engine import RulesEngineClient from .types import DataStreamBaseReq, DataStreamReq, DataStreamResp, EntityType, KeyConflictError, MessageType, RulesEngineError -from .websocket_client import ClientOptions as WSClientOptions, DatastreamWSClient +from .websocket_client import MAX_MESSAGE_SIZE, ClientOptions as WSClientOptions, DatastreamWSClient _hints_cache: Dict[type, Dict[str, Any]] = {} @@ -103,6 +103,9 @@ class DataStreamClientOptions: replicator_health_url: Optional[str] = "http://localhost:8090/ready" replicator_health_check: int = DEFAULT_REPLICATOR_HEALTH_CHECK_MS + # Largest WebSocket message in bytes we accept; None removes the limit. + max_message_size: Optional[int] = MAX_MESSAGE_SIZE + # Event callbacks on_connected: Optional[Callable[[], None]] = None on_disconnected: Optional[Callable[[], None]] = None @@ -139,6 +142,7 @@ def __init__(self, options: DataStreamClientOptions) -> None: self._base_url = options.base_url self._logger = options.logger self._cache_ttl = options.cache_ttl + self._max_message_size = options.max_message_size # Callbacks self._on_connected = options.on_connected @@ -237,6 +241,7 @@ async def start(self) -> None: message_handler=self._handle_message, logger=self._logger, connection_ready_handler=self._handle_connection_ready, + max_message_size=self._max_message_size, on_connected=self._on_ws_connected, on_disconnected=self._on_ws_disconnected, on_ready=self._on_ready, diff --git a/src/schematic/datastream/websocket_client.py b/src/schematic/datastream/websocket_client.py index a4a5130..53540d2 100644 --- a/src/schematic/datastream/websocket_client.py +++ b/src/schematic/datastream/websocket_client.py @@ -40,6 +40,18 @@ MAX_RECONNECT_ATTEMPTS = 10 MIN_RECONNECT_DELAY = 1.0 # seconds MAX_RECONNECT_DELAY = 30.0 # seconds +# A connection only counts as healthy — and so only clears the backoff — once it +# has delivered a message or stayed up this long. Clearing on the handshake +# alone means a failure that always lands *after* connecting (an oversized +# frame, say) reconnects at the backoff floor forever instead of escalating. +HEALTHY_CONNECTION_THRESHOLD = 30.0 # seconds + +# Largest message we accept from the server. The websockets library defaults to +# 1 MiB, which the flags payload of a large environment exceeds; the library +# then closes the connection with a 1009 and we reconnect into the same frame. +# 100 MiB matches the Node SDK's `ws` default and still backstops a runaway +# payload. +MAX_MESSAGE_SIZE = 100 * 1024 * 1024 # 100 MiB # Headers attached to the WebSocket handshake so the backend can distinguish # direct-SDK connections from the schematic-datastream-replicator and correlate @@ -107,6 +119,8 @@ class ClientOptions: max_reconnect_attempts: int = MAX_RECONNECT_ATTEMPTS min_reconnect_delay: float = MIN_RECONNECT_DELAY max_reconnect_delay: float = MAX_RECONNECT_DELAY + # Maximum size in bytes of a message we accept; None removes the limit. + max_message_size: Optional[int] = MAX_MESSAGE_SIZE # Event callbacks — called on state transitions on_connected: Optional[Callable[[], None]] = None @@ -165,6 +179,7 @@ def __init__(self, options: ClientOptions) -> None: self._max_reconnect_attempts = options.max_reconnect_attempts self._min_reconnect_delay = options.min_reconnect_delay self._max_reconnect_delay = options.max_reconnect_delay + self._max_message_size = options.max_message_size # Event callbacks self._on_connected = options.on_connected @@ -268,9 +283,12 @@ async def _connect_and_read(self) -> None: # avoid conflicts with the Go server's gorilla/websocket. ping_interval=None, ping_timeout=None, + # Without this the library caps messages at 1 MiB and closes + # the connection with a 1009 on anything larger. + max_size=self._max_message_size, ) as ws: self._ws = ws - self._reconnect_attempts = 0 + connected_at = asyncio.get_event_loop().time() self._set_connected(True) # Run the ready handler before marking the client ready @@ -295,8 +313,16 @@ async def _connect_and_read(self) -> None: try: async for raw_message in ws: await self._handle_message(raw_message) + # A delivered message proves the connection works, + # so start the next backoff from scratch. + self._reconnect_attempts = 0 finally: self._stop_ping_pong() + # A long-lived connection counts as healthy too, even if + # the server never sent anything. + elapsed = asyncio.get_event_loop().time() - connected_at + if elapsed >= HEALTHY_CONNECTION_THRESHOLD: + self._reconnect_attempts = 0 self._logger.info("WebSocket connection closed") diff --git a/tests/datastream/test_websocket_client.py b/tests/datastream/test_websocket_client.py index de55522..394e110 100644 --- a/tests/datastream/test_websocket_client.py +++ b/tests/datastream/test_websocket_client.py @@ -12,6 +12,7 @@ from schematic.datastream.types import DataStreamBaseReq, DataStreamReq, DataStreamResp, EntityType from schematic.datastream.websocket_client import ( _WS_HEADERS_KWARG, + MAX_MESSAGE_SIZE, ClientOptions, DatastreamWSClient, convert_api_url_to_websocket_url, @@ -627,3 +628,213 @@ async def handler(msg): pass # At high attempt counts, delay should be capped at max + jitter ceiling delay = client._calculate_backoff_delay(20) assert delay <= 5.0 + 1.0 + + +# --------------------------------------------------------------------------- +# Message size limit +# --------------------------------------------------------------------------- + + +async def test_max_size_passed_to_websockets_connect() -> None: + """The connect call carries an explicit max_size, so we don't inherit the + websockets library's 1 MiB default.""" + captured_kwargs: dict = {} + ws = MockWebSocket(block_on_empty=True) + + @asynccontextmanager + async def capturing_connect(*args, **kwargs): + captured_kwargs.update(kwargs) + yield ws + + connected = asyncio.Event() + client, ws, _ = make_client(ws=ws, on_connected=lambda: connected.set()) + + with patch("schematic.datastream.websocket_client.websockets.connect", capturing_connect): + async with run_client(client): + await asyncio.wait_for(connected.wait(), timeout=2.0) + + assert captured_kwargs["max_size"] == MAX_MESSAGE_SIZE + assert MAX_MESSAGE_SIZE > 1024 * 1024 + + +async def test_max_message_size_is_configurable() -> None: + captured_kwargs: dict = {} + ws = MockWebSocket(block_on_empty=True) + + @asynccontextmanager + async def capturing_connect(*args, **kwargs): + captured_kwargs.update(kwargs) + yield ws + + connected = asyncio.Event() + client, ws, _ = make_client(ws=ws, max_message_size=None, on_connected=lambda: connected.set()) + + with patch("schematic.datastream.websocket_client.websockets.connect", capturing_connect): + async with run_client(client): + await asyncio.wait_for(connected.wait(), timeout=2.0) + + assert captured_kwargs["max_size"] is None + + +@asynccontextmanager +async def serve_one_message(payload: str) -> AsyncIterator[str]: + """Run a real websockets server that sends `payload` to each client. + + Yields the ws:// URL to connect to. + """ + import websockets as ws_lib + + async def handler(connection, *_args) -> None: + await connection.send(payload) + try: + await connection.wait_closed() + except Exception: + pass + + server = await ws_lib.serve(handler, "127.0.0.1", 0) + try: + port = next(iter(server.sockets)).getsockname()[1] + yield f"ws://127.0.0.1:{port}/datastream" + finally: + server.close() + await server.wait_closed() + + +def _oversized_payload() -> str: + """A datastream message larger than the websockets library's 1 MiB default.""" + message = json.dumps( + { + "entity_type": "rulesengine.Flags", + "message_type": "full", + "data": {"filler": "x" * (2 * 1024 * 1024)}, + } + ) + assert len(message.encode()) > 1024 * 1024 + return message + + +async def test_oversized_message_reaches_handler_over_a_real_connection() -> None: + """Regression for SCH-7098: a >1 MiB Flags frame used to be dropped by the + websockets library with a 1009, leaving the client in a reconnect loop.""" + payload = _oversized_payload() + received: List[DataStreamResp] = [] + + async def handler(m: DataStreamResp) -> None: + received.append(m) + + async with serve_one_message(payload) as url: + client = DatastreamWSClient( + ClientOptions( + url=url, + api_key="key", + message_handler=handler, + logger=logger, + min_reconnect_delay=0.0, + max_reconnect_delay=0.0, + ) + ) + async with run_client(client): + await wait_until(lambda: len(received) == 1, timeout=10.0) + + assert received[0].entity_type == "rulesengine.Flags" + assert len(received[0].data["filler"]) == 2 * 1024 * 1024 # type: ignore[index] + + +async def test_oversized_message_is_dropped_at_the_old_1mib_limit() -> None: + """Control for the test above: pin max_message_size back to the library + default and the same frame never reaches the handler.""" + payload = _oversized_payload() + received: List[DataStreamResp] = [] + connects: List[int] = [] + + async def handler(m: DataStreamResp) -> None: + received.append(m) + + async with serve_one_message(payload) as url: + client = DatastreamWSClient( + ClientOptions( + url=url, + api_key="key", + message_handler=handler, + logger=logger, + min_reconnect_delay=0.0, + max_reconnect_delay=0.0, + max_message_size=1024 * 1024, + on_connected=lambda: connects.append(1), + ) + ) + async with run_client(client): + await wait_until(lambda: len(connects) >= 2, timeout=10.0) + + assert received == [] + + +# --------------------------------------------------------------------------- +# Backoff escalation +# --------------------------------------------------------------------------- + + +async def test_backoff_escalates_when_failure_follows_a_successful_connect() -> None: + """A failure that always lands after the handshake must still escalate the + backoff — otherwise a poison frame reconnects at the floor forever.""" + attempts_seen: List[int] = [] + + @asynccontextmanager + async def connect_then_fail(*args, **kwargs): + attempts_seen.append(0) + yield MockWebSocket() + raise ConnectionError("closed right after connecting") + + async def handler(m: DataStreamResp) -> None: ... + + client = DatastreamWSClient( + ClientOptions( + url="wss://test.example.com/datastream", + api_key="key", + message_handler=handler, + logger=logger, + min_reconnect_delay=0.0, + max_reconnect_delay=0.0, + max_reconnect_attempts=4, + ) + ) + + with patch("schematic.datastream.websocket_client.websockets.connect", connect_then_fail): + async with run_client(client): + await wait_until(lambda: len(attempts_seen) >= 3, timeout=5.0) + assert client._reconnect_attempts >= 2 + + +async def test_delivered_message_clears_the_backoff() -> None: + """A connection that delivers a message is healthy, so the next failure + starts its backoff from zero again.""" + msg = json.dumps({"entity_type": "rulesengine.Company", "message_type": "full", "data": {"id": "c1"}}) + received: List[DataStreamResp] = [] + connects: List[int] = [] + + async def handler(m: DataStreamResp) -> None: + received.append(m) + + @asynccontextmanager + async def connect(*args, **kwargs): + connects.append(1) + if len(connects) == 1: + raise ConnectionError("first attempt fails") + yield MockWebSocket(messages=[msg], block_on_empty=True) + + client = DatastreamWSClient( + ClientOptions( + url="wss://test.example.com/datastream", + api_key="key", + message_handler=handler, + logger=logger, + min_reconnect_delay=0.0, + max_reconnect_delay=0.0, + max_reconnect_attempts=5, + ) + ) + + with patch("schematic.datastream.websocket_client.websockets.connect", connect): + async with run_client(client): + await wait_until(lambda: len(received) == 1) + assert client._reconnect_attempts == 0