Check single-subscription channels before queueing, not during replay

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)

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