perf(ws): orderbook materialization via model_construct + apply_snapshot single-copy + zero-delta cache preservation

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:

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,
@@ -505,3 +509,156 @@ 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: pytest.MonkeyPatch
+    ) -> 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.
+        """
+        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]
+
+        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 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
+        ), 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, monkeypatch: pytest.MonkeyPatch
+    ) -> 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.
+        """
+        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]
+
+        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"]
+        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