diff --git a/CHANGELOG.md b/CHANGELOG.md index 5d1a05d..cbb9f2d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,30 @@ All notable changes to kalshi-sdk will be documented in this file. ## Unreleased +### WS `run_forever(stop_event=...)` cooperative shutdown (#177) + +`KalshiWebSocket.run_forever()` now accepts an optional +`stop_event: asyncio.Event | None = None` parameter. When set — typically +from a SIGINT handler via `add_signal_handler(SIGINT, stop.set)` — +`run_forever()` clears `_running`, closes the connection, and drains the +recv loop via its existing `not self._running` branch. The recv task is +NOT cancelled, so no `CancelledError` leaks out. + +```python +import asyncio, signal + +stop = asyncio.Event() +asyncio.get_running_loop().add_signal_handler(signal.SIGINT, stop.set) + +async with ws.connect() as session: + await session.subscribe_ticker(tickers=["EXAMPLE-25-T"]) + await session.run_forever(stop_event=stop) +``` + +No behavior change when `stop_event` is omitted — external cancellation +still propagates as before, and the #175 "missing subscription" guard +remains in place. + ### WS `run_forever()` raises on missing subscription (#175) `KalshiWebSocket.run_forever()` previously returned immediately when no diff --git a/docs/websockets.md b/docs/websockets.md index c3304ae..35a4b5a 100644 --- a/docs/websockets.md +++ b/docs/websockets.md @@ -128,6 +128,30 @@ async with ws.connect() as session: has landed in the session — a callback alone doesn't tell the server to send frames, and the previous silent-no-op behavior was a foot-gun (#175). +### Cooperative shutdown + +Pass an ``asyncio.Event`` to ``run_forever(stop_event=...)`` to terminate +the recv loop without raising ``CancelledError``. The canonical pattern +wires the event to ``SIGINT`` so Ctrl+C drains in-flight dispatches, +closes the WebSocket cleanly, and returns: + +```python +import asyncio +import signal + +stop = asyncio.Event() +asyncio.get_running_loop().add_signal_handler(signal.SIGINT, stop.set) + +async with ws.connect() as session: + await session.subscribe_ticker(tickers=["EXAMPLE-25-T"]) + await session.run_forever(stop_event=stop) +``` + +When the event fires, ``run_forever()`` clears ``_running``, closes the +connection, and awaits the recv loop's natural exit. No ``CancelledError`` +leaks out (#177). Without ``stop_event``, external cancellation still +propagates as before. + `on()` works both before and after `connect()`; callbacks registered before the socket opens are buffered and applied when the session starts. diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index f76266c..0a7a822 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -118,14 +118,28 @@ async def _stop(self) -> None: with contextlib.suppress(asyncio.CancelledError): await self._recv_task - # Send sentinels to all active queues - if self._sub_mgr: - for sub in self._sub_mgr.active_subscriptions.values(): - await sub.queue.put_sentinel() + await self._broadcast_sentinels() if self._connection: await self._connection.close() + async def _broadcast_sentinels(self) -> None: + """Put a shutdown sentinel on every active subscription queue. + + Iterator consumers see the sentinel and exit their ``async for`` + loops. ``MessageQueue.put_sentinel`` is idempotent — calling + twice puts two sentinels in the queue, the iterator reads the + first and raises ``StopAsyncIteration`` without ever touching + the second. + + Used by ``_stop()`` (post-cancel cleanup), the cooperative + shutdown branch of ``run_forever(stop_event=...)`` (#177), and + ``_recv_loop``'s fatal-error broadcast paths. + """ + if self._sub_mgr is not None: + for sub in self._sub_mgr.active_subscriptions.values(): + await sub.queue.put_sentinel() + def _ensure_recv_loop(self) -> None: """Start the recv_loop background task if not already running.""" if self._recv_task is None or self._recv_task.done(): @@ -187,9 +201,7 @@ async def _recv_loop(self) -> None: try: await inner except (KalshiBackpressureError, KalshiSubscriptionError): - if self._sub_mgr: - for sub in self._sub_mgr.active_subscriptions.values(): - await sub.queue.put_sentinel() + await self._broadcast_sentinels() except Exception: logger.debug( "Shielded dispatch raised during cancel cleanup", @@ -203,9 +215,7 @@ async def _recv_loop(self) -> None: logger.error( "Fatal WS error in recv loop", exc_info=True, ) - if self._sub_mgr: - for sub in self._sub_mgr.active_subscriptions.values(): - await sub.queue.put_sentinel() + await self._broadcast_sentinels() break except (json.JSONDecodeError, ValidationError, KeyError): # #83: genuinely-non-fatal per-message errors. Log with @@ -224,9 +234,7 @@ async def _recv_loop(self) -> None: "Unexpected error in recv loop; broadcasting sentinels", exc_info=True, ) - if self._sub_mgr: - for sub in self._sub_mgr.active_subscriptions.values(): - await sub.queue.put_sentinel() + await self._broadcast_sentinels() raise async def _process_frame(self, raw: str) -> None: @@ -321,9 +329,7 @@ async def _handle_reconnect(self) -> None: logger.error( "Reconnect failed: %s", reconnect_err, exc_info=True, ) - if self._sub_mgr: - for sub in self._sub_mgr.active_subscriptions.values(): - await sub.queue.put_sentinel() + await self._broadcast_sentinels() self._running = False async def _handle_seq_gap(self, gap: SequenceGap) -> None: @@ -526,7 +532,7 @@ def decorator( return func return decorator - async def run_forever(self) -> None: + async def run_forever(self, stop_event: asyncio.Event | None = None) -> None: """Block until the recv loop terminates. Use with the callback API. Requires at least one prior ``subscribe_*`` (or generic @@ -537,6 +543,26 @@ async def run_forever(self) -> None: have explicitly subscribed to, so a callback without a matching subscribe sees nothing. + :param stop_event: optional ``asyncio.Event`` used to terminate + ``run_forever()`` cooperatively (#177). When set — typically + from a signal handler such as ``add_signal_handler(SIGINT, + stop.set)`` — this method clears ``_running``, closes the + connection, and drains the recv loop. The recv loop sees + ``ConnectionClosed`` on its next read and exits via the + normal ``not self._running`` branch, NOT via cancellation, + so no ``CancelledError`` leaks out. When ``None`` (the + default) the method blocks on ``_recv_task`` directly and + external cancellation still propagates as before. + + External cancellation of ``run_forever()`` itself (e.g., + ``task.cancel()`` on the awaiting task) while ``stop_event`` + is provided still propagates — the cancellation cleans up the + internal ``stop_waiter`` task but does NOT trigger the + cooperative shutdown branch. ``_recv_task`` keeps running + until the session's ``__aexit__`` calls ``_stop()`` for the + full teardown. Use the event for graceful exit; rely on + ``__aexit__`` for hard cancellation. + :raises KalshiSubscriptionError: ``run_forever()`` was called before any ``subscribe_*`` request landed (formerly a silent no-op return — fixed in #175). @@ -550,7 +576,71 @@ async def run_forever(self) -> None: "@ws.on(channel) callback does not subscribe — the server " "only sends frames for channels you explicitly subscribe to." ) - await self._recv_task + if stop_event is None: + await self._recv_task + return + + # Race the recv loop against the stop event. Whichever finishes first + # wins; if stop wins, signal the recv loop to exit cleanly (via + # _running + connection.close()) so the loop's existing + # ConnectionClosed handler breaks instead of attempting reconnect. + # NOTE: we do NOT cancel _recv_task — cancellation would re-introduce + # the very CancelledError leak this hook exists to avoid (#177). + stop_waiter = asyncio.create_task(stop_event.wait()) + try: + done, _ = await asyncio.wait( + {self._recv_task, stop_waiter}, + return_when=asyncio.FIRST_COMPLETED, + ) + if stop_waiter in done: + # Cooperative shutdown takes priority — broadcast sentinels + # at the end even if _recv_task happens to also be in + # `done` (server-side close racing the stop signal). The + # recv loop's clean `not self._running: break` exits do + # NOT broadcast on their own; only its error paths do. So + # the simultaneous-completion case (both tasks done in the + # same wait round) MUST go through this branch to avoid + # leaving iterator consumers hanging. + # + # Clearing _running first means a still-running recv loop's + # ConnectionClosed branch hits `if not self._running: + # break` instead of reconnecting. We do NOT cancel + # _recv_task — cancellation would re-introduce the very + # CancelledError leak this hook exists to avoid (#177). + # + # Nested try/finally so a close() exception still drains + # the recv task AND broadcasts queue sentinels — otherwise + # iterators hang on consumers' `async for` even though the + # connection is gone. _stop() does the same sentinel + # broadcast on __aexit__; duplicated here so cooperative + # shutdown is complete even for callers who use + # run_forever() outside an `async with` block. + self._running = False + try: + if self._connection is not None: + await self._connection.close() + finally: + try: + if not self._recv_task.done(): + await self._recv_task + else: + # Already done. .result() re-raises any + # exception, surfaced as run_forever's exception. + self._recv_task.result() + finally: + await self._broadcast_sentinels() + elif self._recv_task in done: + # Server-side close / crash without stop event firing. The + # recv loop's error paths broadcast sentinels themselves; + # its clean ConnectionClosed break does not, but in that + # case the user didn't request shutdown and __aexit__ will + # handle the teardown. Just re-surface any exception. + self._recv_task.result() + finally: + if not stop_waiter.done(): + stop_waiter.cancel() + with contextlib.suppress(asyncio.CancelledError): + await stop_waiter # ------------------------------------------------------------------ # Orderbook convenience diff --git a/tests/ws/test_client.py b/tests/ws/test_client.py index 0c1bfa7..9bc2b38 100644 --- a/tests/ws/test_client.py +++ b/tests/ws/test_client.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +from typing import Any import pytest @@ -384,6 +385,84 @@ async def test_run_forever_without_subscription_raises( with pytest.raises(KalshiSubscriptionError, match="at least one active subscription"): await session.run_forever() + async def test_run_forever_with_stop_event_returns_cleanly( + self, + fake_ws, # type: ignore[no-untyped-def] + test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#177: stop_event signals cooperative shutdown. run_forever() + closes the connection and drains the recv loop without raising + CancelledError. The loop exits via its existing `not _running` + branch on the next ConnectionClosed, NOT via cancellation.""" + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect() as session: + await session.subscribe_ticker(tickers=["T1"]) + stop = asyncio.Event() + run_task = asyncio.create_task(session.run_forever(stop_event=stop)) + # Give the loop a tick to settle on the recv await. 100 ms is + # generous enough for resource-constrained CI runners while + # still keeping the test fast. + await asyncio.sleep(0.1) + assert not run_task.done() + # Trigger cooperative shutdown. + stop.set() + # Returns cleanly, no exception leaked. + await asyncio.wait_for(run_task, timeout=2.0) + assert run_task.exception() is None + assert session._connection is not None + assert session._connection.state == ConnectionState.CLOSED + + async def test_run_forever_with_pre_set_stop_event_returns_immediately( + self, + fake_ws, # type: ignore[no-untyped-def] + test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#177: a stop_event already set before run_forever() runs should + fire on the first scheduling tick — same cooperative-shutdown path, + just no wait. Guards against the race-free case being broken by + future refactors of the asyncio.wait() arrangement.""" + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect() as session: + await session.subscribe_ticker(tickers=["T1"]) + stop = asyncio.Event() + stop.set() + await asyncio.wait_for(session.run_forever(stop_event=stop), timeout=2.0) + assert session._connection is not None + assert session._connection.state == ConnectionState.CLOSED + + async def test_run_forever_with_stop_event_broadcasts_sentinels( + self, + fake_ws, # type: ignore[no-untyped-def] + test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#177 review fix: cooperative shutdown must broadcast queue + sentinels so iterator consumers exit `async for` cleanly. Without + this an iterator outside an `async with` block would hang on the + empty queue after the connection closed.""" + 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"]) + stop = asyncio.Event() + run_task = asyncio.create_task(session.run_forever(stop_event=stop)) + await asyncio.sleep(0.05) + stop.set() + await asyncio.wait_for(run_task, timeout=2.0) + + # The iterator must terminate without further input. Pre-fix it + # would block on the empty queue waiting for a sentinel that + # only _stop() (on __aexit__) was sending. + collected: list[Any] = [] + async def drain() -> None: + async for msg in stream: + collected.append(msg) + await asyncio.wait_for(drain(), timeout=2.0) + # No frames ever sent through fake_ws, so the iterator should + # have terminated immediately on the sentinel. + assert collected == [] + # --------------------------------------------------------------------------- # Error callback