Check single-subscription channels before queueing, not during replay

[pucedoteth]

What

userEvents and orderUpdates cannot be multiplexed, and subscribe rejects a second one with NotImplementedError. That check only ran on the connected path, so subscribing twice before the socket opens is accepted and queued, then rejected later while on_open replays the queue.

Why it matters

The exception escapes inside the websocket callback, where the caller cannot catch it, and it aborts the replay loop. Every subscription queued behind the duplicate is silently dropped.

Against master:

ws_manager.subscribe({"type": "userEvents"}, cb)             # queued
ws_manager.subscribe({"type": "userEvents"}, cb)             # queued, no error
ws_manager.subscribe({"type": "l2Book", "coin": "ETH"}, cb)  # queued

ws_manager.on_open(None)
# NotImplementedError: Cannot subscribe to userEvents multiple times
#   subscribe frames sent to the server: 1
#   l2Book:eth registered:                False

The ETH order book feed is never subscribed and never reported as missing.

The same two calls after the socket is open raise at the call site:

ws_manager.subscribe({"type": "userEvents"}, cb)
ws_manager.subscribe({"type": "userEvents"}, cb)
# NotImplementedError raised here, catchable, nothing else affected

So identical user code either raises where it is written, or loses an unrelated market-data feed, depending only on whether the socket happened to be open yet. Info(skip_ws=False) subscribes during construction, so the queued path is the common one at startup.

How

Run the check in subscribe for both paths, counting queued entries as well as active ones, so the duplicate is refused where it is requested.

on_open now takes the queue before replaying it. subscribe consults that list, so it has to be drained first, and leaving entries in place would also replay them again on any later on_open.

Behaviour on the connected path is unchanged, and channels that do multiplex still accept several callbacks on the same identifier.

Tests

Added tests/websocket_manager_test.py, covering the duplicate on both paths, the dropped-subscription case, queue replay and clearing, and multiplexing. The stub socket means nothing connects.

Against the unmodified websocket_manager.py:

FAILED test_duplicate_single_subscription_raises_while_queued
FAILED test_rejected_duplicate_does_not_drop_later_subscriptions
FAILED test_on_open_replays_and_clears_the_queue
3 failed, 2 passed

The two that pass either way are the connected-path duplicate and the multiplexing case, so the change is scoped to the broken path. With the fix the full suite is green: 41 passed.

[koriyoshi2041]

There is still a race between on_open and a caller invoking subscribe. on_open sets ws_ready = True and drains self.queued_subscriptions before replaying the local copy. A caller can then register userEvents in active_subscriptions before the websocket thread replays the queued userEvents; the replay raises NotImplementedError and the later queued subscriptions are lost again.

I reproduced this deterministically by replacing the queued list with an iterator that pauses after the swap and before its first yield, then subscribing to userEvents from the caller thread. The opener captured NotImplementedError('Cannot subscribe to userEvents multiple times'), and the queued l2Book was never registered. The current focused tests pass 5/5 and the repository tests pass 41/41, but none schedules a subscribe in that transition window.

Could the ws_ready transition, queue handoff, duplicate check, and active registration be protected by the same lock, or could replay register the handed-off entries without re-entering the public duplicate check? The latter still needs synchronization with concurrent public subscriptions.

[pucedoteth]

You're right, and thanks for the precise repro instructions. Reproduced and fixed in 91eab45.

I first tried to reproduce by pausing inside the first replayed subscribe, which does not show it: by then userEvents is already in active_subscriptions, so the racing caller correctly raises and the replay survives. Pausing where you said, in the iteration itself, reproduces it exactly:

queued: ['userEvents', 'l2Book']
caller subscribe: returned normally (registered into active_subscriptions)
opener raised: NotImplementedError('Cannot subscribe to userEvents multiple times')
active: {'userEvents': 1}
l2Book registered: False

on_open dies and the queued l2Book is gone. Same silent loss this PR set out to fix, through a narrower window, and it predates the PR: the old check was len(self.active_subscriptions[identifier]) != 0, which loses the same race.

The fix

I took your first suggestion, one lock over the whole transition. self.lock is a threading.RLock; on_open holds it across ws_ready = True, the handoff and the replay, and subscribe re-enters it from that same thread.

def on_open(self, _ws):
    with self.lock:
        self.ws_ready = True
        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)

A racing caller now blocks until the replay finishes and then sees the real duplicate, which is the correct outcome for that caller — it is subscribing to userEvents twice — instead of corrupting the replay for everyone else.

I chose this over "replay without re-entering the duplicate check" because, as you noted, that variant still needs synchronisation, and it would also let a genuine double-queue slip through into active_subscriptions.

unsubscribe takes the lock as well. It read-modify-writes active_subscriptions[identifier], so a concurrent subscribe append could be lost. on_message deliberately does not take it: it only reads, and locking there would serialise callbacks against subscribes and deadlock any callback that subscribes from another thread. Say the word if you'd rather I cover that too.

self.ws.send now happens under the lock. It's a small framed write and serialising subscribe frames is desirable anyway, but flagging it as a deliberate choice rather than an oversight.

Tests

test_subscribe_racing_on_open_does_not_drop_the_replay pins your scenario. It parks the websocket thread in the window with a queue whose iteration blocks, starts the caller, then releases, and asserts all three things: on_open did not raise, the caller did get NotImplementedError, and l2Book is still registered.

Removing only the lock and keeping the tests fails exactly that one test:

FAILED tests/websocket_manager_test.py::test_subscribe_racing_on_open_does_not_drop_the_replay
1 failed, 5 passed

With it, 6 passed in that file and 42 passed across tests/. black --check and mypy are clean.

One note on the harness: arming the pause is explicit, because subscribe also walks the queue for its duplicate check and would otherwise trip it during setup.

[koriyoshi2041]

Confirmed at 91eab45d: the shared RLock closes the handoff/replay window, and the new deterministic regression preserves the queued l2Book while making the racing duplicate caller fail. I also reran the focused websocket tests (6/6), full repository suite (42/42), and diff hygiene locally. The original concurrency finding is resolved.