diff --git a/hyperliquid/websocket_manager.py b/hyperliquid/websocket_manager.py index 4c73a688..e6d9ed4a 100644 --- a/hyperliquid/websocket_manager.py +++ b/hyperliquid/websocket_manager.py @@ -82,6 +82,9 @@ def __init__(self, base_url): self.queued_subscriptions: List[Tuple[Subscription, ActiveSubscription]] = [] self.active_subscriptions: Dict[str, List[ActiveSubscription]] = defaultdict(list) ws_url = "ws" + base_url[len("http") :] + "/ws" + # Guards ws_ready, queued_subscriptions and active_subscriptions. Reentrant because + # on_open replays the queue through subscribe() while already holding it. + self.lock = threading.RLock() self.ws = websocket.WebSocketApp(ws_url, on_message=self.on_message, on_open=self.on_open) self.ping_sender = threading.Thread(target=self.send_ping) self.stop_event = threading.Event() @@ -126,37 +129,52 @@ def on_message(self, _ws, message): def on_open(self, _ws): logging.debug("on_open") - self.ws_ready = True - for subscription, active_subscription in self.queued_subscriptions: - self.subscribe(subscription, active_subscription.callback, active_subscription.subscription_id) + # The whole handoff is one critical section. A caller reaching subscribe() between the + # swap and the replay would otherwise see an empty queue and an empty active map, claim + # a single-subscription channel, and make the replay raise -- which aborts it and drops + # every queued subscription behind the one that raised. + with self.lock: + self.ws_ready = True + # Drain the queue before replaying it: subscribe() consults it for the + # single-subscription channels, and leaving entries behind would also + # replay them again on a later on_open. + queued_subscriptions, self.queued_subscriptions = self.queued_subscriptions, [] + for subscription, active_subscription in queued_subscriptions: + self.subscribe(subscription, active_subscription.callback, active_subscription.subscription_id) def subscribe( self, subscription: Subscription, callback: Callable[[Any], None], subscription_id: Optional[int] = None ) -> int: - if subscription_id is None: - self.subscription_id_counter += 1 - subscription_id = self.subscription_id_counter - if not self.ws_ready: - logging.debug("enqueueing subscription") - self.queued_subscriptions.append((subscription, ActiveSubscription(callback, subscription_id))) - else: - logging.debug("subscribing") + with self.lock: + if subscription_id is None: + self.subscription_id_counter += 1 + subscription_id = self.subscription_id_counter identifier = subscription_to_identifier(subscription) if identifier == "userEvents" or identifier == "orderUpdates": # TODO: ideally the userEvent and orderUpdates messages would include the user so that we can multiplex - if len(self.active_subscriptions[identifier]) != 0: + # Queued subscriptions count too, otherwise the duplicate is only caught while on_open replays the + # queue, where it aborts the replay and silently drops every subscription behind it. + already_queued = any(subscription_to_identifier(s) == identifier for s, _ in self.queued_subscriptions) + if len(self.active_subscriptions[identifier]) != 0 or already_queued: raise NotImplementedError(f"Cannot subscribe to {identifier} multiple times") - self.active_subscriptions[identifier].append(ActiveSubscription(callback, subscription_id)) - self.ws.send(json.dumps({"method": "subscribe", "subscription": subscription})) - return subscription_id + if not self.ws_ready: + logging.debug("enqueueing subscription") + self.queued_subscriptions.append((subscription, ActiveSubscription(callback, subscription_id))) + else: + logging.debug("subscribing") + self.active_subscriptions[identifier].append(ActiveSubscription(callback, subscription_id)) + self.ws.send(json.dumps({"method": "subscribe", "subscription": subscription})) + return subscription_id def unsubscribe(self, subscription: Subscription, subscription_id: int) -> bool: if not self.ws_ready: raise NotImplementedError("Can't unsubscribe before websocket connected") - identifier = subscription_to_identifier(subscription) - active_subscriptions = self.active_subscriptions[identifier] - new_active_subscriptions = [x for x in active_subscriptions if x.subscription_id != subscription_id] - if len(new_active_subscriptions) == 0: - self.ws.send(json.dumps({"method": "unsubscribe", "subscription": subscription})) - self.active_subscriptions[identifier] = new_active_subscriptions - return len(active_subscriptions) != len(new_active_subscriptions) + with self.lock: + identifier = subscription_to_identifier(subscription) + active_subscriptions = self.active_subscriptions[identifier] + new_active_subscriptions = [x for x in active_subscriptions if x.subscription_id != subscription_id] + if len(new_active_subscriptions) == 0: + self.ws.send(json.dumps({"method": "unsubscribe", "subscription": subscription})) + # Read-modify-write: without the lock a concurrent subscribe's append is lost here. + self.active_subscriptions[identifier] = new_active_subscriptions + return len(active_subscriptions) != len(new_active_subscriptions) diff --git a/tests/websocket_manager_test.py b/tests/websocket_manager_test.py new file mode 100644 index 00000000..5cc7121a --- /dev/null +++ b/tests/websocket_manager_test.py @@ -0,0 +1,151 @@ +import json +import threading +import time +from collections import defaultdict +from types import SimpleNamespace + +import pytest + +from hyperliquid.websocket_manager import WebsocketManager + + +def make_manager(ws_ready: bool): + """A WebsocketManager with a stub socket, so no connection is opened.""" + ws_manager = WebsocketManager.__new__(WebsocketManager) + ws_manager.subscription_id_counter = 0 + ws_manager.ws_ready = ws_ready + ws_manager.queued_subscriptions = [] + ws_manager.active_subscriptions = defaultdict(list) + ws_manager.lock = threading.RLock() + sent = [] + ws_manager.ws = SimpleNamespace(send=sent.append) + return ws_manager, sent + + +def callback(_msg): + pass + + +def test_duplicate_single_subscription_raises_while_queued(): + # userEvents and orderUpdates cannot be multiplexed. Before this was checked on + # the queued path the duplicate was only caught later, inside on_open. + ws_manager, _ = make_manager(ws_ready=False) + ws_manager.subscribe({"type": "userEvents"}, callback) + + with pytest.raises(NotImplementedError): + ws_manager.subscribe({"type": "userEvents"}, callback) + + +def test_duplicate_single_subscription_raises_when_connected(): + ws_manager, _ = make_manager(ws_ready=True) + ws_manager.subscribe({"type": "orderUpdates"}, callback) + + with pytest.raises(NotImplementedError): + ws_manager.subscribe({"type": "orderUpdates"}, callback) + + +def test_rejected_duplicate_does_not_drop_later_subscriptions(): + # The duplicate used to surface inside on_open, which aborted the replay and + # silently dropped every subscription queued behind it. + ws_manager, sent = make_manager(ws_ready=False) + ws_manager.subscribe({"type": "userEvents"}, callback) + with pytest.raises(NotImplementedError): + ws_manager.subscribe({"type": "userEvents"}, callback) + ws_manager.subscribe({"type": "l2Book", "coin": "ETH"}, callback) + + ws_manager.on_open(None) + + assert len(ws_manager.active_subscriptions["l2Book:eth"]) == 1 + assert [json.loads(msg)["subscription"]["type"] for msg in sent] == ["userEvents", "l2Book"] + + +def test_on_open_replays_and_clears_the_queue(): + ws_manager, sent = make_manager(ws_ready=False) + ws_manager.subscribe({"type": "l2Book", "coin": "BTC"}, callback) + ws_manager.subscribe({"type": "trades", "coin": "ETH"}, callback) + + ws_manager.on_open(None) + + assert len(sent) == 2 + assert len(ws_manager.active_subscriptions["l2Book:btc"]) == 1 + assert len(ws_manager.active_subscriptions["trades:eth"]) == 1 + assert ws_manager.queued_subscriptions == [] + + +def test_multiplexable_channel_still_accepts_several_callbacks(): + ws_manager, _ = make_manager(ws_ready=True) + ws_manager.subscribe({"type": "l2Book", "coin": "ETH"}, callback) + ws_manager.subscribe({"type": "l2Book", "coin": "ETH"}, callback) + + assert len(ws_manager.active_subscriptions["l2Book:eth"]) == 2 + + +class PausingQueue(list): + """A queue whose first armed iteration blocks until released. + + on_open hands the queue off and then iterates it, so pausing in __iter__ parks the + websocket thread in the window after the handoff and before the first replayed + subscription registers anything. Arming is explicit because subscribe() also walks + the queue for its duplicate check. + """ + + def __init__(self, entered, release): + super().__init__() + self.entered = entered + self.release = release + self.armed = False + + def __iter__(self): + if self.armed: + self.armed = False + self.entered.set() + assert self.release.wait(5), "replay was never released" + return super().__iter__() + + +def test_subscribe_racing_on_open_does_not_drop_the_replay(): + # A caller reaching subscribe() between on_open's queue handoff and the replay saw an + # empty queue and an empty active map, so it claimed userEvents. The replay then raised + # out of on_open, aborting it and dropping every subscription queued behind userEvents. + ws_manager, _ = make_manager(ws_ready=False) + entered, release = threading.Event(), threading.Event() + queue = PausingQueue(entered, release) + ws_manager.queued_subscriptions = queue + + ws_manager.subscribe({"type": "userEvents"}, callback) + ws_manager.subscribe({"type": "l2Book", "coin": "ETH"}, callback) + queue.armed = True + + opener: dict = {} + caller: dict = {} + + def run_on_open(): + try: + ws_manager.on_open(None) + except Exception as exc: # noqa: BLE001 - recorded so the assert can report it + opener["error"] = exc + + def run_caller(): + try: + ws_manager.subscribe({"type": "userEvents"}, callback) + except NotImplementedError as exc: + caller["error"] = exc + + opener_thread = threading.Thread(target=run_on_open) + caller_thread = threading.Thread(target=run_caller) + opener_thread.start() + try: + assert entered.wait(5), "on_open never reached the replay" + caller_thread.start() + # Let the caller reach subscribe. With the lock held across the handoff it parks + # there; without it, it claims userEvents and the replay below dies. + time.sleep(0.1) + finally: + release.set() + opener_thread.join(5) + caller_thread.join(5) + + assert "error" not in opener, f"on_open raised {opener.get('error')!r}" + assert "error" in caller, "the racing caller should still see the duplicate" + assert len(ws_manager.active_subscriptions["userEvents"]) == 1 + assert len(ws_manager.active_subscriptions["l2Book:eth"]) == 1