Skip to content

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

Merged
TexasCoding merged 2 commits into
mainfrom
r3/W3-C
May 22, 2026
Merged

polish(ws): seq tracker sync hot path + orphan dispatch correlation + reconnect exc_info + recv hot-loop polish + _stop teardown order#374
TexasCoding merged 2 commits into
mainfrom
r3/W3-C

Conversation

@TexasCoding

Copy link
Copy Markdown
Owner

Summary

Round-3 audit closure wave W3-C: five low-severity WS subsystem polish/perf/correctness fixes. The recv hot loop avoids per-frame coroutine and wait_for allocations; orphan-subscribe acks no longer clobber sids the server has reused; reconnect failures preserve their root cause; _stop() follows the canonical close-then-drain order so the final frames reach iterator consumers.

Issues closed

Behavioral changes

  • SequenceTracker.track_sync(...) is a new sync method on the public surface (the async track wrapper remains for back-compat).
  • KalshiConnectionError from exhausted reconnects now exposes __cause__ (the last underlying failure). Code that exception-matches solely by message is unaffected; code that inspects __cause__ will start seeing populated values.
  • _stop() ordering change is observable as: the last in-flight frames now reach iterator consumers BEFORE the close sentinel, instead of being silently dropped.

Tests

  • tests/ws/test_sequence.pytest_issue_330_* covers track_sync happy path, forward gap, reset, duplicate, non-sequenced passthrough, and the async wrapper compat.
  • tests/ws/test_dispatch.pytest_issue_354_orphan_unsubscribe_ack_does_not_clobber_reused_sid exercises the orphan→reuse race; test_issue_354_legitimate_server_unsubscribe_still_tears_down is the regression guard for the cleanup path.
  • tests/ws/test_connection.pytest_issue_355_reconnect_logs_exc_info_and_chains_last_exc asserts both LogRecord.exc_info on per-attempt logs and __cause__ on the terminal raise.
  • tests/ws/test_client.pyTestIssue356RecvHotLoopTimeout covers behavior + source-level guard; TestIssue357StopOrdering covers runtime order + source-level guard.

Source

Round-3 independent audit closure plan, wave W3.

…o; recv wait_for; _stop order

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
@TexasCoding TexasCoding added polish Code-quality and DX improvements; non-functional ws WebSocket-related performance Performance / hot-path concern labels May 22, 2026
@claude

claude Bot commented May 22, 2026

Copy link
Copy Markdown

Code Review — PR #374: WS subsystem polish/perf (W3-C)

Reviewed all five changes. Overall this is solid, well-motivated work with good test coverage. A few issues worth addressing before merge.


Overview

Five targeted fixes across the WS subsystem: a sync fast-path for the sequence tracker, orphan-subscribe ack correlation, reconnect exception chaining, a per-frame allocation cut in the recv loop, and a teardown ordering fix. All five are behaviorally correct and the PR summary accurately describes what changed.


Issues

1. _stop() has no exception guard around connection.close() — subscribers can hang indefinitely

File: kalshi/ws/client.py, _stop()

If self._connection.close() raises (e.g. the underlying socket is already dead), _stop() propagates the exception before it reaches self._recv_task cancellation and _broadcast_sentinels(). Any async for consumer waiting on a subscription queue will hang forever — no sentinel is ever put.

# current
if self._connection:
    await self._connection.close()          # raises → everything below is skipped

if self._recv_task and not self._recv_task.done():
    ...
await self._broadcast_sentinels()           # never reached

Suggested fix — suppress errors from close the same way recv-task errors are suppressed:

if self._connection:
    with contextlib.suppress(Exception):
        await self._connection.close()

The reconnect path in connection.py already suppresses close errors for the same reason. _stop() should match.


2. _process_frame reaches into SequenceTracker._on_gap (private attribute)

File: kalshi/ws/client.py line 469

if gap is not None and self._seq_tracker._on_gap is not None:
    await self._seq_tracker._on_gap(gap)

This is package-internal code so it won't break users, but it couples the hot-path caller to SequenceTracker's internal representation. A simple way to keep the perf win while closing the coupling:

# sequence.py — add one method
async def fire_gap(self, gap: SequenceGap) -> None:
    if self._on_gap is not None:
        await self._on_gap(gap)

Then _process_frame becomes:

if gap is not None:
    await self._seq_tracker.fire_gap(gap)

No allocation overhead (the None guard stays inside SequenceTracker), and the private attribute stays private.


3. _pending_orphan_unsub is not exported or type-annotated in __init__

File: kalshi/ws/dispatch.py line 156–157

self._pending_orphan_unsub: set[int] = set()

This is an instance attribute initialized in __init__, which is correct. Not a blocker, but the inline comment cross-referencing _handle_server_unsubscribe already explains the invariant well — no additional change needed here.


Minor observations

  • last_exc = None when ws_max_retries=0 (connection.py): raise ... from None explicitly suppresses implicit context, which is semantically correct (there was no attempt to chain). Edge case, not a bug.
  • asyncio.timeout requires Python 3.11+: pyproject.toml requires >=3.12, so this is fine.
  • Source-level guard tests (inspect.getsource): Unusual but justified here — these are perf-critical invariants that are invisible to behavioral tests. The pattern is consistent with what was already in the test suite.
  • Log message improvement in reconnect ("attempt %d/%d failed"): Good UX win, easy to overlook.

Test coverage

Coverage is thorough:

  • test_sequence.py — happy path, gap, reset, duplicate, non-sequenced passthrough, async wrapper compat. ✓
  • test_dispatch.py — orphan-clobber race + regression guard for the legitimate teardown path. ✓
  • test_connection.py — asserts both exc_info on log records and __cause__ on the terminal raise. ✓
  • test_client.py — functional drain test + source-level guard for both #356 and #357. ✓

The only gap: there is no test covering what happens if connection.close() raises during _stop() (the issue raised above).


Summary

Item Severity Verdict
_stop() unguarded close() can skip sentinels Medium Fix before merge
_process_frame accesses _on_gap directly Low Worth a quick fire_gap method
Tests for _stop() error path Low Nice to have
Everything else Looks good

The sequencing logic, orphan correlation, exception chaining, and asyncio.timeout swap are all correct. The _stop() guard is the one thing worth fixing before this ships.

@TexasCoding

Copy link
Copy Markdown
Owner Author

Round-2 review addressed in 952e02a:

  • _stop() close() guarded: wrapped _connection.close() in try/except with logger.warning(..., exc_info=True); moved teardown + _broadcast_sentinels() into finally so a raising transport can no longer strand iterator consumers on closed queues.
  • fire_gap encapsulation: added SequenceTracker.fire_gap(gap) and switched _process_frame to use it instead of poking _seq_tracker._on_gap directly.
  • Test added: test_issue_357_stop_broadcasts_sentinels_when_close_raises mocks _connection.close() to raise and asserts sentinels still fire on every active subscription queue.

Verification:

uv run ruff check kalshi/ws/ tests/ws/    # All checks passed!
uv run mypy kalshi/ws/                     # no issues, 20 files
uv run pytest tests/ws/ -v -k 'TestIssue357'  # 3 passed
uv run pytest tests/ws/ -q                 # 395 passed

@claude

claude Bot commented May 22, 2026

Copy link
Copy Markdown

Code Review — PR #374: WS Subsystem Polish Wave W3-C

Overall: Well-scoped, high-quality wave. Each fix is small, clearly motivated by an issue, and independently testable. The changes are correctness improvements with measurable performance upside on the recv hot path. A few observations below.


sequence.py#330: track_sync fast path

Positive:

  • Clean separation: sync bookkeeping in track_sync, async dispatch responsibility delegated to the caller. The happy path (most frames) now allocates zero coroutines.
  • fire_gap is a good encapsulation — callers don't need to reach into _on_gap private state.
  • track() backward-compat wrapper is correct: it calls track_sync then awaits the gap if present.

Observation — fire_gap vs direct _on_gap call:
In client.py, _process_frame calls self._seq_tracker.fire_gap(gap) rather than checking self._on_gap is not None directly. That's the right call, but it adds one extra method dispatch per gap event. Negligible since gaps are rare, and the encapsulation is worth it.

Minor — track_sync docstring:

See :meth:track for the full state-machine semantics.

The :meth: cross-reference lacks backtick quoting for Sphinx (:meth:\track``). Non-blocking, but will render as plain text in generated docs.


client.py#356: asyncio.timeout in recv loop

Correct and safe. Project requires Python 3.12+, so asyncio.timeout is fully available. The context-manager form avoids the per-iteration Task wrapper that wait_for creates, which is the right choice for a thousands-of-frames-per-second hot path.

Consistency note: _stop() still uses asyncio.wait_for(self._recv_task, timeout=2.0) for the task join. That's fine — _stop() is not a hot path — but worth knowing if a future pass reviews task-join patterns.


client.py#357: _stop() close-then-drain ordering

Positive:

  • The try/finally structure guarantees sentinels broadcast even when close() raises — previously iterator consumers could hang indefinitely.
  • Corrects a real deadlock risk: the old ordering waited on _recv_task with a live connection, meaning the recv loop could never observe ConnectionClosed to exit.

One edge case — asyncio.CancelledError during close():
CancelledError is a BaseException in Python 3.8+, so it is NOT caught by except Exception. The finally block still runs (correct), but the CancelledError propagates outward. This means _broadcast_sentinels() and the recv-task join both run, but the _running = False line after the try/finally block does not, and the exception unwinds _stop(). Whether this is a problem depends on the caller of _stop() — but since contextlib.suppress is the common caller pattern, it's probably fine. Not a blocker, just worth a comment if this is reviewed again later.


connection.py#355: exc_info + __cause__

Clean. Correctly chains the last failure and includes exc_info=True on per-attempt logs.

Edge case — ws_max_retries = 0:
If the retry loop never runs (zero retries configured), last_exc remains None and the raise becomes raise KalshiConnectionError(...) from None. In Python, from None explicitly suppresses the exception context and sets __suppress_context__ = True. This is semantically odd — there was no previous exception to suppress — though in practice ws_max_retries = 0 is likely never a valid config. A simple guard would close this:

raise KalshiConnectionError(
    f"Max reconnect attempts ({self._config.ws_max_retries}) exceeded"
) from last_exc  # fine if last_exc is None; Python treats `from None` as suppress

Or, to avoid any surprise: only attach from last_exc when it's not None. Very minor; the current code is not wrong, just slightly surprising.


dispatch.py#354: Orphan unsubscribe correlation

Positive:

  • Marking the pending set BEFORE sending the unsub is the right ordering — prevents a race where the server ack arrives before the local marker is set.
  • Short-circuiting before the seq reset / orderbook removal is correct: the orphan ack should never clobber live state.

_pending_orphan_unsub persistence across reconnects:
MessageDispatcher is NOT recreated on reconnect — _handle_reconnect() calls _connection.reconnect() without rebuilding the dispatcher. The _pending_orphan_unsub set persists across reconnects. In practice this is benign (messages from the old connection can't arrive on the new WebSocket connection), but if the server reuses a SID and an unsubscribed ack arrives on the new connection for a SID that happened to be in the orphan set from before reconnect, it would incorrectly short-circuit. This window is effectively zero in practice, but it's worth either clearing the set in reset_all() / on reconnect, or noting this assumption explicitly.

The client_id is None log path refactor:
The refactor from early-return to conditional ternary + elif is equivalent and cleaner. No concerns.


Tests

Thorough. Coverage hits all stated behavioral changes:

  • test_issue_330_*: 6 cases covering sync fast path, gap, reset, duplicate, non-sequenced, and async wrapper backward compat. The explicit check that called == [] after track_sync (on_gap NOT invoked) is exactly the right invariant to lock in.
  • test_issue_354_*: Both the race scenario AND the regression guard for legitimate teardown. Good that both exist — the guard prevents future optimizers from accidentally merging the two paths.
  • test_issue_355_*: Verifies LogRecord.exc_info presence (not just message string) and __cause__ population. Solid.
  • test_issue_356_* + test_issue_357_*: Source-level guard tests using inspect.getsource are fragile in compiled-only environments, but this is an established pattern in this repo and the runtime behavioral tests accompany them — acceptable.

One gap: no test exercises the _pending_orphan_unsub set behavior across a reconnect cycle. Given the observation above, a test that reconnects and verifies the set is either empty or correctly populated post-reconnect would close that loop.


Summary

Area Assessment
Correctness ✅ All five fixes are correct for the stated race/ordering/performance issues
Performance track_sync + asyncio.timeout eliminate hot-path allocations
Security ✅ No concerns
Test coverage ✅ Comprehensive; one gap noted (orphan set across reconnect)
Conventions ✅ Follows project patterns: try/finally teardown, exc_info=True, explicit __cause__ chaining

Blockers: none. The above are observations/minor nits. Ready to merge as-is.

@TexasCoding
TexasCoding merged commit 8bef6f5 into main May 22, 2026
5 checks passed
@TexasCoding
TexasCoding deleted the r3/W3-C branch May 22, 2026 23:49
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment