polish(ws): seq tracker sync hot path + orphan dispatch correlation + reconnect exc_info + recv hot-loop polish + _stop teardown order

kalshi/ws/client.py

@@ -175,29 +175,51 @@ 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()
-        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.
-            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()
 
-        if self._connection:
-            await self._connection.close()
+        # #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. 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 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
@@ -282,13 +304,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 +467,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:
+                await self._seq_tracker.fire_gap(gap)
             if not ok:
                 # Gap detected — skip dispatching this message.
                 # The gap handler will trigger resync.

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.

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,20 +139,22 @@ 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()
             logger.info(
                 "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.",

kalshi/ws/sequence.py

@@ -46,76 +46,103 @@ 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:
+    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``.
+
+        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
+
+    async def fire_gap(self, gap: SequenceGap) -> None:
+        """Invoke the configured ``on_gap`` callback, if any.
 
-        return False
+        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.