From 3456bba9fd4528629bcd5ee033edb7400f08648a Mon Sep 17 00:00:00 2001 From: Jeff West Date: Thu, 21 May 2026 06:26:14 -0500 Subject: [PATCH 1/3] fix(ws): roll back seq watermark on pre-validation failure (#241) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A malformed orderbook_delta/snapshot or order_group_updates frame (ValidationError, KeyError, programming bug) previously left the seq tracker over-advanced past a frame that was never applied and never delivered. The next legitimate frame's seq matched the over-advanced expected value, gap detection silently skipped, and the local orderbook desynced from the server with no resync trigger. Money risk on any strategy reading from the local book. Fix: - Widen `_process_frame`'s try/except to cover pre-validation + apply + dispatch as one unit. Any exception rolls back the seq watermark to `prev_seq` and re-raises, so the recv-loop's existing malformed-frame handler logs + continues the loop. - Propagate ValidationError out of `MessageDispatcher.dispatch` instead of swallowing it. order_group_updates (sequenced but validated inside dispatch) now also reaches the rollback path. Tests: - tests/ws/test_seq_rollback.py — end-to-end: snapshot seq=1, malformed delta seq=2, good delta seq=3 -> gap handler fires with expected=2 received=3 (both orderbook_delta and order_group_updates). - tests/ws/test_dispatch.py — update existing test that asserted the old swallow contract; it now asserts the re-raise behavior. Closes #241 --- kalshi/ws/client.py | 52 ++++++---- kalshi/ws/dispatch.py | 8 +- tests/ws/test_dispatch.py | 13 ++- tests/ws/test_seq_rollback.py | 188 ++++++++++++++++++++++++++++++++++ 4 files changed, 237 insertions(+), 24 deletions(-) create mode 100644 tests/ws/test_seq_rollback.py diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index 44087eeb..2ef3ce6f 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -314,28 +314,29 @@ async def _process_frame(self, raw: str) -> None: # Only meaningful once we know we'll dispatch (and might roll back). tracked = True - # Check for orderbook messages — validate ONCE for the local - # manager, then hand the typed message off to dispatch via - # pre_validated so the dispatcher routes the same instance to - # the queue without re-running Pydantic. + # Pre-validate orderbook frames (so the manager applies typed data), + # then dispatch. Both steps must roll the seq watermark back if they + # fail — otherwise a malformed frame on a sequenced channel silently + # advances the watermark and the next legitimate frame's gap is + # missed, leaving the local orderbook desynced from the server (#241). pre_validated: BaseModel | None = None - if msg_type == "orderbook_snapshot" and self._orderbook_mgr: - snapshot = OrderbookSnapshotMessage.model_validate(data) - # #199: mutate in place — skip the O(n log n) sort + 2N - # OrderbookLevel allocation that the public ``apply_snapshot`` - # wrapper performs. Consumers materialize via ``get(ticker)`` - # on demand. - # Record per-sid ticker ownership so all-markets subscriptions - # can be torn down on gap recovery / unsubscribe (#189, #206). - self._orderbook_mgr._apply_snapshot_inplace(snapshot, sid=snapshot.sid) - pre_validated = snapshot - elif msg_type == "orderbook_delta" and self._orderbook_mgr: - delta = OrderbookDeltaMessage.model_validate(data) - # #199: in-place mutation only — see snapshot branch above. - self._orderbook_mgr._apply_delta_inplace(delta) - pre_validated = delta - try: + if msg_type == "orderbook_snapshot" and self._orderbook_mgr: + snapshot = OrderbookSnapshotMessage.model_validate(data) + # #199: mutate in place — skip the O(n log n) sort + 2N + # OrderbookLevel allocation that the public ``apply_snapshot`` + # wrapper performs. Consumers materialize via ``get(ticker)`` + # on demand. + # Record per-sid ticker ownership so all-markets subscriptions + # can be torn down on gap recovery / unsubscribe (#189, #206). + self._orderbook_mgr._apply_snapshot_inplace(snapshot, sid=snapshot.sid) + pre_validated = snapshot + elif msg_type == "orderbook_delta" and self._orderbook_mgr: + delta = OrderbookDeltaMessage.model_validate(data) + # #199: in-place mutation only — see snapshot branch above. + self._orderbook_mgr._apply_delta_inplace(delta) + pre_validated = delta + await self._dispatcher.dispatch(data, pre_validated=pre_validated) except KalshiBackpressureError as exc: # Dispatch failed -> the consumer never saw this message. Roll @@ -354,6 +355,17 @@ async def _process_frame(self, raw: str) -> None: if affected is not None: await self._sub_mgr.broadcast_error(affected.client_id, exc) raise + except Exception: + # #241: ValidationError, KeyError, or any programming bug between + # the seq advance and dispatch leaves the watermark over-advanced + # past a frame that was NEVER applied locally and NEVER delivered + # to the consumer. Roll back so the dropped seq surfaces as a + # forward gap on the next legitimate frame (gap handler then + # drives a real resubscribe + snapshot). Re-raise so the + # recv-loop's existing malformed-frame handler logs + continues. + if tracked and sid is not None and self._seq_tracker: + self._seq_tracker.rollback(sid, prev_seq) + raise async def _handle_reconnect(self) -> None: """Reconnect after ConnectionClosed and re-establish subscriptions. diff --git a/kalshi/ws/dispatch.py b/kalshi/ws/dispatch.py index 2fc5eaa1..a85c0661 100644 --- a/kalshi/ws/dispatch.py +++ b/kalshi/ws/dispatch.py @@ -229,7 +229,13 @@ async def dispatch( "Failed to parse %s message: %s", msg_type, type(exc).__name__, ) - return + # #241: re-raise so `_process_frame` can roll back the seq + # watermark for sequenced channels (order_group_updates, + # orderbook_*). Without this, a malformed sequenced frame + # silently advances the watermark and the next legitimate + # frame's gap is missed. The recv-loop's existing + # malformed-frame handler catches and continues the loop. + raise # Route to subscription queue sid = data.get("sid") diff --git a/tests/ws/test_dispatch.py b/tests/ws/test_dispatch.py index 79cf881e..f368e7c7 100644 --- a/tests/ws/test_dispatch.py +++ b/tests/ws/test_dispatch.py @@ -462,14 +462,21 @@ async def test_server_unsubscribe_unknown_sid_no_crash(self) -> None: raw = json.dumps({"type": "unsubscribed", "msg": {"sid": 999}}) await dispatcher.dispatch(json.loads(raw)) # should not crash - async def test_dispatch_message_without_sid(self) -> None: - """Messages without sid are logged but don't crash.""" + async def test_dispatch_malformed_message_raises_for_seq_rollback(self) -> None: + """#241: dispatch must propagate ValidationError so `_process_frame` + can roll back the seq watermark for sequenced channels. Previously + the dispatcher swallowed the error and returned, which silently + advanced the watermark past a never-delivered frame. + """ + from pydantic import ValidationError mgr = FakeSubManager() dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] + # ticker requires `sid`; omitting it triggers ValidationError. raw = json.dumps( {"type": "ticker", "msg": ticker_payload_dict(market_ticker="T", market_id="x")} ) - await dispatcher.dispatch(json.loads(raw)) # should not crash + with pytest.raises(ValidationError): + await dispatcher.dispatch(json.loads(raw)) @pytest.mark.asyncio diff --git a/tests/ws/test_seq_rollback.py b/tests/ws/test_seq_rollback.py new file mode 100644 index 00000000..6c2256df --- /dev/null +++ b/tests/ws/test_seq_rollback.py @@ -0,0 +1,188 @@ +"""Regression tests for seq-rollback on pre-validation failure (#241). + +#241: a malformed sequenced frame (ValidationError, KeyError, ...) must +NOT advance the seq watermark — otherwise the next legitimate frame's +seq matches the (over-advanced) expected value and gap detection +silently skips, leaving the local orderbook desynced from the server. + +These tests drive the fake server end-to-end so the actual +`_recv_loop` + `_process_frame` interaction is exercised, not a mocked +shape. Determinism: poll on observable state with bounded wait_for +rather than fixed sleeps. +""" +from __future__ import annotations + +import asyncio + +import pytest + +from kalshi.config import KalshiConfig +from kalshi.ws.client import KalshiWebSocket +from kalshi.ws.sequence import SequenceGap + + +@pytest.mark.asyncio +class TestValidationErrorDoesNotAdvanceSeq: + """#241: ValidationError on sequenced frame must roll back the watermark.""" + + async def test_malformed_orderbook_delta_keeps_watermark( + self, + fake_ws, # type: ignore[no-untyped-def] + test_auth, # type: ignore[no-untyped-def] + ) -> None: + """Snapshot seq=1, malformed delta seq=2, good delta seq=3. + + With the bug: seq=2's track() advanced last_seq to 2 BEFORE the + ValidationError raised, so seq=3 == expected and gap detection + never fires. The local orderbook is now silently desynced. + + With the fix: ValidationError rolls last_seq back to 1, so seq=3 + > expected=2 triggers a forward gap and the resync handler runs. + """ + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect() as session: + # Capture gaps WITHOUT triggering the real resubscribe machinery + # (would race the test for sids and resubscribe acks). + gaps: list[SequenceGap] = [] + + async def capture(gap: SequenceGap) -> None: + gaps.append(gap) + + assert session._seq_tracker is not None + session._seq_tracker._on_gap = capture + + await session.subscribe_orderbook_delta( + tickers=["T1"], maxsize=100, + ) + assert session._sub_mgr is not None + sid = next( + iter(session._sub_mgr.active_subscriptions.values()) + ).server_sid + assert sid is not None + + # 1. Clean snapshot, seq=1. + await fake_ws.send_to_all({ + "type": "orderbook_snapshot", "sid": sid, "seq": 1, + "msg": { + "market_ticker": "T1", "market_id": "x", + "yes": [["0.50", "100"]], "no": [], + }, + }) + + async def snapshot_landed() -> None: + while session._seq_tracker.peek(sid) != 1: + await asyncio.sleep(0.01) + await asyncio.wait_for(snapshot_landed(), timeout=2.0) + + # 2. Malformed delta seq=2 — bad `side` fails the Literal check. + await fake_ws.send_to_all({ + "type": "orderbook_delta", "sid": sid, "seq": 2, + "msg": { + "market_ticker": "T1", "market_id": "x", + "price": "0.51", "delta": "10", "side": "invalid", + }, + }) + + # 3. Good delta seq=3. Without #241 fix, last_seq is already 2 + # so seq=3 == expected — no gap, silent desync. With the fix, + # last_seq is still 1 → expected=2, received=3 → forward gap. + await fake_ws.send_to_all({ + "type": "orderbook_delta", "sid": sid, "seq": 3, + "msg": { + "market_ticker": "T1", "market_id": "x", + "price": "0.52", "delta": "5", "side": "yes", + }, + }) + + async def gap_fired() -> None: + while not gaps: + await asyncio.sleep(0.01) + await asyncio.wait_for(gap_fired(), timeout=2.0) + + assert len(gaps) == 1, ( + "#241 regression: malformed sequenced frame silently " + "advanced the seq watermark; next legitimate frame's " + "gap was missed." + ) + gap = gaps[0] + assert gap.sid == sid + assert gap.kind == "gap" + assert gap.expected == 2 + assert gap.received == 3 + + async def test_malformed_order_group_update_keeps_watermark( + self, + fake_ws, # type: ignore[no-untyped-def] + test_auth, # type: ignore[no-untyped-def] + ) -> None: + """Same shape for order_group_updates — dispatch-layer mirror (#241). + + order_group_updates is sequenced but validated inside the + dispatcher (no orderbook pre-validation), so the dispatcher + must propagate the ValidationError up to `_process_frame` for + rollback. Previously the dispatcher swallowed it -> silent + seq advance. + """ + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect() as session: + gaps: list[SequenceGap] = [] + + async def capture(gap: SequenceGap) -> None: + gaps.append(gap) + + assert session._seq_tracker is not None + session._seq_tracker._on_gap = capture + + await session._do_subscribe( + "order_group_updates", params={}, maxsize=100, + ) + assert session._sub_mgr is not None + sid = next( + iter(session._sub_mgr.active_subscriptions.values()) + ).server_sid + assert sid is not None + + # First good message, seq=1 — establishes the watermark. + # OrderGroupMessage requires `msg` with a typed payload; send a + # minimal-but-valid envelope. + await fake_ws.send_to_all({ + "type": "order_group_updates", "sid": sid, "seq": 1, + "msg": { + "event_type": "created", + "order_group_id": "g1", + "ts_ms": 1700000000000, + }, + }) + + async def first_landed() -> None: + while session._seq_tracker.peek(sid) != 1: + await asyncio.sleep(0.01) + await asyncio.wait_for(first_landed(), timeout=2.0) + + # Malformed seq=2 — wrong `msg` shape (string where dict expected). + await fake_ws.send_to_all({ + "type": "order_group_updates", "sid": sid, "seq": 2, + "msg": "garbage-not-a-dict", + }) + + # Good seq=3 — should trigger gap detection iff seq=2 didn't + # advance the watermark. + await fake_ws.send_to_all({ + "type": "order_group_updates", "sid": sid, "seq": 3, + "msg": { + "event_type": "limit_updated", + "order_group_id": "g1", + "ts_ms": 1700000000001, + }, + }) + + async def gap_fired() -> None: + while not gaps: + await asyncio.sleep(0.01) + await asyncio.wait_for(gap_fired(), timeout=2.0) + + assert len(gaps) == 1 + assert gaps[0].expected == 2 + assert gaps[0].received == 3 From 1cc516642603a003dbd6ceb245de5d6a461bf9c4 Mon Sep 17 00:00:00 2001 From: Jeff West Date: Thu, 21 May 2026 06:40:08 -0500 Subject: [PATCH 2/3] perf(ws): drop per-frame asyncio.Task + shield via cooperative pause (#245) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Hot path before: every WS frame allocated a Task, two Futures, and a contextvars copy via `asyncio.create_task(self._process_frame(raw))` + `asyncio.shield(inner)`, purely to make dispatch uninterruptible between recv and the next iteration. At >1k orderbook deltas/sec this dominated the recv-loop allocation budget and defeated the #199 in-place orderbook fast path the recv loop carefully preserves. Fix: - Replace cancel-based pause with cooperative pause. The recv loop polls `connection.recv()` with a 50ms timeout and checks a `_pause_pending` flag at the top of each iteration; on observing the request it sets `_pause_granted` and parks on `_resume_signal`. F-P-04 invariant (no in-flight frame dropped on pause) is preserved structurally: pause is only granted at the safe checkpoint between frames, so `_pause_recv_loop` blocks until any in-flight `_process_frame` returns. - Dispatch runs inline. No Task, no shield Future, no contextvars copy per frame. Only the wait_for TimerHandle remains, which is cancelled before firing on every non-quiet recv. - `_ensure_recv_loop` is now dual-purpose: resume a parked loop if one exists, otherwise spawn a fresh task. `_stop` sets `_resume_signal` so a parked loop wakes up and observes `_running=False` cleanly. Tests: - tests/ws/test_recv_loop_hardening.py — F-P-04 test renamed and updated to assert the new cooperative contract: pause_recv_loop blocks until in-flight dispatch finishes (positive assertion via pause_task.done() check before the dispatch barrier is released). Quiet-channel latency: bounded by _RECV_POLL_S (50ms). Busy-channel latency: zero — frame arrives before the timeout fires. Closes #245 --- kalshi/ws/client.py | 163 +++++++++++++++++---------- tests/ws/test_recv_loop_hardening.py | 27 +++-- 2 files changed, 123 insertions(+), 67 deletions(-) diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index 2ef3ce6f..3993aeb8 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -55,6 +55,14 @@ {1002, 1003, 1007, 1008, 1009, 1010} ) | frozenset(range(4000, 5000)) +# #245: cooperative-pause polling interval. Bounds the latency between a +# subscribe-driven pause request and the recv loop parking itself on a +# quiet channel. On a busy channel pause is granted between frames at +# zero added latency. 50ms is well under any subscribe ack timeout +# (5s default) and ~5x the typical Kalshi market-burst inter-frame +# spacing, so polling adds negligible CPU on idle. +_RECV_POLL_S: float = 0.05 + class KalshiWebSocket: """WebSocket client for real-time Kalshi market data. @@ -90,6 +98,18 @@ def __init__( self._running = False self._subscribe_lock = asyncio.Lock() self._pending_callbacks: list[tuple[str, Callable[..., Awaitable[None]]]] = [] + # #245: cooperative pause replaces the per-frame + # ``asyncio.create_task + asyncio.shield`` that previously made + # ``_process_frame`` uninterruptible. The recv loop now polls + # ``connection.recv()`` with a short timeout so it can observe + # ``_pause_pending`` between frames without an external cancel. + # On a busy channel pause is granted in microseconds (between + # frames); on a quiet channel it costs at most ``_RECV_POLL_S`` + # of latency. No Task / no shield Future / no contextvars copy + # per frame. + self._pause_pending = False + self._pause_granted = asyncio.Event() + self._resume_signal = asyncio.Event() # #209: pluggable JSON loader (None -> stdlib json.loads). Resolved # once here so the recv hot path doesn't keep dereferencing config. self._json_loads: Callable[[bytes | str], Any] = ( @@ -129,10 +149,22 @@ async def _start(self) -> None: async def _stop(self) -> None: """Stop the receive loop and close the connection.""" self._running = False + # #245: wake the recv loop if it's parked on `_resume_signal` + # (a pause request that never reached resume). Setting the + # signal lets it observe ``_running=False`` and exit cleanly. + self._resume_signal.set() if self._recv_task and not self._recv_task.done(): - self._recv_task.cancel() - with contextlib.suppress(asyncio.CancelledError): - await self._recv_task + # Closing the connection unblocks any in-flight ``recv()`` + # via ConnectionClosed; the loop then sees ``_running=False`` + # and exits at the top of its while. The poll-interval + # bounds the latency at ``_RECV_POLL_S`` even if close + # races the wait. + with contextlib.suppress(asyncio.CancelledError, Exception): + await asyncio.wait_for(self._recv_task, timeout=2.0) + if not self._recv_task.done(): + self._recv_task.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await self._recv_task await self._broadcast_sentinels() @@ -156,23 +188,42 @@ async def _broadcast_sentinels(self) -> None: await sub.queue.put_sentinel() def _ensure_recv_loop(self) -> None: - """Start the recv_loop background task if not already running.""" + """Start the recv_loop background task or resume a parked one. + + #245: ``_pause_recv_loop`` cooperatively parks the loop on + ``_resume_signal``; this method resumes it by setting the signal. + Falls back to creating a fresh task when no loop is running + (initial subscribe, or post-stop restart). + """ + if self._pause_pending: + self._pause_pending = False + self._pause_granted.clear() + self._resume_signal.set() + return if self._recv_task is None or self._recv_task.done(): self._recv_task = asyncio.create_task(self._recv_loop()) async def _pause_recv_loop(self) -> None: - """Cancel the recv_loop so subscribe can safely call recv. - - F-P-04: cancellation only fires while the loop is awaiting - ``connection.recv()``. The recv-then-process critical section is - shielded from cancellation so any frame already off the socket is - fully dispatched before the loop exits. + """Cooperatively park the recv_loop so subscribe can call recv. + + #245: replaces the previous cancel-based pause (which required + a per-frame ``asyncio.shield`` to keep the in-flight frame from + being dropped). The recv loop polls ``connection.recv()`` with + a ``_RECV_POLL_S`` timeout and checks ``_pause_pending`` between + polls. On a busy channel pause is granted in microseconds; on a + quiet channel the latency is bounded by ``_RECV_POLL_S``. + + F-P-04 invariant preserved: a frame in flight when pause is + requested finishes dispatching before pause is granted, because + the recv loop only sets ``_pause_granted`` at a safe point + between frames (no shield required). """ - if self._recv_task and not self._recv_task.done(): - self._recv_task.cancel() - with contextlib.suppress(asyncio.CancelledError): - await self._recv_task - self._recv_task = None + if self._recv_task is None or self._recv_task.done(): + return + self._pause_granted.clear() + self._resume_signal.clear() + self._pause_pending = True + await self._pause_granted.wait() async def _recv_loop(self) -> None: """Background task: read frames, dispatch, handle reconnect.""" @@ -180,11 +231,29 @@ async def _recv_loop(self) -> None: assert self._dispatcher is not None while self._running: + # #245: cooperative pause checkpoint. Set at a safe point + # between frames so the F-P-04 invariant (no in-flight frame + # dropped) holds without any shielding. + if self._pause_pending: + self._pause_granted.set() + await self._resume_signal.wait() + self._resume_signal.clear() + if not self._running: + break + continue + try: - raw = await self._connection.recv() + # Poll with a short timeout so ``_pause_pending`` and + # ``_running`` get re-checked on quiet channels. On a + # busy channel the timeout is cancelled before firing, + # so polling adds no overhead in the hot path. + raw = await asyncio.wait_for( + self._connection.recv(), timeout=_RECV_POLL_S, + ) + except asyncio.TimeoutError: + continue except asyncio.CancelledError: - # F-P-04: cancellation while awaiting recv = no frame read, - # no data lost. Safe to exit. + # External cancel (``_stop`` fallback). Exit cleanly. break except ConnectionClosed as e: # #197: classify the close code. The server returns codes @@ -216,57 +285,37 @@ async def _recv_loop(self) -> None: break continue - # F-P-04: shield the recv->dispatch critical section. Once a - # frame has been read off the socket, we MUST dispatch it before - # honoring cancellation, otherwise the frame is silently dropped - # (seq trackers miss the gap, ticker frames vanish, etc.). - # - # IMPORTANT: hold the inner task in a local. `asyncio.shield(coro)` - # creates a detached background task and returns control on outer - # cancel — without the explicit `await inner`, dispatch keeps - # running after `break` and races a subsequent subscribe. - inner = asyncio.create_task(self._process_frame(raw)) + # #245: inline dispatch. No Task / no shield / no contextvars + # copy per frame. F-P-04 protection comes from the cooperative + # pause: ``_pause_recv_loop`` blocks until the recv loop + # reaches the safe checkpoint above, so a pause request can + # never cancel us in the middle of ``_process_frame``. try: - await asyncio.shield(inner) + await self._process_frame(raw) except asyncio.CancelledError: - # Block until the shielded dispatch actually finishes, then - # honor cancellation. Loop exit now strictly post-dispatch. - # KalshiBackpressureError / KalshiSubscriptionError must - # still broadcast sentinels even when we're being cancelled — - # the consumer-visible invariant (no hanging iterators) holds - # unconditionally. Parse errors are best-effort logged. - try: - await inner - except (KalshiBackpressureError, KalshiSubscriptionError): - await self._broadcast_sentinels() - except Exception: - logger.debug( - "Shielded dispatch raised during cancel cleanup", - exc_info=True, - ) + # Only ``_stop``'s fallback cancel can reach here. Exit + # cleanly; sentinels are broadcast by the caller. break except (KalshiBackpressureError, KalshiSubscriptionError): - # #83: these signal real consumer-visible problems. Propagate - # via sentinel-broadcast so iterators see the failure and the - # loop exits cleanly rather than spinning on the same error. - logger.error( - "Fatal WS error in recv loop", exc_info=True, - ) + # #83: consumer-visible problem. Broadcast sentinels so + # iterators see the failure and the loop exits cleanly + # rather than spinning on the same error. + logger.error("Fatal WS error in recv loop", exc_info=True) await self._broadcast_sentinels() break except (json.JSONDecodeError, ValidationError, KeyError): - # #83: genuinely-non-fatal per-message errors. Log with - # traceback so users can debug, then continue. - logger.warning( - "Skipping malformed WS frame", exc_info=True, - ) + # #83: genuinely-non-fatal per-message errors. #241: when + # the failing frame was on a sequenced channel, the seq + # watermark was already rolled back inside + # ``_process_frame``'s own except, so a next-seq gap will + # still trigger a real resync. + logger.warning("Skipping malformed WS frame", exc_info=True) continue except Exception: # #83 escape hatch: anything else (AttributeError from a # user-supplied callback, TypeError, programming bug, ...) # MUST broadcast sentinels before propagating — otherwise - # consumers hang on their queues with no signal. Re-raise - # after broadcast so the task failure is still visible. + # consumers hang on their queues with no signal. logger.error( "Unexpected error in recv loop; broadcasting sentinels", exc_info=True, diff --git a/tests/ws/test_recv_loop_hardening.py b/tests/ws/test_recv_loop_hardening.py index e45a1a10..e5fb422b 100644 --- a/tests/ws/test_recv_loop_hardening.py +++ b/tests/ws/test_recv_loop_hardening.py @@ -163,15 +163,17 @@ async def test_unsubscribe_pushes_sentinel( @pytest.mark.asyncio class TestPauseDoesNotDropFrames: - async def test_inflight_frame_dispatched_when_recv_task_cancelled( + async def test_inflight_frame_dispatched_when_pause_requested( self, fake_ws, test_auth, # type: ignore[no-untyped-def] ) -> None: - """F-P-04: a frame mid-dispatch when the recv task is cancelled must - still be delivered. We monkey-patch the dispatcher's dispatch() with - a barrier so we can force cancellation precisely while it's running. - Without the asyncio.shield in _process_frame, the frame is lost. + """F-P-04: a frame mid-dispatch when pause is requested must still + be delivered. Pre-#245 this was enforced by an ``asyncio.shield`` + wrapping every frame; post-#245 the recv loop only sets + ``_pause_granted`` at the top-of-loop checkpoint, so + ``_pause_recv_loop`` cooperatively waits for the in-flight frame + to finish before returning — no shield, no per-frame Task. """ config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) ws = KalshiWebSocket(auth=test_auth, config=config) @@ -207,14 +209,19 @@ async def slow_dispatch( ) await asyncio.wait_for(dispatch_started.wait(), timeout=2.0) - # Now request a pause (cancellation). The shield must keep the - # dispatch alive. + # Now request a pause cooperatively. ``_pause_recv_loop`` must + # block until the recv loop reaches its safe checkpoint, which + # cannot happen until ``slow_dispatch`` returns. pause_task = asyncio.create_task(session._pause_recv_loop()) - # Give the cancel a chance to propagate + # Give the pause request a chance to propagate. await asyncio.sleep(0.05) + assert not pause_task.done(), ( + "#245: pause must wait for the in-flight frame to finish " + "dispatching before being granted" + ) - # Release the dispatch barrier — shielded coroutine runs to - # completion and puts the message on the queue. + # Release the dispatch barrier — process_frame returns, recv + # loop reaches the safe checkpoint, pause is granted. allow_dispatch_finish.set() await asyncio.wait_for(pause_task, timeout=2.0) From f544f82bada8b665a2ca32edbc1656bc595cfef3 Mon Sep 17 00:00:00 2001 From: Jeff West Date: Thu, 21 May 2026 06:45:28 -0500 Subject: [PATCH 3/3] perf(ws): cache materialized Orderbook on _BookState (#244) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every `orderbook_delta` consumed via the high-level `subscribe_book` iterator previously paid an O(n log n) sort on both sides plus 2N `OrderbookLevel` validations + 1 `Orderbook` validation per delta. For a 200-level book at 100 deltas/sec that was ~40k Decimal validations/sec per ticker, completely defeating the #199 in-place hot path the recv loop carefully preserves. Approach (smallest surface area): memoize the materialized Orderbook on `_BookState._cached`. `_apply_delta_inplace` / `_apply_snapshot_inplace` invalidate via `state.invalidate()` when a side actually mutates. Subsequent `to_orderbook()` / `get()` calls return the same cached instance until the next mutation. Snapshot re-application allocates a fresh `_BookState` so the cache resets naturally without any explicit teardown. Trade-off: the previous `is not` identity guarantee between concurrent reads is no longer a contract (was an implementation detail of the always-rebuild path). The mutation-safety invariant holds at the *value* level — handed-out Orderbooks are never mutated in place by the manager. Tests: - tests/ws/test_orderbook_perf.py — identity-stability assertions: consecutive `get(ticker)` returns the same object; delta / snapshot invalidate the cache; `apply_delta`'s return matches the subsequent `get()` (same cached instance). - tests/ws/test_orderbook.py — `test_returned_book_is_snapshot_not_live_view` updated to assert the value-level mutation invariant (previously-handed-out book unchanged after later deltas) rather than the identity-distinctness mechanism. Closes #244 --- kalshi/ws/client.py | 2 +- kalshi/ws/orderbook.py | 30 ++++++++++++-- tests/ws/test_orderbook.py | 22 ++++++++--- tests/ws/test_orderbook_perf.py | 69 +++++++++++++++++++++++++++++++++ 4 files changed, 113 insertions(+), 10 deletions(-) create mode 100644 tests/ws/test_orderbook_perf.py diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index 3993aeb8..f32fc9b1 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -250,7 +250,7 @@ async def _recv_loop(self) -> None: raw = await asyncio.wait_for( self._connection.recv(), timeout=_RECV_POLL_S, ) - except asyncio.TimeoutError: + except TimeoutError: continue except asyncio.CancelledError: # External cancel (``_stop`` fallback). Exit cleanly. diff --git a/kalshi/ws/orderbook.py b/kalshi/ws/orderbook.py index 3a5e3803..a72be611 100644 --- a/kalshi/ws/orderbook.py +++ b/kalshi/ws/orderbook.py @@ -33,20 +33,34 @@ class _BookState: (no ``market_tickers`` param) can be cleanly torn down on gap recovery or unsubscribe without the caller having to know which tickers the server happened to forward through that sid. + + #244: ``_cached`` memoizes the materialized :class:`Orderbook` so the + high-frequency ``subscribe_book`` iterator path doesn't pay an + O(n log n) sort + 2N :class:`OrderbookLevel` validations + 1 + :class:`Orderbook` validation on every delta. The cache is invalidated + by :meth:`_apply_delta_inplace` / :meth:`_apply_snapshot_inplace` + whenever a side mutates. Identity-stable: consecutive ``get(ticker)`` + calls without intervening mutations return the *same* object. """ ticker: str yes: dict[Decimal, Decimal] = field(default_factory=dict) no: dict[Decimal, Decimal] = field(default_factory=dict) sid: int | None = None + # #244: lazy memoization of the materialized public Orderbook view. + # ``None`` means "stale — rebuild on next read". + _cached: Orderbook | None = field(default=None, repr=False, compare=False) def to_orderbook(self) -> Orderbook: - """Build a fresh ``Orderbook`` snapshot from the current state. + """Build (or return cached) ``Orderbook`` snapshot from current state. Levels are emitted price-ascending to match the historical wire-order contract that the previous list-backed implementation maintained via - ``list.sort`` after every insert. + ``list.sort`` after every insert. #244: the result is cached and + invalidated only when ``_apply_*_inplace`` mutates a side. """ + if self._cached is not None: + return self._cached yes_levels = [ OrderbookLevel(price=price, quantity=qty) for price, qty in sorted(self.yes.items()) @@ -55,7 +69,13 @@ def to_orderbook(self) -> Orderbook: OrderbookLevel(price=price, quantity=qty) for price, qty in sorted(self.no.items()) ] - return Orderbook(ticker=self.ticker, yes=yes_levels, no=no_levels) + ob = Orderbook(ticker=self.ticker, yes=yes_levels, no=no_levels) + self._cached = ob + return ob + + def invalidate(self) -> None: + """Drop the cached :class:`Orderbook` view; next read re-materializes.""" + self._cached = None class OrderbookManager: @@ -162,8 +182,12 @@ def _apply_delta_inplace(self, msg: OrderbookDeltaMessage) -> bool: del levels[price] else: levels[price] = new_qty + # #244: a side mutated; drop the cached Orderbook view so + # the next `to_orderbook`/`get` rebuilds with the new side. + state.invalidate() elif delta > 0: levels[price] = delta + state.invalidate() return True diff --git a/tests/ws/test_orderbook.py b/tests/ws/test_orderbook.py index dac70edc..67a37153 100644 --- a/tests/ws/test_orderbook.py +++ b/tests/ws/test_orderbook.py @@ -161,9 +161,14 @@ def test_negative_delta_partial(self) -> None: def test_returned_book_is_snapshot_not_live_view(self) -> None: """Regression for #85: a book handed to a consumer must not mutate - when subsequent deltas are applied. Pin both the identity invariant - (each call returns a distinct ``Orderbook``) and the value invariant - (the previously-returned book's state is preserved verbatim). + when subsequent deltas are applied. + + Post-#244 the materialized view is cached on ``_BookState`` and + only rebuilt when ``_apply_*_inplace`` invalidates it. The + mutation-safety invariant is preserved at the *value* level: + previously-handed-out books are immutable snapshots whose lists + the manager never touches. Identity between concurrent reads is + now a side-effect of caching, not a contract. """ mgr = OrderbookManager() snap_book = mgr.apply_snapshot(make_snapshot(yes=[["0.50", "100"]])) @@ -173,18 +178,23 @@ def test_returned_book_is_snapshot_not_live_view(self) -> None: delta_book = mgr.apply_delta(make_delta(price="0.50", delta="50", side="yes")) assert delta_book is not None - # Distinct instances. + # Different snapshot generations -> different instances. assert snap_book is not delta_book # The first-handed-out book is frozen at its emission-time state. assert snap_book.yes[0].quantity == Decimal("100") assert delta_book.yes[0].quantity == Decimal("150") - # And ``get()`` is also a snapshot, not a live view. + # ``get()`` returns the current cached view (== delta_book here), + # but a *subsequent* delta must not mutate it in place. get_book = mgr.get("T") assert get_book is not None - assert get_book is not delta_book mgr.apply_delta(make_delta(price="0.50", delta="25", side="yes")) assert get_book.yes[0].quantity == Decimal("150") # unchanged + # And the next get re-materializes with the new state. + fresh = mgr.get("T") + assert fresh is not None + assert fresh is not get_book + assert fresh.yes[0].quantity == Decimal("175") def test_apply_delta_is_constant_time_in_book_depth(self) -> None: """Regression for #87: apply_delta must touch at most a constant diff --git a/tests/ws/test_orderbook_perf.py b/tests/ws/test_orderbook_perf.py new file mode 100644 index 00000000..5a5951ea --- /dev/null +++ b/tests/ws/test_orderbook_perf.py @@ -0,0 +1,69 @@ +"""Regression tests for #244 — OrderbookManager.get caches materialized view. + +#244: every ``orderbook_delta`` consumed via the high-level +``subscribe_book`` iterator previously paid an O(n log n) sort + 2N +``OrderbookLevel`` validations + 1 ``Orderbook`` validation per delta. +``_BookState.to_orderbook`` now memoizes the result and only rebuilds +when ``_apply_*_inplace`` invalidates the cache. + +Smoke tests, not microbenches: assert object identity to prove no +re-validation happens between mutations, and assert identity-bust +across mutations so stale data is never served. +""" +from __future__ import annotations + +from decimal import Decimal + +from kalshi.ws.orderbook import OrderbookManager +from tests.ws.test_orderbook import make_delta, make_snapshot + + +class TestOrderbookCaching: + def test_consecutive_get_returns_same_object_identity(self) -> None: + mgr = OrderbookManager() + mgr._apply_snapshot_inplace( + make_snapshot(yes=[["0.50", "100"]], no=[["0.50", "50"]]) + ) + first = mgr.get("T") + second = mgr.get("T") + third = mgr.get("T") + assert first is not None + # #244: identity-stable — no re-materialization without mutation. + assert first is second + assert second is third + + def test_delta_invalidates_cache(self) -> None: + mgr = OrderbookManager() + mgr._apply_snapshot_inplace( + make_snapshot(yes=[["0.50", "100"]]) + ) + before = mgr.get("T") + # Apply a delta that mutates the yes side. + mgr._apply_delta_inplace( + make_delta(price="0.51", delta="10", side="yes") + ) + after = mgr.get("T") + # Identity MUST differ — cache was invalidated. + assert before is not after + # And the new view sees the new level. + assert any(lvl.price == Decimal("0.51") for lvl in after.yes) # type: ignore[union-attr] + + def test_snapshot_resets_cache(self) -> None: + mgr = OrderbookManager() + mgr._apply_snapshot_inplace(make_snapshot(yes=[["0.50", "100"]])) + first = mgr.get("T") + # Re-snapshot replaces state entirely; new state has fresh cache. + mgr._apply_snapshot_inplace(make_snapshot(yes=[["0.60", "200"]], seq=2)) + second = mgr.get("T") + assert first is not second + assert any(lvl.price == Decimal("0.60") for lvl in second.yes) # type: ignore[union-attr] + + def test_public_apply_delta_returns_cached_view(self) -> None: + """``apply_delta`` is identity-stable with subsequent ``get`` calls.""" + mgr = OrderbookManager() + mgr.apply_snapshot(make_snapshot(yes=[["0.50", "100"]])) + after_delta = mgr.apply_delta( + make_delta(price="0.51", delta="10", side="yes") + ) + from_get = mgr.get("T") + assert after_delta is from_get