diff --git a/kalshi/ws/orderbook.py b/kalshi/ws/orderbook.py index 969501d..16b5638 100644 --- a/kalshi/ws/orderbook.py +++ b/kalshi/ws/orderbook.py @@ -2,6 +2,7 @@ from __future__ import annotations import logging +from dataclasses import dataclass, field from decimal import Decimal from kalshi.models.markets import Orderbook, OrderbookLevel @@ -13,6 +14,42 @@ logger = logging.getLogger("kalshi.ws") +@dataclass +class _BookState: + """Internal mutable state for one ticker's orderbook. + + Levels are stored price-indexed (``dict[Decimal, Decimal]``) so delta + application is O(1) regardless of book depth. The sorted + ``list[OrderbookLevel]`` exposed via :class:`Orderbook` is materialized + lazily in :meth:`to_orderbook` (O(n log n), but only when a snapshot is + actually emitted to a consumer). + + Kept separate from the public :class:`Orderbook` model so the manager can + mutate freely without leaking changes to previously-handed-out snapshots. + """ + + ticker: str + yes: dict[Decimal, Decimal] = field(default_factory=dict) + no: dict[Decimal, Decimal] = field(default_factory=dict) + + def to_orderbook(self) -> Orderbook: + """Build a fresh ``Orderbook`` snapshot from the current state. + + Levels are emitted price-ascending to match the historical wire-order + contract that the previous list-backed implementation maintained via + ``list.sort`` after every insert. + """ + yes_levels = [ + OrderbookLevel(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()) + ] + return Orderbook(ticker=self.ticker, yes=yes_levels, no=no_levels) + + class OrderbookManager: """Maintains local orderbook state from WebSocket stream. @@ -22,6 +59,10 @@ class OrderbookManager: for quantities; both parse directly into ``Decimal`` without any cents-to-dollars conversion. + Each call to :meth:`apply_snapshot`, :meth:`apply_delta`, or :meth:`get` + returns a fresh :class:`Orderbook` instance. Consumers may safely retain + references; subsequent updates will not mutate previously-returned books. + Usage: mgr = OrderbookManager() book = mgr.apply_snapshot(snapshot_msg) # Initialize @@ -30,38 +71,36 @@ class OrderbookManager: """ def __init__(self) -> None: - self._books: dict[str, Orderbook] = {} + self._books: dict[str, _BookState] = {} def apply_snapshot(self, msg: OrderbookSnapshotMessage) -> Orderbook: """Initialize (or reset) a book from a full snapshot.""" ticker = msg.msg.market_ticker - yes_levels = [ - OrderbookLevel(price=Decimal(p), quantity=Decimal(q)) - for p, q in msg.msg.yes - ] - no_levels = [ - OrderbookLevel(price=Decimal(p), quantity=Decimal(q)) - for p, q in msg.msg.no - ] - book = Orderbook(ticker=ticker, yes=yes_levels, no=no_levels) - self._books[ticker] = book + yes_levels = {Decimal(p): Decimal(q) for p, q in msg.msg.yes} + no_levels = {Decimal(p): Decimal(q) for p, q in msg.msg.no} + state = _BookState(ticker=ticker, yes=yes_levels, no=no_levels) + self._books[ticker] = state logger.debug( "Orderbook snapshot: %s (%d yes, %d no levels)", ticker, len(yes_levels), len(no_levels), ) - return book + return state.to_orderbook() def apply_delta(self, msg: OrderbookDeltaMessage) -> Orderbook | None: """Apply an incremental delta to an existing book. + O(1) on the update itself (price-indexed dict lookup); the returned + Orderbook is materialized at O(n log n) only because consumers expect + a sorted level list. + Returns the updated Orderbook, or None if no book exists for this ticker (delta arrived before snapshot -- should not happen in normal flow). """ ticker = msg.msg.market_ticker - book = self._books.get(ticker) - if book is None: + state = self._books.get(ticker) + if state is None: logger.warning("Delta for unknown ticker %s (no snapshot yet)", ticker) return None @@ -69,35 +108,30 @@ def apply_delta(self, msg: OrderbookDeltaMessage) -> Orderbook | None: delta = msg.msg.delta # Decimal via FixedPointCount side = msg.msg.side - levels = book.yes if side == "yes" else book.no - - # Find existing level at this price - existing_idx = -1 - for i, level in enumerate(levels): - if level.price == price: - existing_idx = i - break + levels = state.yes if side == "yes" else state.no + existing_qty = levels.get(price) - if existing_idx >= 0: - existing = levels[existing_idx] - new_qty = existing.quantity + delta + if existing_qty is not None: + new_qty = existing_qty + delta if new_qty <= 0: - # Remove the level - levels.pop(existing_idx) + del levels[price] else: - levels[existing_idx] = OrderbookLevel(price=price, quantity=new_qty) - else: - if delta > 0: - # Add new level - levels.append(OrderbookLevel(price=price, quantity=delta)) - # Keep sorted by price - levels.sort(key=lambda lv: lv.price) + levels[price] = new_qty + elif delta > 0: + levels[price] = delta - return book + return state.to_orderbook() def get(self, ticker: str) -> Orderbook | None: - """Get current book state (non-blocking).""" - return self._books.get(ticker) + """Get current book state (non-blocking). + + Returns a fresh :class:`Orderbook` snapshot; the caller is free to + retain it without seeing future mutations leak in. + """ + state = self._books.get(ticker) + if state is None: + return None + return state.to_orderbook() def remove(self, ticker: str) -> None: """Remove a book (e.g., on unsubscribe).""" diff --git a/tests/ws/test_orderbook.py b/tests/ws/test_orderbook.py index 33cb23a..42e40ea 100644 --- a/tests/ws/test_orderbook.py +++ b/tests/ws/test_orderbook.py @@ -159,6 +159,101 @@ def test_negative_delta_partial(self) -> None: assert book is not None assert book.yes[0].quantity == Decimal("70") # 100 - 30 + def test_returned_book_is_snapshot_not_live_view(self) -> None: + """Regression for #85: a book handed to a consumer must not mutate + when subsequent deltas are applied. Pin both the identity invariant + (each call returns a distinct ``Orderbook``) and the value invariant + (the previously-returned book's state is preserved verbatim). + """ + mgr = OrderbookManager() + snap_book = mgr.apply_snapshot(make_snapshot(yes=[["0.50", "100"]])) + + # Apply a delta that would have mutated the level in-place under the + # old implementation (quantity 100 -> 150). + delta_book = mgr.apply_delta(make_delta(price="0.50", delta="50", side="yes")) + assert delta_book is not None + + # Distinct instances. + assert snap_book is not delta_book + # The first-handed-out book is frozen at its emission-time state. + assert snap_book.yes[0].quantity == Decimal("100") + assert delta_book.yes[0].quantity == Decimal("150") + + # And ``get()`` is also a snapshot, not a live view. + get_book = mgr.get("T") + assert get_book is not None + assert get_book is not delta_book + mgr.apply_delta(make_delta(price="0.50", delta="25", side="yes")) + assert get_book.yes[0].quantity == Decimal("150") # unchanged + + def test_apply_delta_is_constant_time_in_book_depth(self) -> None: + """Regression for #87: apply_delta must touch at most a constant + number of levels per delta, not scan the whole book. + + We swap the internal price->qty mapping for a counting dict that + records every key access (``get`` / ``__getitem__`` / + ``__setitem__`` / ``__delitem__`` / ``__contains__``). The old + list-scan implementation iterated the level container once per + level (O(n) accesses); the new dict-backed implementation does + O(1) accesses regardless of book depth. + + This is an *invariant* test (access count grows like O(1), not + O(n)), not a wall-clock budget -- stable on slow CI. + """ + from kalshi.ws.orderbook import _BookState + + class CountingDict(dict[Decimal, Decimal]): + def __init__(self, *a: object, **kw: object) -> None: + super().__init__(*a, **kw) # type: ignore[arg-type] + self.accesses = 0 + + def get(self, key: object, default: object = None) -> object: # type: ignore[override] + self.accesses += 1 + return super().get(key, default) # type: ignore[arg-type] + + def __getitem__(self, key: Decimal) -> Decimal: + self.accesses += 1 + return super().__getitem__(key) + + def __setitem__(self, key: Decimal, value: Decimal) -> None: + self.accesses += 1 + super().__setitem__(key, value) + + def __delitem__(self, key: Decimal) -> None: + self.accesses += 1 + super().__delitem__(key) + + def __contains__(self, key: object) -> bool: + self.accesses += 1 + return super().__contains__(key) + + # Two probe sizes -- a small book and a large book. If the + # implementation is O(n), the second number is ~50x the first. + # If O(1), they're identical. + def measure(n: int) -> int: + mgr = OrderbookManager() + mgr.apply_snapshot( + make_snapshot(yes=[[f"0.{i:04d}", "100"] for i in range(1, n + 1)]) + ) + counting = CountingDict(mgr._books["T"].yes) + mgr._books["T"] = _BookState( + ticker="T", yes=counting, no=mgr._books["T"].no + ) + mgr.apply_delta( + make_delta(price=f"0.{n // 2:04d}", delta="1", side="yes") + ) + return counting.accesses + + small = measure(10) + large = measure(500) + + # Allow generous slack but reject anything that scales with n. + # O(1) -> equal counts; O(n) -> large is ~50x small. + assert large <= small + 2, ( + f"apply_delta access count grew from {small} (n=10) to " + f"{large} (n=500) -- expected O(1), looks O(n)" + ) + def test_fractional_delta_matches_wire_format(self) -> None: """Live capture on demo shows delta_fp arrives as ``"1.00"`` (with trailing .00), not bare ``"1"``. Lock in that Decimal arithmetic