From 84ce4382c6f680243740a489f93936128bc6ea9c Mon Sep 17 00:00:00 2001 From: Jeff West Date: Fri, 22 May 2026 09:05:16 -0500 Subject: [PATCH 1/2] polish(ws): seq sync fast path; orphan correlation; reconnect exc_info; recv wait_for; _stop order MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Five low-severity polish/perf/correctness fixes against the WS subsystem, all sourced from the round-3 audit closure plan (wave W3). #330 (perf): ``SequenceTracker.track`` was ``async def`` solely so the rare gap path could ``await self._on_gap``. Every legitimate sequenced frame paid coroutine-object allocation + ``__await__`` teardown on the hot recv loop. Split into a synchronous ``track_sync(...) -> (bool, SequenceGap | None)`` that does all bookkeeping without awaiting; the recv loop calls it and only awaits ``_on_gap`` when a gap is returned. The original ``async def track`` is kept as a thin backwards-compat wrapper for callers and tests that exercise the old surface. #354 (correctness): ``_handle_orphan_subscribed`` sent ``unsubscribe`` for orphan sids with no correlation, and ``_handle_server_unsubscribe`` tore down state for ANY ``unsubscribed`` envelope unconditionally — so if the server reused that sid for a freshly-completed subscribe before the orphan ack landed, the new sub's seq watermark and orderbook books got clobbered. Track orphan sids in a ``_pending_orphan_unsub`` set; the unsubscribe handler short-circuits when the ack matches a pending orphan, leaving the now-legitimate owner's state intact. #355 (observability): ``ConnectionManager.reconnect`` logged each attempt failure at DEBUG with no ``exc_info`` and dropped the last exception, so the terminal ``KalshiConnectionError`` surfaced with no ``__cause__``. Operators saw "Max reconnect attempts exceeded" with no hint whether the root cause was auth, TLS, or DNS. Capture the last exception, include ``exc_info=True`` on the DEBUG attempt log, and chain the final raise via ``from last_exc`` — matches the REST transport's pattern. #356 (perf): the recv hot loop's per-frame ``asyncio.wait_for(recv(), timeout=_RECV_POLL_S)`` allocated a TimerHandle + Task wrap each iteration. Replaced with ``async with asyncio.timeout(_RECV_POLL_S)`` (Python 3.11+) which reuses the loop's underlying timer plumbing with a single context-manager allocation per frame. Behaviorally equivalent: ``TimeoutError`` is still raised on the quiet-channel poll-and-recheck path. #357 (correctness): ``_stop()`` previously cancelled the recv task and broadcast queue sentinels BEFORE closing the connection — so frames buffered inside the ``websockets`` library at the moment of ``_stop()`` could neither be drained (recv loop gone) nor signalled (queues already closed). Reordered to close-then-drain: close the connection first so in-flight ``recv()`` raises ``ConnectionClosed`` and the loop drains naturally; broadcast sentinels last so iterator consumers see the drained tail before the close sentinel terminates their ``async for``. Closes #330, #354, #355, #356, #357 --- kalshi/ws/client.py | 51 +++++++++++----- kalshi/ws/connection.py | 22 ++++++- kalshi/ws/dispatch.py | 44 +++++++++++--- kalshi/ws/sequence.py | 93 ++++++++++++++++------------ tests/ws/test_client.py | 118 ++++++++++++++++++++++++++++++++++++ tests/ws/test_connection.py | 41 +++++++++++++ tests/ws/test_dispatch.py | 75 +++++++++++++++++++++++ tests/ws/test_sequence.py | 74 ++++++++++++++++++++++ 8 files changed, 456 insertions(+), 62 deletions(-) diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index bc3aaa37..41b9824b 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -175,18 +175,32 @@ async def _start(self) -> None: raise async def _stop(self) -> None: - """Stop the receive loop and close the connection.""" + """Stop the receive loop and close the connection. + + #357: close-then-drain ordering. Closing the connection first + unblocks any in-flight ``recv()`` via ``ConnectionClosed`` and + lets the recv loop drain whatever frames the ``websockets`` + library has buffered before the close-ack arrives. Sentinels are + broadcast LAST so iterator consumers see those final drained + frames before the close sentinel, instead of having the recv + task cancelled and the queues closed out from under them. + """ 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() + + # #357: close the connection FIRST. The in-flight ``recv()`` + # raises ``ConnectionClosed``; the loop drains any buffered + # frames, sees ``_running=False``, and exits at the top of its + # while. + if self._connection: + await self._connection.close() + if self._recv_task and not self._recv_task.done(): - # 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. + # 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(): @@ -194,11 +208,11 @@ async def _stop(self) -> None: with contextlib.suppress(asyncio.CancelledError, Exception): await self._recv_task + # Sentinels last (#357): iterator consumers have already seen + # the post-close drained frames; the sentinel now terminates + # their ``async for`` cleanly. await self._broadcast_sentinels() - if self._connection: - await self._connection.close() - # Reset state AFTER awaiting close so teardown still has manager refs (#297). self._connection = None self._sub_mgr = None @@ -282,13 +296,16 @@ async def _recv_loop(self) -> None: continue try: + # #356: ``asyncio.timeout`` (3.11+) reuses the loop's + # underlying TimerHandle plumbing with a single context- + # manager allocation, avoiding the per-frame Task wrap + # that ``asyncio.wait_for`` performs on every iteration. # 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, - ) + # so polling adds minimal overhead in the hot path. + async with asyncio.timeout(_RECV_POLL_S): + raw = await self._connection.recv() except TimeoutError: continue except asyncio.CancelledError: @@ -442,9 +459,15 @@ async def _process_frame(self, raw: str) -> None: if sub: channel = sub.channel prev_seq = self._seq_tracker.peek(sid) - ok = await self._seq_tracker.track( + # #330: sync fast path. Avoids per-frame coroutine allocation + # on the happy path (seq == expected, the case on every + # legitimate sequenced frame). Only await on_gap when the + # tracker actually reports a gap. + ok, gap = self._seq_tracker.track_sync( sid, seq, msg_type if msg_type else channel ) + if gap is not None and self._seq_tracker._on_gap is not None: + await self._seq_tracker._on_gap(gap) if not ok: # Gap detected — skip dispatching this message. # The gap handler will trigger resync. diff --git a/kalshi/ws/connection.py b/kalshi/ws/connection.py index b5c0779c..5e856af0 100644 --- a/kalshi/ws/connection.py +++ b/kalshi/ws/connection.py @@ -170,6 +170,10 @@ async def reconnect(self) -> None: self._ws = None await self._set_state(ConnectionState.RECONNECTING) + # #355: capture the last failure so the final raise can chain it + # via ``from`` and operators can see the real root cause (auth, + # TLS, DNS) instead of a bare "max retries exceeded". + last_exc: BaseException | None = None for attempt in range(self._config.ws_max_retries): # #221 P2.1: match REST transport's AWS Full Jitter formula. # The previous ``base * 2**attempt + uniform(0, 0.5)`` reduced @@ -195,13 +199,25 @@ async def reconnect(self) -> None: self._ws = await self._open_socket() await self._set_state(ConnectionState.CONNECTED) return - except Exception: - logger.debug("Reconnect attempt %d failed", attempt + 1) + except Exception as exc: + # #355: keep DEBUG level so happy-path retries stay quiet, + # but include ``exc_info`` so the underlying failure (auth, + # TLS, DNS) is visible when DEBUG logging is enabled. The + # final raise below chains the last failure via ``from`` + # so it surfaces on the traceback even at default log + # levels — matches the REST transport's pattern. + last_exc = exc + logger.debug( + "Reconnect attempt %d/%d failed", + attempt + 1, + self._config.ws_max_retries, + exc_info=True, + ) continue await self._set_state(ConnectionState.CLOSED) raise KalshiConnectionError( f"Max reconnect attempts ({self._config.ws_max_retries}) exceeded" - ) + ) from last_exc async def close(self) -> None: """Gracefully close the WebSocket connection with code 1000. diff --git a/kalshi/ws/dispatch.py b/kalshi/ws/dispatch.py index cc5b2fbd..1c775821 100644 --- a/kalshi/ws/dispatch.py +++ b/kalshi/ws/dispatch.py @@ -59,6 +59,14 @@ def __init__( self._seq_tracker = seq_tracker self._orderbook_mgr = orderbook_mgr self._callbacks: dict[str, Callable[[Any], Awaitable[None]]] = {} + # #354: sids the client never owned but proactively unsubscribed + # (orphan subscribe acks). When the server's ``unsubscribed`` ack + # for one of these arrives, ``_handle_server_unsubscribe`` MUST + # short-circuit instead of running full teardown — otherwise a + # sid that the server has already reused for a freshly-completed + # subscribe would have its seq watermark and orderbook state + # clobbered by the late orphan ack. + self._pending_orphan_unsub: set[int] = set() def register_callback( self, channel: str, callback: Callable[[Any], Awaitable[None]] @@ -106,6 +114,21 @@ async def _handle_server_unsubscribe(self, data: dict[str, Any]) -> None: logger.debug("server unsubscribed envelope without sid: %s", data) return + # #354: correlate orphan-unsubscribe acks. If we proactively sent + # an unsubscribe for a sid the client never owned (orphan + # subscribe ack), the corresponding ``unsubscribed`` ack lands + # here. Short-circuit unconditionally: even when the server has + # since reused the sid for a freshly-completed legitimate + # subscribe (now in ``_sid_to_client``), this ack is for the + # ORPHAN, not the new owner. Running the teardown below would + # clobber the new sub's seq watermark and orderbook books. + if sid in self._pending_orphan_unsub: + self._pending_orphan_unsub.discard(sid) + logger.debug( + "server unsubscribed: orphan-correlation ack for sid %d", sid, + ) + return + client_id = self._sub_mgr._sid_to_client.pop(sid, None) if self._seq_tracker is not None: self._seq_tracker.reset(sid) @@ -116,13 +139,11 @@ async def _handle_server_unsubscribe(self, data: dict[str, Any]) -> None: if self._orderbook_mgr is not None: self._orderbook_mgr.remove_by_sid(sid) - if client_id is None: - logger.debug( - "server unsubscribed for unknown sid %d (already reaped)", sid - ) - return - - sub = self._sub_mgr._subscriptions.pop(client_id, None) + sub = ( + self._sub_mgr._subscriptions.pop(client_id, None) + if client_id is not None + else None + ) if sub is not None: # Wake any held iterator so async-for exits cleanly. await sub.queue.put_sentinel() @@ -130,6 +151,10 @@ async def _handle_server_unsubscribe(self, data: dict[str, Any]) -> None: "Server-initiated unsubscribe: channel=%s sid=%d client_id=%d", sub.channel, sid, client_id, ) + elif client_id is None: + logger.debug( + "server unsubscribed for unknown sid %d (already reaped)", sid, + ) async def _handle_orphan_subscribed(self, data: dict[str, Any]) -> None: """Release a server-side sid whose subscribe ack has no client mapping. @@ -164,6 +189,11 @@ async def _handle_orphan_subscribed(self, data: dict[str, Any]) -> None: # A normal subscribe completed and installed the mapping # before this handler ran — nothing to do. return + # #354: mark the sid as orphan-unsubscribed BEFORE sending so the + # eventual ``unsubscribed`` ack short-circuits in + # ``_handle_server_unsubscribe`` rather than clobbering a sid the + # server may have reused. + self._pending_orphan_unsub.add(sid) logger.warning( "Orphan subscribed ack for sid=%d (no client mapping); " "sending unsubscribe to release server-side state.", diff --git a/kalshi/ws/sequence.py b/kalshi/ws/sequence.py index 4d58c5d9..3da88160 100644 --- a/kalshi/ws/sequence.py +++ b/kalshi/ws/sequence.py @@ -46,76 +46,93 @@ def should_track(self, channel: str) -> bool: """Whether this channel type has sequence numbers.""" return channel in SEQUENCED_CHANNELS - async def track(self, sid: int, seq: int | None, channel: str) -> bool: - """Track a message's sequence number. - - Returns ``True`` if the caller should dispatch the message, - ``False`` if it should be dropped. - - Three-way distinction for sequenced channels (#196): - - 1. ``seq == expected`` — happy path. Advance watermark, dispatch. - 2. ``seq > expected`` — forward gap. Fire ``on_gap`` with - ``kind="gap"``, advance watermark to ``seq``, drop the - message (caller's resync handler will resubscribe). - 3. ``seq == last`` — exact duplicate. **Drop** without - dispatch. Previously dispatched-without-tracking, which - let a re-delivered delta be applied twice to the orderbook. - 4. ``seq < last`` — server-side reset that reused the sid (or - seq=1 after a high watermark). Fire ``on_gap`` with - ``kind="reset"``, rewind watermark to ``seq``, drop the - message so the resync handler can resubscribe from a clean - snapshot. Without this, every subsequent post-reset message - would be treated as another duplicate until the old - watermark was re-crossed, silently corrupting state. - - Non-sequenced channels and ``seq=None`` always return ``True``. + def track_sync( + self, sid: int, seq: int | None, channel: str + ) -> tuple[bool, SequenceGap | None]: + """Sync fast path for the recv hot loop (#330). + + Returns ``(dispatch_ok, gap)``. The bookkeeping (watermark + advance / rewind / duplicate detection) is done synchronously + here; if ``gap`` is non-None the caller MUST ``await`` the + ``on_gap`` callback itself. The happy path + (``seq == expected``) returns ``(True, None)`` and allocates no + coroutine, eliminating per-frame ``async def`` overhead on + thousands-of-frames-per-second sequenced channels. + + See :meth:`track` for the full state-machine semantics. """ if not self.should_track(channel) or seq is None: - return True + return True, None last = self._last_seq.get(sid) if last is None: # First message for this sid self._last_seq[sid] = seq - return True + return True, None expected = last + 1 if seq == expected: self._last_seq[sid] = seq - return True + return True, None if seq == last: # Exact duplicate. Drop instead of dispatching, so a redelivered # delta cannot be applied twice to the orderbook. logger.debug("Duplicate seq %d for sid %d (last=%d); dropping", seq, sid, last) - return False + return False, None if seq < last: # Server-side reset that preserved the sid (or seq=1 after a - # high watermark). Treat as a gap with kind="reset" so the - # gap handler clears local state and resubscribes from a - # fresh snapshot. Rewind the watermark so subsequent - # post-reset messages aren't silently dropped as "still old". + # high watermark). Rewind the watermark and signal the gap so + # the caller's handler can resubscribe from a clean snapshot. gap = SequenceGap(sid=sid, expected=expected, received=seq, kind="reset") logger.warning( "Sequence reset on sid=%d: last=%d got=%d", sid, last, seq, ) self._last_seq[sid] = seq - if self._on_gap is not None: - await self._on_gap(gap) - return False + return False, gap # Forward gap (seq > expected). gap = SequenceGap(sid=sid, expected=expected, received=seq, kind="gap") logger.warning("Sequence gap: sid=%d expected=%d got=%d", sid, expected, seq) self._last_seq[sid] = seq # Accept the new seq to continue tracking + return False, gap - if self._on_gap is not None: - await self._on_gap(gap) + async def track(self, sid: int, seq: int | None, channel: str) -> bool: + """Track a message's sequence number. + + Returns ``True`` if the caller should dispatch the message, + ``False`` if it should be dropped. - return False + Three-way distinction for sequenced channels (#196): + + 1. ``seq == expected`` — happy path. Advance watermark, dispatch. + 2. ``seq > expected`` — forward gap. Fire ``on_gap`` with + ``kind="gap"``, advance watermark to ``seq``, drop the + message (caller's resync handler will resubscribe). + 3. ``seq == last`` — exact duplicate. **Drop** without + dispatch. Previously dispatched-without-tracking, which + let a re-delivered delta be applied twice to the orderbook. + 4. ``seq < last`` — server-side reset that reused the sid (or + seq=1 after a high watermark). Fire ``on_gap`` with + ``kind="reset"``, rewind watermark to ``seq``, drop the + message so the resync handler can resubscribe from a clean + snapshot. Without this, every subsequent post-reset message + would be treated as another duplicate until the old + watermark was re-crossed, silently corrupting state. + + Non-sequenced channels and ``seq=None`` always return ``True``. + + Backwards-compat async wrapper around :meth:`track_sync` (#330). + The recv hot path should call ``track_sync`` directly and await + ``on_gap`` only when a gap is returned, to avoid per-frame + coroutine allocation. + """ + ok, gap = self.track_sync(sid, seq, channel) + if gap is not None and self._on_gap is not None: + await self._on_gap(gap) + return ok def peek(self, sid: int) -> int | None: """Return the current last-seen seq for ``sid``, or None if untracked. diff --git a/tests/ws/test_client.py b/tests/ws/test_client.py index ed09fe58..08b84a07 100644 --- a/tests/ws/test_client.py +++ b/tests/ws/test_client.py @@ -924,3 +924,121 @@ async def test_issue_332_next_subscribe_after_backpressure_starts_fresh( ) msg = await asyncio.wait_for(stream.__anext__(), timeout=2.0) assert msg.msg.market_ticker == "T2" + + +@pytest.mark.asyncio +class TestIssue356RecvHotLoopTimeout: + async def test_issue_356_recv_loop_uses_asyncio_timeout_not_wait_for( + self, + fake_ws, # type: ignore[no-untyped-def] + test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#356: per-frame ``asyncio.wait_for`` is replaced by + ``async with asyncio.timeout(...)`` in the recv hot loop. Verify + the loop still drains frames correctly under the new construct. + """ + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect() as session: + stream = await session.subscribe_ticker(tickers=["T1"]) + # Send a few frames in quick succession — the timeout-context + # path must drain them all without TimeoutError leaking out. + for i in range(5): + await fake_ws.send_to_all( + { + "type": "ticker", + "sid": 1, + "msg": ticker_payload_dict( + market_ticker="T1", + market_id="x", + yes_bid_dollars=str(50 + i), + ), + } + ) + received: list[Any] = [] + for _ in range(5): + received.append( + await asyncio.wait_for(stream.__anext__(), timeout=2.0) + ) + assert len(received) == 5 + + async def test_issue_356_recv_loop_source_uses_asyncio_timeout(self) -> None: + """#356 guard: ensure the source actually calls ``asyncio.timeout`` + and no longer wraps recv() in ``asyncio.wait_for``. Lock the + intent in source so a future revert can't regress silently. + """ + import inspect + + src = inspect.getsource(KalshiWebSocket._recv_loop) + assert "asyncio.timeout(_RECV_POLL_S)" in src, ( + "_recv_loop must use asyncio.timeout for the per-frame poll" + ) + assert "wait_for(self._connection.recv()" not in src, ( + "_recv_loop must NOT wrap recv() in asyncio.wait_for" + ) + + +@pytest.mark.asyncio +class TestIssue357StopOrdering: + async def test_issue_357_stop_closes_connection_before_sentinels( + self, + fake_ws, # type: ignore[no-untyped-def] + test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#357: ``_stop()`` must close the WS connection BEFORE broadcasting + queue sentinels. Verified by recording the order in which + ``ConnectionManager.close`` and ``MessageQueue.put_sentinel`` + run during teardown. + """ + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + + order: list[str] = [] + async with ws.connect() as session: + stream = await session.subscribe_ticker(tickers=["T1"]) + assert ws._connection is not None + assert ws._sub_mgr is not None + + real_close = ws._connection.close + + async def tracking_close() -> None: + order.append("close") + await real_close() + + ws._connection.close = tracking_close # type: ignore[method-assign] + + for sub in ws._sub_mgr.active_subscriptions.values(): + real_sentinel = sub.queue.put_sentinel + + async def tracking_sentinel( + _real: Any = real_sentinel, + ) -> None: + order.append("sentinel") + await _real() + + sub.queue.put_sentinel = tracking_sentinel # type: ignore[method-assign] + + # Drain the iterator so its task completes before teardown + # races; one quick read is enough. + assert stream is not None + + # On _stop teardown: close MUST run before any sentinel. + assert order, "Expected at least close + sentinel to run during _stop" + assert order[0] == "close", ( + f"_stop must close connection before broadcasting sentinels, got {order!r}" + ) + + async def test_issue_357_stop_source_close_precedes_broadcast(self) -> None: + """#357 source-level guard: ``_stop()`` calls ``_connection.close`` + before ``_broadcast_sentinels`` (so a future reorder cannot regress + the runtime ordering silently).""" + import inspect + + src = inspect.getsource(KalshiWebSocket._stop) + close_idx = src.find("self._connection.close()") + broadcast_idx = src.find("self._broadcast_sentinels()") + assert close_idx != -1, "_stop must close the connection" + assert broadcast_idx != -1, "_stop must broadcast sentinels" + assert close_idx < broadcast_idx, ( + "_stop must close the connection BEFORE broadcasting sentinels (#357)" + ) diff --git a/tests/ws/test_connection.py b/tests/ws/test_connection.py index 6e807c6d..85a2d0ce 100644 --- a/tests/ws/test_connection.py +++ b/tests/ws/test_connection.py @@ -372,6 +372,47 @@ async def test_reconnect_max_retries_exceeded( await mgr.reconnect() assert mgr.state == ConnectionState.CLOSED + async def test_issue_355_reconnect_logs_exc_info_and_chains_last_exc( + self, + test_auth: object, + caplog: pytest.LogCaptureFixture, + ) -> None: + """#355: Per-attempt DEBUG log MUST include ``exc_info`` and the + final ``KalshiConnectionError`` MUST chain the last attempt's + exception via ``__cause__`` — without it operators cannot tell + DNS vs TLS vs auth from a ten-retry burn. + """ + import logging as _logging + + config = KalshiConfig( + ws_base_url="ws://127.0.0.1:1", + timeout=1.0, + retry_base_delay=0.001, + retry_max_delay=0.005, + ws_max_retries=2, + ) + mgr = ConnectionManager(auth=test_auth, config=config) # type: ignore[arg-type] + with ( + caplog.at_level(_logging.DEBUG, logger="kalshi.ws"), + pytest.raises(KalshiConnectionError) as excinfo, + ): + await mgr.reconnect() + + # The final raise chains the last attempt's exception. + assert excinfo.value.__cause__ is not None, ( + "Max-retries raise must chain the last attempt's exception" + ) + + # Each attempt log carries exc_info (LogRecord.exc_info populated). + attempt_records = [ + r for r in caplog.records + if "Reconnect attempt" in r.getMessage() and "failed" in r.getMessage() + ] + assert attempt_records, "Expected DEBUG 'Reconnect attempt N/M failed' records" + assert all(r.exc_info is not None for r in attempt_records), ( + "Per-attempt failure log MUST include exc_info=True" + ) + async def test_reconnect_eventually_succeeds( self, fake_ws: FakeKalshiWS, test_auth: object ) -> None: diff --git a/tests/ws/test_dispatch.py b/tests/ws/test_dispatch.py index b4d4abc0..6239b506 100644 --- a/tests/ws/test_dispatch.py +++ b/tests/ws/test_dispatch.py @@ -673,3 +673,78 @@ def test_message_models_multivariate_lookup_key() -> None: # multivariate_market_lifecycle is sibling (different message type) -- must stay assert "multivariate_market_lifecycle" in MESSAGE_MODELS assert "multivariate" not in MESSAGE_MODELS # the original short form, now replaced + + +@pytest.mark.asyncio +async def test_issue_354_orphan_unsubscribe_ack_does_not_clobber_reused_sid() -> None: + """#354: An orphan-subscribe ack triggers an unsubscribe; the server's + eventual ``unsubscribed`` for that orphan sid MUST NOT tear down state + if the server has since reused the sid for a freshly-completed + subscribe. The dispatcher correlates by tracking orphan unsubscribes + in a pending set; the second ``unsubscribed`` arrives, sees the + marker, and short-circuits without resetting the now-owned sid. + """ + sent: list[dict[str, object]] = [] + + class _FakeConnection: + async def send(self, msg: dict[str, object]) -> None: + sent.append(msg) + + mgr = FakeSubManager() + mgr._connection = _FakeConnection() # type: ignore[attr-defined] + mgr._get_msg_id = lambda: 1 # type: ignore[attr-defined] + tracker = SequenceTracker() + dispatcher = MessageDispatcher( + sub_mgr=mgr, # type: ignore[arg-type] + seq_tracker=tracker, + ) + + # Step 1: orphan subscribed for sid=42 — dispatcher marks it pending + # and emits unsubscribe. + await dispatcher.dispatch( + {"type": "subscribed", "id": 1, "msg": {"sid": 42}} + ) + assert 42 in dispatcher._pending_orphan_unsub + assert sent and sent[-1]["params"] == {"sids": [42]} + + # Step 2: server reuses sid=42 for a new (legitimate) subscription; + # client now owns it and the seq tracker has fresh state. + new_sub = mgr.add(42, "orderbook_delta", client_id=99) + tracker.track_sync(42, 5, "orderbook_delta") + assert tracker.peek(42) == 5 + + # Step 3: late ``unsubscribed`` for the orphan sid lands. The guard + # must short-circuit instead of resetting the seq watermark or + # popping the new mapping. + await dispatcher.dispatch( + {"type": "unsubscribed", "id": 0, "msg": {"sid": 42}} + ) + assert 42 not in dispatcher._pending_orphan_unsub # marker consumed + assert mgr._sid_to_client.get(42) == 99 # mapping preserved + assert tracker.peek(42) == 5 # seq state preserved + assert new_sub.queue._closed is False # iterator NOT terminated + + +@pytest.mark.asyncio +async def test_issue_354_legitimate_server_unsubscribe_still_tears_down() -> None: + """#354 regression guard: the orphan-correlation guard MUST NOT + short-circuit a legitimate server-initiated unsubscribe for a sid we + own. The cleanup path (seq reset, sentinel) must still run. + """ + mgr = FakeSubManager() + sub = mgr.add(7, "orderbook_delta") + tracker = SequenceTracker() + tracker.track_sync(7, 3, "orderbook_delta") + dispatcher = MessageDispatcher( + sub_mgr=mgr, # type: ignore[arg-type] + seq_tracker=tracker, + ) + + await dispatcher.dispatch( + {"type": "unsubscribed", "id": 0, "msg": {"sid": 7}} + ) + + assert 7 not in mgr._sid_to_client + assert tracker.peek(7) is None # seq state reset + # Sentinel pushed so any iterator on `sub.queue` exits. + assert sub.queue._closed is True diff --git a/tests/ws/test_sequence.py b/tests/ws/test_sequence.py index bc88c17d..18b2108e 100644 --- a/tests/ws/test_sequence.py +++ b/tests/ws/test_sequence.py @@ -100,3 +100,77 @@ async def test_none_seq_on_sequenced_channel(self) -> None: """seq=None on a sequenced channel is treated as OK (snapshot/first).""" tracker = SequenceTracker() assert await tracker.track(1, None, "orderbook_delta") is True + + +def test_issue_330_track_sync_happy_path_no_gap() -> None: + """#330: sync fast path returns (True, None) for in-order frames without + allocating any coroutine. Verified by calling it as plain sync.""" + tracker = SequenceTracker() + assert tracker.track_sync(1, 1, "orderbook_delta") == (True, None) + assert tracker.track_sync(1, 2, "orderbook_delta") == (True, None) + assert tracker.track_sync(1, 3, "orderbook_delta") == (True, None) + + +def test_issue_330_track_sync_forward_gap_returns_gap() -> None: + """#330: forward gap returns (False, SequenceGap(kind='gap')) without + invoking on_gap — the caller awaits it. on_gap MUST NOT fire from + track_sync even when one is registered on the tracker.""" + called: list[SequenceGap] = [] + + async def on_gap(g: SequenceGap) -> None: + called.append(g) + + tracker = SequenceTracker(on_gap=on_gap) + tracker.track_sync(1, 1, "orderbook_delta") + ok, gap = tracker.track_sync(1, 5, "orderbook_delta") + assert ok is False + assert gap is not None + assert gap.kind == "gap" + assert gap.expected == 2 + assert gap.received == 5 + # Hot path responsibility: track_sync NEVER awaits on_gap. + assert called == [] + + +def test_issue_330_track_sync_reset_returns_gap() -> None: + """#330: backwards seq returns (False, SequenceGap(kind='reset')) and + rewinds the watermark, without awaiting the callback.""" + tracker = SequenceTracker() + tracker.track_sync(1, 10, "orderbook_delta") + ok, gap = tracker.track_sync(1, 1, "orderbook_delta") + assert ok is False + assert gap is not None and gap.kind == "reset" + assert tracker.peek(1) == 1 + + +def test_issue_330_track_sync_duplicate_no_gap() -> None: + """#330: exact duplicate returns (False, None) — drop without gap.""" + tracker = SequenceTracker() + tracker.track_sync(1, 1, "orderbook_delta") + tracker.track_sync(1, 2, "orderbook_delta") + ok, gap = tracker.track_sync(1, 2, "orderbook_delta") + assert ok is False + assert gap is None + + +def test_issue_330_track_sync_non_sequenced_passthrough() -> None: + """#330: non-sequenced channels and seq=None return (True, None).""" + tracker = SequenceTracker() + assert tracker.track_sync(1, None, "ticker") == (True, None) + assert tracker.track_sync(1, 5, "fill") == (True, None) + + +@pytest.mark.asyncio +async def test_issue_330_async_wrapper_still_awaits_on_gap() -> None: + """#330: the async track() wrapper must remain a working back-compat + surface — gap callbacks still fire when callers use the old API.""" + gaps: list[SequenceGap] = [] + + async def on_gap(g: SequenceGap) -> None: + gaps.append(g) + + tracker = SequenceTracker(on_gap=on_gap) + await tracker.track(1, 1, "orderbook_delta") + result = await tracker.track(1, 5, "orderbook_delta") + assert result is False + assert len(gaps) == 1 and gaps[0].kind == "gap" From 952e02aefa6ccd397aceaf210d2ad587824c6184 Mon Sep 17 00:00:00 2001 From: Jeff West Date: Fri, 22 May 2026 09:17:34 -0500 Subject: [PATCH 2/2] fix(ws): guard _stop() close() so sentinels always broadcast; encapsulate fire_gap --- kalshi/ws/client.py | 46 ++++++++++++++++++++++++----------------- kalshi/ws/sequence.py | 10 +++++++++ tests/ws/test_client.py | 42 +++++++++++++++++++++++++++++++++++++ 3 files changed, 79 insertions(+), 19 deletions(-) diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index 41b9824b..4e788ed2 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -194,24 +194,32 @@ async def _stop(self) -> None: # #357: close the connection FIRST. The in-flight ``recv()`` # raises ``ConnectionClosed``; the loop drains any buffered # frames, sees ``_running=False``, and exits at the top of its - # while. - if self._connection: - await self._connection.close() - - if self._recv_task and not self._recv_task.done(): - # 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() + # while. Guard the close in try/finally so a misbehaving + # transport (e.g. ``close()`` raising) cannot strand iterator + # consumers waiting on queues that never receive a sentinel. + try: + if self._connection: + await self._connection.close() + except Exception: + logger.warning( + "_connection.close() raised during _stop", exc_info=True + ) + finally: + if self._recv_task and not self._recv_task.done(): + # The poll-interval bounds the latency at ``_RECV_POLL_S`` + # even if close races the wait. with contextlib.suppress(asyncio.CancelledError, Exception): - await self._recv_task - - # Sentinels last (#357): iterator consumers have already seen - # the post-close drained frames; the sentinel now terminates - # their ``async for`` cleanly. - await self._broadcast_sentinels() + 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 + + # Sentinels last (#357): iterator consumers have already seen + # the post-close drained frames; the sentinel now terminates + # their ``async for`` cleanly. Inside ``finally`` so a raising + # close() still wakes consumers instead of hanging them. + await self._broadcast_sentinels() # Reset state AFTER awaiting close so teardown still has manager refs (#297). self._connection = None @@ -466,8 +474,8 @@ async def _process_frame(self, raw: str) -> None: ok, gap = self._seq_tracker.track_sync( sid, seq, msg_type if msg_type else channel ) - if gap is not None and self._seq_tracker._on_gap is not None: - await self._seq_tracker._on_gap(gap) + if gap is not None: + await self._seq_tracker.fire_gap(gap) if not ok: # Gap detected — skip dispatching this message. # The gap handler will trigger resync. diff --git a/kalshi/ws/sequence.py b/kalshi/ws/sequence.py index 3da88160..9fb77f63 100644 --- a/kalshi/ws/sequence.py +++ b/kalshi/ws/sequence.py @@ -134,6 +134,16 @@ async def track(self, sid: int, seq: int | None, channel: str) -> bool: await self._on_gap(gap) return ok + async def fire_gap(self, gap: SequenceGap) -> None: + """Invoke the configured ``on_gap`` callback, if any. + + Encapsulates ``_on_gap`` for callers using the sync fast path + (``track_sync``) so they do not need to reach into private state + to dispatch the gap notification. + """ + if self._on_gap is not None: + await self._on_gap(gap) + def peek(self, sid: int) -> int | None: """Return the current last-seen seq for ``sid``, or None if untracked. diff --git a/tests/ws/test_client.py b/tests/ws/test_client.py index 08b84a07..3ed771c8 100644 --- a/tests/ws/test_client.py +++ b/tests/ws/test_client.py @@ -1042,3 +1042,45 @@ async def test_issue_357_stop_source_close_precedes_broadcast(self) -> None: assert close_idx < broadcast_idx, ( "_stop must close the connection BEFORE broadcasting sentinels (#357)" ) + + async def test_issue_357_stop_broadcasts_sentinels_when_close_raises( + self, + fake_ws, # type: ignore[no-untyped-def] + test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#357 round-2: ``_stop()`` must broadcast sentinels even when + ``_connection.close()`` raises, otherwise iterator consumers hang + on queues whose recv loop is dead. The close() call is wrapped in + ``try/except``; the sentinel broadcast lives in ``finally``. + """ + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + + sentinels: list[str] = [] + + async with ws.connect() as session: + await session.subscribe_ticker(tickers=["T1"]) + assert ws._connection is not None + assert ws._sub_mgr is not None + + async def raising_close() -> None: + raise RuntimeError("transport boom") + + ws._connection.close = raising_close # type: ignore[method-assign] + + for sid, sub in ws._sub_mgr.active_subscriptions.items(): + real_sentinel = sub.queue.put_sentinel + + async def tracking_sentinel( + _real: Any = real_sentinel, _sid: int = sid, + ) -> None: + sentinels.append(f"sid={_sid}") + await _real() + + sub.queue.put_sentinel = tracking_sentinel # type: ignore[method-assign] + + # __aexit__ -> _stop(); close() raises but sentinels MUST still fire. + assert sentinels, ( + "Sentinels must broadcast even when _connection.close() raises; " + "otherwise iterator consumers hang waiting on the closed queue." + )