From 226b7eb78043f002e7183ce096031b002632455b Mon Sep 17 00:00:00 2001 From: Mohith26 <109003762+Mohith26@users.noreply.github.com> Date: Sun, 8 Nov 2026 15:46:38 -0500 Subject: [PATCH] WsClient: keep local order books sorted and verify begin_nonce continuity - asks/bids stay sorted (asks ascending, bids descending) when new price levels arrive; previously new levels were appended to the end of the lists - level updates use bisect over a parallel price index instead of scanning and copying the list per delta (4-18x faster in test/bench_ws_order_book.py) - nonce/begin_nonce/offset/last_updated_at on the local book track the last applied update instead of staying at the snapshot values - each update is checked with the documented continuity rule (begin_nonce == previous nonce); on a gap the book is dropped, the market is resubscribed for a fresh snapshot and on_order_book_gap is called. verify_order_book_continuity=False turns the check off --- examples/README.md | 11 ++ lighter/ws_client.py | 122 +++++++++++++++-- test/bench_ws_order_book.py | 62 +++++++++ test/test_ws_client_order_book.py | 213 ++++++++++++++++++++++++++++++ 4 files changed, 400 insertions(+), 8 deletions(-) create mode 100644 test/bench_ws_order_book.py create mode 100644 test/test_ws_client_order_book.py diff --git a/examples/README.md b/examples/README.md index 26a400a..5241c1d 100644 --- a/examples/README.md +++ b/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: diff --git a/lighter/ws_client.py b/lighter/ws_client.py index 0638875..6ae934e 100644 --- a/lighter/ws_client.py +++ b/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: diff --git a/test/bench_ws_order_book.py b/test/bench_ws_order_book.py new file mode 100644 index 0000000..899a413 --- /dev/null +++ b/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") diff --git a/test/test_ws_client_order_book.py b/test/test_ws_client_order_book.py new file mode 100644 index 0000000..fe88414 --- /dev/null +++ b/test/test_ws_client_order_book.py @@ -0,0 +1,213 @@ +import asyncio +import json +import random +from decimal import Decimal + +from lighter.ws_client import WsClient + + +class FakeWs: + def __init__(self): + self.sent = [] + + def send(self, payload): + self.sent.append(json.loads(payload)) + + +class FakeAsyncWs: + def __init__(self): + self.sent = [] + + async def send(self, payload): + self.sent.append(json.loads(payload)) + + +def make_client(**kwargs): + kwargs.setdefault("on_order_book_update", None) + return WsClient(ws_url="wss://example.invalid/stream", order_book_ids=[1], **kwargs) + + +def snapshot(asks, bids, nonce=100, offset=10): + return { + "type": "subscribed/order_book", + "channel": "order_book:1", + "offset": offset, + "order_book": { + "code": 0, + "asks": [{"price": p, "size": s} for p, s in asks], + "bids": [{"price": p, "size": s} for p, s in bids], + "offset": offset, + "nonce": nonce, + "begin_nonce": 0, + "last_updated_at": 1, + }, + } + + +def update(asks, bids, begin_nonce, nonce, offset=11): + return { + "type": "update/order_book", + "channel": "order_book:1", + "offset": offset, + "order_book": { + "code": 0, + "asks": [{"price": p, "size": s} for p, s in asks], + "bids": [{"price": p, "size": s} for p, s in bids], + "offset": offset, + "nonce": nonce, + "begin_nonce": begin_nonce, + "last_updated_at": 2, + }, + } + + +def levels(side): + return [(o["price"], o["size"]) for o in side] + + +def test_snapshot_is_sorted_and_zero_levels_dropped(): + client = make_client() + client.on_message(FakeWs(), snapshot( + asks=[("101.0", "1"), ("100.5", "2"), ("102.0", "0")], + bids=[("99.0", "1"), ("99.5", "3"), ("98.0", "0.000")], + )) + book = client.order_book_states["1"] + assert levels(book["asks"]) == [("100.5", "2"), ("101.0", "1")] + assert levels(book["bids"]) == [("99.5", "3"), ("99.0", "1")] + + +def test_new_levels_inside_the_book_keep_best_prices_first(): + client = make_client() + ws = FakeWs() + client.on_message(ws, snapshot(asks=[("100.0", "1"), ("102.0", "1")], bids=[("98.0", "1"), ("96.0", "1")])) + # Before this change these levels were appended at the end of the lists. + client.on_message(ws, update(asks=[("101.0", "5"), ("99.5", "2")], bids=[("97.0", "4"), ("99.0", "3")], + begin_nonce=100, nonce=105)) + book = client.order_book_states["1"] + assert levels(book["asks"]) == [("99.5", "2"), ("100.0", "1"), ("101.0", "5"), ("102.0", "1")] + assert levels(book["bids"]) == [("99.0", "3"), ("98.0", "1"), ("97.0", "4"), ("96.0", "1")] + assert book["asks"][0]["price"] == "99.5" and book["bids"][0]["price"] == "99.0" + + +def test_size_changes_and_removals(): + client = make_client() + ws = FakeWs() + client.on_message(ws, snapshot(asks=[("100.0", "1"), ("101.0", "1")], bids=[("99.0", "1")])) + client.on_message(ws, update(asks=[("100.0", "0.00000"), ("101.0", "7"), ("105.0", "0")], + bids=[("99.00", "2")], begin_nonce=100, nonce=101)) + book = client.order_book_states["1"] + assert levels(book["asks"]) == [("101.0", "7")] + # "99.00" and "99.0" are the same price level + assert levels(book["bids"]) == [("99.0", "2")] + + +def test_book_metadata_tracks_the_latest_update(): + client = make_client() + ws = FakeWs() + client.on_message(ws, snapshot(asks=[], bids=[], nonce=100, offset=10)) + client.on_message(ws, update(asks=[("1", "1")], bids=[], begin_nonce=100, nonce=103, offset=14)) + book = client.order_book_states["1"] + assert (book["nonce"], book["begin_nonce"], book["offset"], book["last_updated_at"]) == (103, 100, 14, 2) + + +def test_gap_resubscribes_and_drops_deltas_until_fresh_snapshot(): + gaps = [] + updates = [] + client = make_client(on_order_book_gap=lambda m, e, g: gaps.append((m, e, g)), + on_order_book_update=lambda m, ob: updates.append(m)) + ws = FakeWs() + client.on_message(ws, snapshot(asks=[("100.0", "1")], bids=[("99.0", "1")], nonce=100)) + client.on_message(ws, update(asks=[("100.0", "2")], bids=[], begin_nonce=100, nonce=110)) + # nonces 110..120 were missed + client.on_message(ws, update(asks=[("100.0", "9")], bids=[], begin_nonce=120, nonce=125)) + assert gaps == [("1", 110, 120)] + assert ws.sent == [{"type": "unsubscribe", "channel": "order_book/1"}, + {"type": "subscribe", "channel": "order_book/1"}] + assert levels(client.order_book_states["1"]["asks"]) == [("100.0", "2")] + # deltas that arrive before the new snapshot are ignored, without another resubscribe + client.on_message(ws, update(asks=[("100.0", "3")], bids=[], begin_nonce=125, nonce=126)) + assert len(ws.sent) == 2 and len(gaps) == 1 + client.on_message(ws, snapshot(asks=[("100.5", "4")], bids=[("99.0", "1")], nonce=130)) + client.on_message(ws, update(asks=[("100.5", "5")], bids=[], begin_nonce=130, nonce=131)) + assert levels(client.order_book_states["1"]["asks"]) == [("100.5", "5")] + assert updates == ["1", "1", "1", "1"] # snapshot, update, snapshot, update + + +def test_resync_unsubscribe_ack_is_not_reported_as_unhandled(): + unhandled = [] + client = make_client(on_unhandled_message=unhandled.append) + ws = FakeWs() + client.on_message(ws, snapshot(asks=[("100.0", "1")], bids=[], nonce=100)) + client.on_message(ws, update(asks=[], bids=[], begin_nonce=150, nonce=151)) + # shape observed from the testnet server after an unsubscribe + client.on_message(ws, {"type": "unsubscribed", "channel": "order_book:1"}) + assert unhandled == [] + client.on_message(ws, snapshot(asks=[("100.0", "1")], bids=[], nonce=160)) + client.on_message(ws, {"type": "unsubscribed", "channel": "order_book:1"}) + assert unhandled == [{"type": "unsubscribed", "channel": "order_book:1"}] + + +def test_continuity_check_can_be_disabled(): + client = make_client(verify_order_book_continuity=False) + ws = FakeWs() + client.on_message(ws, snapshot(asks=[("100.0", "1")], bids=[], nonce=100)) + client.on_message(ws, update(asks=[("100.0", "2")], bids=[], begin_nonce=999, nonce=1000)) + assert levels(client.order_book_states["1"]["asks"]) == [("100.0", "2")] + assert ws.sent == [] + + +def test_async_path_awaits_resubscribe(): + client = make_client() + ws = FakeAsyncWs() + + async def run(): + await client.on_message_async(ws, json.dumps(snapshot(asks=[("1", "1")], bids=[], nonce=5))) + await client.on_message_async(ws, json.dumps(update(asks=[("1", "2")], bids=[], begin_nonce=7, nonce=8))) + + asyncio.run(run()) + assert ws.sent == [{"type": "unsubscribe", "channel": "order_book/1"}, + {"type": "subscribe", "channel": "order_book/1"}] + + +def test_state_set_by_caller_is_indexed_on_first_update(): + client = make_client() + client.order_book_states["1"] = {"asks": [{"price": "3", "size": "1"}, {"price": "1", "size": "1"}], + "bids": [], "nonce": 1} + client.update_order_book_state("1", {"asks": [{"price": "2", "size": "1"}], "bids": []}) + assert [o["price"] for o in client.order_book_states["1"]["asks"]] == ["1", "2", "3"] + + +def test_matches_reference_book_over_random_updates(): + rng = random.Random(0) + client = make_client() + ws = FakeWs() + prices = [f"{100 + i * 0.5:.1f}" for i in range(200)] + ref = {"asks": {}, "bids": {}} + client.on_message(ws, snapshot(asks=[], bids=[], nonce=0)) + nonce = 0 + for _ in range(3000): + asks, bids = [], [] + for _ in range(rng.randint(1, 6)): + side = rng.choice(["asks", "bids"]) + price = rng.choice(prices) + size = "0" if rng.random() < 0.3 else str(rng.randint(1, 50)) + (asks if side == "asks" else bids).append((price, size)) + client.on_message(ws, update(asks=asks, bids=bids, begin_nonce=nonce, nonce=nonce + 1)) + nonce += 1 + for side, rows in (("asks", asks), ("bids", bids)): + for price, size in rows: + if Decimal(size) == 0: + ref[side].pop(Decimal(price), None) + else: + ref[side][Decimal(price)] = size + book = client.order_book_states["1"] + assert [(Decimal(p), s) for p, s in levels(book["asks"])] == sorted(ref["asks"].items()) + assert [(Decimal(p), s) for p, s in levels(book["bids"])] == sorted(ref["bids"].items(), reverse=True) + assert ws.sent == [] + + +def test_update_orders_is_unchanged_for_existing_callers(): + client = make_client() + existing = [{"price": "1", "size": "1"}] + client.update_orders([{"price": "2", "size": "3"}, {"price": "1", "size": "0"}], existing) + assert existing == [{"price": "2", "size": "3"}]