WsClient: keep local order books sorted and verify begin_nonce continuity

examples/README.md

@@ -122,6 +122,17 @@ More details can be found in the following example:
 - `sub_account_transfer_eth.py`
 - `sub_account_transfer_usdc.py`
 
+## Local order books over websocket
+`WsClient` keeps a local copy of every subscribed order book (`client.order_book_states[market_id]`):
+- `asks` are sorted by ascending price and `bids` by descending price, so `asks[0]` and `bids[0]` are the best levels.
+- `nonce`, `begin_nonce`, `offset` and `last_updated_at` reflect the last applied update.
+- Each update is checked for continuity: its `begin_nonce` must equal the `nonce` of the previous update.
+  If it doesn't, an update was missed, so the client drops the book, resubscribes to get a fresh snapshot and
+  calls `on_order_book_gap(market_id, expected_begin_nonce, received_begin_nonce)` if you pass one.
+  Pass `verify_order_book_continuity=False` to turn the check off.
+
+Example: `websocket/ws.py`
+
 ## Public Pools
 Public pools behave just like subaccounts, except that anyone can join them.  
 You can create / modify a public pool using the SDK. Check out the following example:

lighter/ws_client.py

@@ -1,4 +1,6 @@
 import json
+from bisect import bisect_left
+from decimal import Decimal
 from websockets.sync.client import connect
 from websockets.client import connect as connect_async
 from lighter.configuration import Configuration
@@ -18,6 +20,8 @@ def __init__(
         ws_url=None,
         on_shutdown=None,
         on_unhandled_message=None,
+        on_order_book_gap=None,
+        verify_order_book_continuity=True,
     ):
         if ws_url is not None:
             self.base_url = ws_url.rstrip("/")
@@ -36,11 +40,20 @@ def __init__(
 
         self.order_book_states = {}
         self.account_states = {}
+        # Sorted price keys per market and side, kept parallel to the level lists in
+        # order_book_states (asks ascending, bids descending), so updates are O(log n).
+        self._order_book_keys = {}
+        # Markets whose book broke continuity and are waiting for a fresh snapshot.
+        self._resyncing_markets = set()
+        self._pending_resyncs = []
+        self._async_mode = False
 
         self.on_order_book_update = on_order_book_update
         self.on_account_update = on_account_update
         self.on_shutdown = on_shutdown
         self.on_unhandled_message = on_unhandled_message
+        self.on_order_book_gap = on_order_book_gap
+        self.verify_order_book_continuity = verify_order_book_continuity
 
         self.ws = None
 
@@ -65,9 +78,14 @@ def on_message(self, ws, message):
             ws.send(json.dumps({"type": "pong"}))
         elif message_type == "shutdown":
             self.handle_shutdown(message)
+        elif message_type == "unsubscribed" and self._is_resync_ack(message):
+            pass
         else:
             self.handle_unhandled_message(message)
 
+        if not self._async_mode:
+            self._flush_resyncs(ws)
+
     async def on_message_async(self, ws, message):
         message = json.loads(message)
         message_type = message.get("type")
@@ -78,7 +96,9 @@ async def on_message_async(self, ws, message):
             # Respond to ping with pong
             await ws.send(json.dumps({"type": "pong"}))
         else:
+            self._async_mode = True
             self.on_message(ws, message)
+            await self._flush_resyncs_async(ws)
 
     def handle_connected(self, ws):
         for market_id in self.subscriptions["order_books"]:
@@ -106,23 +126,109 @@ async def handle_connected_async(self, ws):
 
     def handle_subscribed_order_book(self, message):
         market_id = message["channel"].split(":")[1]
-        self.order_book_states[market_id] = message["order_book"]
+        order_book = message["order_book"]
+        # The snapshot replaces any previous book, including one that lost continuity.
+        self._resyncing_markets.discard(market_id)
+        order_book["asks"] = [o for o in order_book.get("asks", []) if Decimal(o["size"]) > 0]
+        order_book["bids"] = [o for o in order_book.get("bids", []) if Decimal(o["size"]) > 0]
+        order_book["asks"].sort(key=lambda o: Decimal(o["price"]))
+        order_book["bids"].sort(key=lambda o: -Decimal(o["price"]))
+        self._order_book_keys[market_id] = {
+            "asks": [Decimal(o["price"]) for o in order_book["asks"]],
+            "bids": [-Decimal(o["price"]) for o in order_book["bids"]],
+        }
+        self.order_book_states[market_id] = order_book
         if self.on_order_book_update:
             self.on_order_book_update(market_id, self.order_book_states[market_id])
 
     def handle_update_order_book(self, message):
         market_id = message["channel"].split(":")[1]
-        self.update_order_book_state(market_id, message["order_book"])
+        if market_id in self._resyncing_markets or market_id not in self.order_book_states:
+            # Waiting for a fresh snapshot; deltas on a broken book would be wrong.
+            return
+        order_book = message["order_book"]
+        if self.verify_order_book_continuity and not self._is_continuous(market_id, order_book):
+            self._handle_order_book_gap(market_id, order_book)
+            return
+        self.update_order_book_state(market_id, order_book)
         if self.on_order_book_update:
             self.on_order_book_update(market_id, self.order_book_states[market_id])
 
-    def update_order_book_state(self, market_id, order_book):
-        self.update_orders(
-            order_book["asks"], self.order_book_states[market_id]["asks"]
-        )
-        self.update_orders(
-            order_book["bids"], self.order_book_states[market_id]["bids"]
+    def _is_continuous(self, market_id, order_book):
+        """Lighter's documented check: an update's begin_nonce equals the previous nonce."""
+        expected = self.order_book_states[market_id].get("nonce")
+        begin = order_book.get("begin_nonce")
+        if expected is None or begin is None:
+            return True
+        return begin == expected
+
+    def _handle_order_book_gap(self, market_id, order_book):
+        expected = self.order_book_states[market_id].get("nonce")
+        got = order_book.get("begin_nonce")
+        logger.warning(
+            "Order book %s lost continuity (expected begin_nonce %s, got %s); resubscribing for a fresh snapshot",
+            market_id, expected, got,
         )
+        self._resyncing_markets.add(market_id)
+        self._pending_resyncs.append(market_id)
+        if self.on_order_book_gap:
+            self.on_order_book_gap(market_id, expected, got)
+
+    def _is_resync_ack(self, message):
+        """True for the server's ack of an unsubscribe we sent while resyncing a book."""
+        channel = message.get("channel", "")
+        return channel.startswith("order_book:") and channel.split(":")[1] in self._resyncing_markets
+
+    def _resync_messages(self, market_id):
+        channel = f"order_book/{market_id}"
+        return [
+            json.dumps({"type": "unsubscribe", "channel": channel}),
+            json.dumps({"type": "subscribe", "channel": channel}),
+        ]
+
+    def _flush_resyncs(self, ws):
+        while self._pending_resyncs:
+            for payload in self._resync_messages(self._pending_resyncs.pop(0)):
+                ws.send(payload)
+
+    async def _flush_resyncs_async(self, ws):
+        while self._pending_resyncs:
+            for payload in self._resync_messages(self._pending_resyncs.pop(0)):
+                await ws.send(payload)
+
+    def update_order_book_state(self, market_id, order_book):
+        state = self.order_book_states[market_id]
+        keys = self._order_book_keys.get(market_id)
+        if keys is None:
+            # State set directly by the caller: index it once.
+            state["asks"].sort(key=lambda o: Decimal(o["price"]))
+            state["bids"].sort(key=lambda o: -Decimal(o["price"]))
+            keys = self._order_book_keys[market_id] = {
+                "asks": [Decimal(o["price"]) for o in state["asks"]],
+                "bids": [-Decimal(o["price"]) for o in state["bids"]],
+            }
+        self._apply_levels(order_book.get("asks", []), state["asks"], keys["asks"], 1)
+        self._apply_levels(order_book.get("bids", []), state["bids"], keys["bids"], -1)
+        for field in ("nonce", "begin_nonce", "offset", "last_updated_at"):
+            if field in order_book:
+                state[field] = order_book[field]
+
+    @staticmethod
+    def _apply_levels(new_orders, levels, keys, sign):
+        """Apply price-level deltas in place, keeping levels sorted (sign=1 asks, -1 bids)."""
+        for new_order in new_orders:
+            key = sign * Decimal(new_order["price"])
+            i = bisect_left(keys, key)
+            exists = i < len(keys) and keys[i] == key
+            if Decimal(new_order["size"]) == 0:
+                if exists:
+                    del keys[i]
+                    del levels[i]
+            elif exists:
+                levels[i]["size"] = new_order["size"]
+            else:
+                keys.insert(i, key)
+                levels.insert(i, new_order)
 
     def update_orders(self, new_orders, existing_orders):
         for new_order in new_orders:

test/bench_ws_order_book.py

@@ -0,0 +1,62 @@
+"""Benchmark WsClient order book updates: indexed sorted levels vs. the previous list scan.
+
+Run: python test/bench_ws_order_book.py
+Not collected by pytest (file name does not start with test_).
+"""
+import random
+import statistics
+import time
+
+from lighter.ws_client import WsClient
+
+
+def make_stream(levels, n_updates, seed=0):
+    rng = random.Random(seed)
+    mid = 100_000.0
+    asks = [{"price": f"{mid + 0.1 * (i + 1):.1f}", "size": "1.00000"} for i in range(levels)]
+    bids = [{"price": f"{mid - 0.1 * (i + 1):.1f}", "size": "1.00000"} for i in range(levels)]
+    width = levels * 1.2
+    stream = []
+    for _ in range(n_updates):
+        upd = {"asks": [], "bids": []}
+        for _ in range(rng.randint(1, 8)):
+            side = rng.choice(["asks", "bids"])
+            off = 0.1 * rng.randint(1, int(width))
+            price = mid + off if side == "asks" else mid - off
+            size = "0.00000" if rng.random() < 0.3 else f"{rng.uniform(0.001, 5):.5f}"
+            upd[side].append({"price": f"{price:.1f}", "size": size})
+        stream.append(upd)
+    return asks, bids, stream
+
+
+def run_new(asks, bids, stream):
+    client = WsClient(ws_url="wss://example.invalid/stream", order_book_ids=[1],
+                      on_order_book_update=None, verify_order_book_continuity=False)
+    client.handle_subscribed_order_book({"channel": "order_book:1", "order_book": {
+        "asks": [dict(o) for o in asks], "bids": [dict(o) for o in bids]}})
+    t = time.perf_counter()
+    for upd in stream:
+        client.update_order_book_state("1", {"asks": [dict(o) for o in upd["asks"]],
+                                             "bids": [dict(o) for o in upd["bids"]]})
+    return time.perf_counter() - t
+
+
+def run_old(asks, bids, stream):
+    client = WsClient(ws_url="wss://example.invalid/stream", order_book_ids=[1], on_order_book_update=None)
+    state = {"asks": [dict(o) for o in asks], "bids": [dict(o) for o in bids]}
+    t = time.perf_counter()
+    for upd in stream:
+        # the previous update_order_book_state body
+        client.update_orders([dict(o) for o in upd["asks"]], state["asks"])
+        client.update_orders([dict(o) for o in upd["bids"]], state["bids"])
+    return time.perf_counter() - t
+
+
+if __name__ == "__main__":
+    print(f"{'levels/side':>11} {'updates':>8} {'old median s':>13} {'new median s':>13} {'speedup':>8}")
+    for levels in (50, 200, 1000):
+        n = 3_000
+        asks, bids, stream = make_stream(levels, n)
+        old = statistics.median(run_old(asks, bids, stream) for _ in range(5))
+        new = statistics.median(run_new(asks, bids, stream) for _ in range(5))
+        print(f"{levels:>11} {n:>8} {old:>13.4f} {new:>13.4f} {old / new:>7.1f}x")