From 8749ed047758c1f6844a0a66e79fb8401a381e8d Mon Sep 17 00:00:00 2001 From: Jeff West Date: Fri, 22 May 2026 08:21:23 -0500 Subject: [PATCH 1/2] perf(ws): orderbook materialization + apply_snapshot copy + zero-delta cache MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #327: `_BookState.to_orderbook` materialized levels via the validating `OrderbookLevel(...)` / `Orderbook(...)` constructors, re-running the per-level `DollarDecimal` / `FixedPointCount` validators on data that is already SDK-canonical (`Decimal` price -> `Decimal` quantity, sourced from `_levels_to_dict` or arithmetic on already-validated Decimals). Switching to `model_construct` for both the per-level and outer model removes 2N+1 Pydantic validations on every cache miss while preserving the wire-order contract via the existing `sorted(...)` walk. #344: public `OrderbookManager.apply_snapshot` previously routed through `_apply_snapshot_inplace` (which identity-adopts `msg.msg.yes/no` into the new `_BookState`) and then overwrote `state.yes` / `state.no` with `dict(msg.msg.yes)` / `dict(msg.msg.no)` to enforce the defensive-copy contract — two `_BookState.__setattr__` round-trips per side for one useful copy. Extracted the shared bookkeeping into `_install_snapshot` so the public path can install copies directly while the recv-loop `_apply_snapshot_inplace` keeps the v2.6 #296 identity-adoption perf win unchanged. #347: `_apply_delta_inplace` invalidated the #244 materialization cache unconditionally whenever the level already existed, including the `delta=0` case where the price-level dict is identical before and after. The cache invalidation now only fires when the dict actually mutates (`new_qty != existing_qty` or the level is deleted at zero/ negative new quantity), so `get()` keeps returning the same `Orderbook` instance across no-op deltas — which is what the #244 cache was designed to do. Closes #327, #344, #347 --- kalshi/ws/orderbook.py | 141 +++++++++++++++++++------------ tests/ws/test_orderbook.py | 166 +++++++++++++++++++++++++++++++++++++ 2 files changed, 256 insertions(+), 51 deletions(-) diff --git a/kalshi/ws/orderbook.py b/kalshi/ws/orderbook.py index a70f5e9..05675c8 100644 --- a/kalshi/ws/orderbook.py +++ b/kalshi/ws/orderbook.py @@ -59,16 +59,30 @@ def to_orderbook(self) -> Orderbook: contract that the previous list-backed implementation maintained via ``list.sort`` after every insert. #244: the result is cached and invalidated only when ``_apply_*_inplace`` mutates a side. + + #327: levels are materialized via :meth:`BaseModel.model_construct` + rather than the validating constructor. The ``yes`` / ``no`` dicts + are already SDK-canonical (``Decimal`` price -> ``Decimal`` quantity) + — they came either through ``_levels_to_dict`` on the snapshot wire + decode or through arithmetic on already-validated Decimals in + :meth:`OrderbookManager._apply_delta_inplace`. Re-running the + per-level ``DollarDecimal`` / ``FixedPointCount`` validators on every + materialization (O(n) cost on a hot path the #244 cache exists to + protect) was pure overhead. """ if self._cached is not None: return self._cached yes_levels = [ - OrderbookLevel(price=price, quantity=qty) for price, qty in sorted(self.yes.items()) + OrderbookLevel.model_construct(price=price, quantity=qty) + for price, qty in sorted(self.yes.items()) ] no_levels = [ - OrderbookLevel(price=price, quantity=qty) for price, qty in sorted(self.no.items()) + OrderbookLevel.model_construct(price=price, quantity=qty) + for price, qty in sorted(self.no.items()) ] - ob = Orderbook(ticker=self.ticker, yes=yes_levels, no=no_levels) + ob = Orderbook.model_construct( + ticker=self.ticker, yes=yes_levels, no=no_levels + ) self._cached = ob return ob @@ -113,28 +127,21 @@ def __init__(self) -> None: # number of tickers owned by the sid, not O(n) over every book. self._sid_tickers: dict[int, set[str]] = {} - def _apply_snapshot_inplace( - self, msg: OrderbookSnapshotMessage, sid: int | None = None - ) -> None: - """In-place state mutation for a snapshot. No materialization. - - Recv-loop hot path entry point (#199). The public - :meth:`apply_snapshot` wraps this with a ``to_orderbook()`` call - for direct callers that still want a fully built ``Orderbook``. - - ``sid`` records which subscription produced this snapshot so the - ticker can be torn down later via :meth:`remove_by_sid` — required - for all-markets subscriptions that don't pin tickers up front. - If omitted, defaults to ``msg.sid`` from the envelope. - - Ownership note: ``msg.msg.yes`` / ``.no`` are adopted by identity - into the new ``_BookState``; callers must treat them as consumed. + def _install_snapshot( + self, + ticker: str, + sid_val: int | None, + yes: dict[Decimal, Decimal], + no: dict[Decimal, Decimal], + ) -> _BookState: + """Replace the book for ``ticker`` with a freshly seeded state. + + The provided ``yes`` / ``no`` dicts are adopted by identity into + the new :class:`_BookState`; the caller is responsible for + defensive-copying when its input must not alias the manager's + internal state. See :meth:`apply_snapshot` vs. + :meth:`_apply_snapshot_inplace` for the two callers. """ - ticker = msg.msg.market_ticker - sid_val = sid if sid is not None else msg.sid - yes_levels = msg.msg.yes - no_levels = msg.msg.no - # If a prior book existed under a different sid (e.g. server moved # the ticker between subs), unindex it from the old sid first so # remove_by_sid on the stale id doesn't drop the live book. @@ -146,7 +153,7 @@ def _apply_snapshot_inplace( if not old_bucket: self._sid_tickers.pop(prior.sid, None) - state = _BookState(ticker=ticker, yes=yes_levels, no=no_levels, sid=sid_val) + state = _BookState(ticker=ticker, yes=yes, no=no, sid=sid_val) self._books[ticker] = state if sid_val is not None: bucket = self._sid_tickers.get(sid_val) @@ -157,10 +164,33 @@ def _apply_snapshot_inplace( logger.debug( "Orderbook snapshot: %s (%d yes, %d no levels, sid=%s)", ticker, - len(yes_levels), - len(no_levels), + len(yes), + len(no), sid_val, ) + return state + + def _apply_snapshot_inplace( + self, msg: OrderbookSnapshotMessage, sid: int | None = None + ) -> None: + """In-place state mutation for a snapshot. No materialization. + + Recv-loop hot path entry point (#199). The public + :meth:`apply_snapshot` wraps this with a ``to_orderbook()`` call + for direct callers that still want a fully built ``Orderbook``. + + ``sid`` records which subscription produced this snapshot so the + ticker can be torn down later via :meth:`remove_by_sid` — required + for all-markets subscriptions that don't pin tickers up front. + If omitted, defaults to ``msg.sid`` from the envelope. + + Ownership note: ``msg.msg.yes`` / ``.no`` are adopted by identity + into the new ``_BookState``; callers must treat them as consumed. + """ + sid_val = sid if sid is not None else msg.sid + self._install_snapshot( + msg.msg.market_ticker, sid_val, msg.msg.yes, msg.msg.no + ) def _apply_delta_inplace(self, msg: OrderbookDeltaMessage) -> bool: """In-place state mutation for a delta. No materialization. @@ -186,41 +216,50 @@ def _apply_delta_inplace(self, msg: OrderbookDeltaMessage) -> bool: new_qty = existing_qty + delta if new_qty <= 0: del levels[price] - else: + # #244: a side mutated; drop the cached Orderbook view so + # the next `to_orderbook`/`get` rebuilds with the new side. + state.invalidate() + elif new_qty != existing_qty: 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() + state.invalidate() + # #347: ``delta == 0`` (or otherwise nets to no change) leaves + # the price-level dict identical; preserve the #244 cache so + # consumers polling ``get()`` keep getting the same instance. elif delta > 0: levels[price] = delta state.invalidate() + # #347: ``delta <= 0`` at a non-existent level is a no-op too — + # the original code already skipped invalidation here; keep it so. return True def apply_snapshot(self, msg: OrderbookSnapshotMessage, sid: int | None = None) -> Orderbook: """Initialize (or reset) a book from a full snapshot. - Public wrapper: mutates in place via :meth:`_apply_snapshot_inplace` - and materializes a fresh :class:`Orderbook`. The recv loop bypasses - this and calls :meth:`_apply_snapshot_inplace` directly to avoid the - unused O(n log n) materialization (#199). - - ``msg`` is safe to reuse after this call: the manager defensively - copies the adopted dicts into ``_BookState`` so subsequent delta - mutations do not leak back through ``msg.msg.yes`` / ``.no``. The - recv loop, which never holds ``msg`` past dispatch, skips this - defensive copy via the direct inplace path. + Public wrapper: seeds the book with a defensive single-copy of the + caller's ``msg.msg.yes`` / ``.no`` dicts and materializes a fresh + :class:`Orderbook`. The recv loop bypasses this and calls + :meth:`_apply_snapshot_inplace` directly to avoid both the + defensive copy and the O(n log n) materialization (#199, #296). + + ``msg`` is safe to reuse after this call: the manager owns the + copied dicts, so subsequent delta mutations cannot leak back + through ``msg.msg.yes`` / ``.no``. + + #344: previously this routed through ``_apply_snapshot_inplace`` + (which identity-adopts) and then overwrote ``state.yes`` / ``.no`` + with a ``dict(...)`` copy — two assignments per side for one + useful copy. We now install the copies directly via + :meth:`_install_snapshot`, so each side is materialized into the + ``_BookState`` exactly once. """ - self._apply_snapshot_inplace(msg, sid=sid) - ticker = msg.msg.market_ticker - state = self._books[ticker] - # Defensive copy on the public path so the caller-supplied ``msg`` - # keeps its original ``.yes`` / ``.no`` dicts; ``_apply_snapshot_inplace`` - # adopted them by identity above, so we copy via ``msg.msg`` to make - # the "copy the caller's input" intent explicit. The recv-loop bypass - # keeps the identity-adopt perf win where it matters. - state.yes = dict(msg.msg.yes) - state.no = dict(msg.msg.no) + sid_val = sid if sid is not None else msg.sid + state = self._install_snapshot( + msg.msg.market_ticker, + sid_val, + dict(msg.msg.yes), + dict(msg.msg.no), + ) return state.to_orderbook() def apply_delta(self, msg: OrderbookDeltaMessage) -> Orderbook | None: diff --git a/tests/ws/test_orderbook.py b/tests/ws/test_orderbook.py index 4ca1753..b87ebef 100644 --- a/tests/ws/test_orderbook.py +++ b/tests/ws/test_orderbook.py @@ -505,3 +505,169 @@ def test_public_apply_snapshot_does_not_alias_input_msg(self) -> None: # Caller's held reference is untouched — the public-API safety contract. assert held_yes[Decimal("0.30")] == Decimal("5") + +class TestWave2CorrectnessRegressions: + """Round-3 W2 perf-correctness regressions for #327, #344, #347.""" + + def test_issue_327_to_orderbook_no_revalidation( + self, monkeypatch: object + ) -> None: + """``_BookState.to_orderbook`` materializes via ``model_construct``. + + The validating constructors (``OrderbookLevel(...)`` / + ``Orderbook(...)``) re-run the per-level ``DollarDecimal`` / + ``FixedPointCount`` validators on data that is already SDK-canonical + — wasted work on the high-frequency ``subscribe_book`` path that + the #244 cache exists to protect. + + We spy on the validating ``__init__`` methods (which Pydantic's + ``model_construct`` bypasses) and assert zero calls during a + ``apply_snapshot`` + ``apply_delta`` cycle. + """ + import pytest + + from kalshi.models.markets import Orderbook, OrderbookLevel + from kalshi.ws import orderbook as ob_mod + + calls = {"level": 0, "ob": 0} + orig_level_init = OrderbookLevel.__init__ + orig_ob_init = Orderbook.__init__ + + def spy_level_init(self: object, **kw: object) -> None: + calls["level"] += 1 + orig_level_init(self, **kw) # type: ignore[arg-type] + + def spy_ob_init(self: object, **kw: object) -> None: + calls["ob"] += 1 + orig_ob_init(self, **kw) # type: ignore[arg-type] + + mp = pytest.MonkeyPatch() + mp.setattr(OrderbookLevel, "__init__", spy_level_init) + mp.setattr(Orderbook, "__init__", spy_ob_init) + try: + mgr = ob_mod.OrderbookManager() + book = mgr.apply_snapshot( + make_snapshot( + yes=[["0.50", "100"], ["0.55", "200"], ["0.60", "300"]], + no=[["0.40", "150"], ["0.35", "75"]], + ) + ) + assert len(book.yes) == 3 and len(book.no) == 2 + # Force re-materialization by invalidating and reading again. + delta_book = mgr.apply_delta( + make_delta(price="0.50", delta="50", side="yes") + ) + assert delta_book is not None + finally: + mp.undo() + + assert ( + calls["level"] == 0 + ), f"OrderbookLevel validating __init__ ran {calls['level']} times" + assert ( + calls["ob"] == 0 + ), f"Orderbook validating __init__ ran {calls['ob']} times" + + def test_issue_344_apply_snapshot_public_single_copy(self) -> None: + """Public ``apply_snapshot`` assigns each side dict exactly once. + + Pre-fix the public path identity-adopted ``msg.msg.yes`` / ``.no`` + via :meth:`_apply_snapshot_inplace` and then overwrote with + ``dict(msg.msg.yes)`` / ``dict(msg.msg.no)`` — two assignments + per side for one useful copy. We spy on ``_BookState.__setattr__`` + and assert exactly one ``yes`` and one ``no`` assignment. + """ + import pytest + + from kalshi.ws import orderbook as ob_mod + + side_assignments: list[tuple[str, int]] = [] + orig_setattr = ob_mod._BookState.__setattr__ + + def spy_setattr(self: object, name: str, value: object) -> None: + if name in ("yes", "no"): + side_assignments.append((name, id(value))) + orig_setattr(self, name, value) # type: ignore[arg-type] + + mp = pytest.MonkeyPatch() + mp.setattr(ob_mod._BookState, "__setattr__", spy_setattr) + try: + msg = make_snapshot( + yes=[["0.50", "100"], ["0.55", "200"]], + no=[["0.45", "150"]], + ) + mgr = ob_mod.OrderbookManager() + book = mgr.apply_snapshot(msg) + assert book.ticker == "T" + finally: + mp.undo() + + yes_assigns = [a for a in side_assignments if a[0] == "yes"] + no_assigns = [a for a in side_assignments if a[0] == "no"] + assert ( + len(yes_assigns) == 1 + ), f"expected 1 yes assignment, got {len(yes_assigns)}: {yes_assigns}" + assert ( + len(no_assigns) == 1 + ), f"expected 1 no assignment, got {len(no_assigns)}: {no_assigns}" + + # And the defensive-copy contract still holds. + state = mgr._books["T"] + assert state.yes is not msg.msg.yes + assert state.no is not msg.msg.no + assert state.yes == {Decimal("0.50"): Decimal("100"), Decimal("0.55"): Decimal("200")} + + def test_issue_347_zero_delta_preserves_cache(self) -> None: + """A delta that nets to no change must not invalidate the #244 cache. + + ``_apply_delta_inplace`` previously called ``state.invalidate()`` + unconditionally when the level already existed, even when ``delta`` + was zero (``new_qty == existing_qty``). Cache-identity for a + no-op delta is preserved by post-fix. + """ + mgr = OrderbookManager() + mgr.apply_snapshot(make_snapshot(yes=[["0.50", "100"]])) + state = mgr._books["T"] + + # Populate the #244 cache. + first = mgr.get("T") + assert first is not None + cached = state._cached + assert cached is first + + # Apply a zero delta: the dict must be byte-identical afterward. + before_snapshot = dict(state.yes) + result = mgr.apply_delta(make_delta(price="0.50", delta="0", side="yes")) + assert result is not None + assert state.yes == before_snapshot + + # Cache MUST be preserved across the zero-delta and identity-stable + # on subsequent reads. + assert state._cached is cached, "zero delta should not invalidate cache" + assert result is cached + again = mgr.get("T") + assert again is cached + + # A non-zero delta still invalidates (control case). + mgr.apply_delta(make_delta(price="0.50", delta="5", side="yes")) + fresh = mgr.get("T") + assert fresh is not None + assert fresh is not cached + assert fresh.yes[0].quantity == Decimal("105") + + def test_issue_347_zero_delta_at_missing_level_preserves_cache(self) -> None: + """A non-positive delta at a missing price level is also a no-op.""" + mgr = OrderbookManager() + mgr.apply_snapshot(make_snapshot(yes=[["0.50", "100"]])) + state = mgr._books["T"] + first = mgr.get("T") + assert first is not None + cached = state._cached + assert cached is first + + # Delta of 0 at a price we don't track: previously already a no-op, + # codify that the cache continues to ride through. + mgr.apply_delta(make_delta(price="0.99", delta="0", side="yes")) + assert state._cached is cached + again = mgr.get("T") + assert again is cached From d215125ee778cb8f4a45a3f82c3a54feac1f2aa6 Mon Sep 17 00:00:00 2001 From: Jeff West Date: Fri, 22 May 2026 08:34:25 -0500 Subject: [PATCH 2/2] polish(tests): drop unused monkeypatch fixture; hoist test imports --- tests/ws/test_orderbook.py | 73 +++++++++++++++++--------------------- 1 file changed, 32 insertions(+), 41 deletions(-) diff --git a/tests/ws/test_orderbook.py b/tests/ws/test_orderbook.py index b87ebef..2f1c11d 100644 --- a/tests/ws/test_orderbook.py +++ b/tests/ws/test_orderbook.py @@ -4,6 +4,10 @@ from decimal import Decimal +import pytest + +from kalshi.models.markets import Orderbook, OrderbookLevel +from kalshi.ws import orderbook as ob_mod from kalshi.ws.models.orderbook_delta import ( OrderbookDeltaMessage, OrderbookDeltaPayload, @@ -510,7 +514,7 @@ class TestWave2CorrectnessRegressions: """Round-3 W2 perf-correctness regressions for #327, #344, #347.""" def test_issue_327_to_orderbook_no_revalidation( - self, monkeypatch: object + self, monkeypatch: pytest.MonkeyPatch ) -> None: """``_BookState.to_orderbook`` materializes via ``model_construct``. @@ -524,11 +528,6 @@ def test_issue_327_to_orderbook_no_revalidation( ``model_construct`` bypasses) and assert zero calls during a ``apply_snapshot`` + ``apply_delta`` cycle. """ - import pytest - - from kalshi.models.markets import Orderbook, OrderbookLevel - from kalshi.ws import orderbook as ob_mod - calls = {"level": 0, "ob": 0} orig_level_init = OrderbookLevel.__init__ orig_ob_init = Orderbook.__init__ @@ -541,25 +540,22 @@ def spy_ob_init(self: object, **kw: object) -> None: calls["ob"] += 1 orig_ob_init(self, **kw) # type: ignore[arg-type] - mp = pytest.MonkeyPatch() - mp.setattr(OrderbookLevel, "__init__", spy_level_init) - mp.setattr(Orderbook, "__init__", spy_ob_init) - try: - mgr = ob_mod.OrderbookManager() - book = mgr.apply_snapshot( - make_snapshot( - yes=[["0.50", "100"], ["0.55", "200"], ["0.60", "300"]], - no=[["0.40", "150"], ["0.35", "75"]], - ) - ) - assert len(book.yes) == 3 and len(book.no) == 2 - # Force re-materialization by invalidating and reading again. - delta_book = mgr.apply_delta( - make_delta(price="0.50", delta="50", side="yes") + monkeypatch.setattr(OrderbookLevel, "__init__", spy_level_init) + monkeypatch.setattr(Orderbook, "__init__", spy_ob_init) + + mgr = ob_mod.OrderbookManager() + book = mgr.apply_snapshot( + make_snapshot( + yes=[["0.50", "100"], ["0.55", "200"], ["0.60", "300"]], + no=[["0.40", "150"], ["0.35", "75"]], ) - assert delta_book is not None - finally: - mp.undo() + ) + assert len(book.yes) == 3 and len(book.no) == 2 + # Force re-materialization by invalidating and reading again. + delta_book = mgr.apply_delta( + make_delta(price="0.50", delta="50", side="yes") + ) + assert delta_book is not None assert ( calls["level"] == 0 @@ -568,7 +564,9 @@ def spy_ob_init(self: object, **kw: object) -> None: calls["ob"] == 0 ), f"Orderbook validating __init__ ran {calls['ob']} times" - def test_issue_344_apply_snapshot_public_single_copy(self) -> None: + def test_issue_344_apply_snapshot_public_single_copy( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: """Public ``apply_snapshot`` assigns each side dict exactly once. Pre-fix the public path identity-adopted ``msg.msg.yes`` / ``.no`` @@ -577,10 +575,6 @@ def test_issue_344_apply_snapshot_public_single_copy(self) -> None: per side for one useful copy. We spy on ``_BookState.__setattr__`` and assert exactly one ``yes`` and one ``no`` assignment. """ - import pytest - - from kalshi.ws import orderbook as ob_mod - side_assignments: list[tuple[str, int]] = [] orig_setattr = ob_mod._BookState.__setattr__ @@ -589,18 +583,15 @@ def spy_setattr(self: object, name: str, value: object) -> None: side_assignments.append((name, id(value))) orig_setattr(self, name, value) # type: ignore[arg-type] - mp = pytest.MonkeyPatch() - mp.setattr(ob_mod._BookState, "__setattr__", spy_setattr) - try: - msg = make_snapshot( - yes=[["0.50", "100"], ["0.55", "200"]], - no=[["0.45", "150"]], - ) - mgr = ob_mod.OrderbookManager() - book = mgr.apply_snapshot(msg) - assert book.ticker == "T" - finally: - mp.undo() + monkeypatch.setattr(ob_mod._BookState, "__setattr__", spy_setattr) + + msg = make_snapshot( + yes=[["0.50", "100"], ["0.55", "200"]], + no=[["0.45", "150"]], + ) + mgr = ob_mod.OrderbookManager() + book = mgr.apply_snapshot(msg) + assert book.ticker == "T" yes_assigns = [a for a in side_assignments if a[0] == "yes"] no_assigns = [a for a in side_assignments if a[0] == "no"]