fix: keep WebSocket subscription changes in sync and send them promptly by chipkent · Pull Request #2 · chipkent/massive-client-python

chipkent · 2026-10-23T23:04:41Z

Preview only. This is a draft in the fork showing the exact text of the PR planned for massive-com/client-python. It will be closed, not merged, once the upstream PR is opened.

Fixes #<issue: F1a>
Fixes #<issue: F1b>
Fixes #<issue: F1c>

This PR builds on the PR that adds free-threaded CI jobs. The PR that fixes reconnects builds on
this one.

Summary

WebSocketClient lets callers change subscriptions while connect() runs. This PR fixes four bugs
in how those changes reach the server and adds two methods.

In this description:

client.scheduled_subs is the set of topics the client wants;
client.subs is the set it believes it has sent.

Compatibility

Existing behaviour is unchanged apart from the bug fixes below.

Public API: no public name, signature, default or attribute type changes.
- scheduled_subs and subs are still sets, and schedule_resub is still a bool.
- Code that changes scheduled_subs directly and sets schedule_resub still works.
Existing methods: subscribe(), unsubscribe() and unsubscribe_all() keep their send
order: the subscribe message first, then the unsubscribe message.
Topics: they are stored exactly as given, including those passed to the constructor, and a
malformed topic is logged as before.
Receive path: the code that receives messages is unchanged, and so is its throughput.

Bug fixes

1. Subscription changes made while `connect()` runs can be lost (#<issue: F1a>)

Bug: a change made while connect() was sending an earlier change was marked as sent but never
sent. The server then streamed a different set of topics from client.scheduled_subs, while
client.subs said everything was sent.

Fix: subscribe(), unsubscribe() and unsubscribe_all() now take a lock. connect() copies
the wanted topics under that lock and records as sent only the topics it actually sent.

The methods replace scheduled_subs with a new set instead of changing it in place, so another
thread can iterate it safely. One consequence: a reference to the old set no longer sees later
changes.

2. Subscription changes are sent up to 1 s late (#<issue: F1a>)

Bug: a change waited until the receive loop came round again, after the next incoming message
or up to the 1-second receive timeout.

Fix: each connection now has one task that sends subscription changes. It starts after
authentication by sending all wanted topics, so a reconnect sends them again.

Waking it: a change wakes the task through loop.call_soon_threadsafe(), which never blocks
the caller. The receive loop also wakes it whenever it sees schedule_resub set.
One sender: being the only sender, the task never interleaves messages, and it merges a burst
of changes into one subscribe message and one unsubscribe message.
Safety while receiving: the task sends while the receive loop waits for the next message.
websockets allows this: in its asyncio connection, only two recv() calls at the same time
conflict.
Errors and shutdown: a send error other than a closed connection is raised by connect(),
as before. The task is stopped when the connection ends.

3. Wildcard topics match by string prefix (#<issue: F1b>)

Subscribing to a wildcard such as A.* replaces the other topics on the same channel. The client
matched by string prefix instead of by channel:

subscribe("A.*") also unsubscribed AM.AAPL, and likewise XA.* and XAS.*, or C.*,
CA.* and CAS.*;
calling subscribe("A.*") a second time unsubscribed A.*, because the wildcard matched itself.

The client now compares the channel exactly and skips the wildcard itself, in both subscribe()
and unsubscribe().

As before, only topics already sent are checked, so subscribe("T.AAPL", "T.*") still subscribes
to both. This PR leaves that unchanged to keep existing behaviour.

4. Changes made around a reconnect are lost

Bug: after the connection closed with an error, connect() reset the wanted topics to those
already sent.

A subscribe() or unsubscribe() made just before the close, or while reconnecting, was lost.
An unsubscribed topic could come back.

Fix: the wanted topics are now kept, and the new connection subscribes to all of them.

Calling connect() again: after close() or a cancel, the client also assumed the old
connection's topics were still subscribed and sent nothing. Each connect() call now starts with
an empty client.subs.

The reconnect issue lists this too; this PR is the one that fixes it.

New methods (#<issue: F1c>)

Both methods are safe to call from any thread while connect() runs.

`set_subscriptions(desired)`

Replaces the wanted topics in one step. The client then sends one unsubscribe message for the
removed topics, followed by one subscribe message for the added ones.

Removals first. Sending removals first lets a connection at the server's topic limit swap
topics. If subscribe() or unsubscribe() changes are waiting to be sent at the same time, they
go out in the same two messages, removals first.

The cost. Data that both a removed and an added topic cover, for example T.AAPL replaced by
T.*, is not delivered between the two messages.

`resubscribe(*topics)`

Sends a subscribe message again for topics the client already counts as subscribed, for example
ones the server never acknowledged.

Malformed topics are logged and skipped.
A topic the client is not subscribed to (and has not scheduled) raises ValueError, and then
nothing is sent.

Tests

Fake server

test_websocket/fake_server.py is a scriptable fake Massive server for these tests. It listens on
a free port and records, for each connection, every message the client sent and the topics it
streams. It can:

script the reply to auth;
refuse handshakes;
close or drop connections;
run a hook on each client message. The hook may close or drop the connection, and
test_websocket/test_fake_server.py checks that it can.

`test_websocket/test_subscriptions.py`

31 tests, which check what the fake server received. They cover:

Concurrency: 16 threads × 10,000 subscribe()/unsubscribe() calls end with the server
streaming exactly client.scheduled_subs.
Latency: with the receive timeout patched to 30 s, changes from any thread reach the server
within 0.5 s.
Sending while receiving: changes sent while a stream of data is being received lose and
reorder no messages.
Cancel and errors: cancelling a connected connect() returns promptly, and a send error is
raised by connect().
Compatibility:
- writing scheduled_subs and schedule_resub directly still works;
- send order of the existing methods is unchanged;
- topics are kept as given.
Wildcards match on the exact channel, and subscribing to a wildcard twice keeps it.
New methods: the messages set_subscriptions() and resubscribe() send.
Reconnects: changes made during an outage are kept, and connect() works again after
close() or a cancel.

Throughput

Throughput is unchanged. Median µs per message against a local server:

Workload	Python	`master`	This PR
20 trades per message, parsed	3.10	123.5	122.3
20 trades per message, parsed	3.14	62.4	62.6
one tiny message, raw	3.10	60.5	58.8
one tiny message, raw	3.14	10.3	10.4

Test runs

make test passes on 3.10 to 3.14 on Ubuntu and Windows, and on 3.13t and 3.14t with
PYTHON_GIL=0.
make lint passes on 3.10.

Known limitation

On Python 3.10 and 3.11, asyncio.wait_for() can ignore a cancellation that arrives just as the
recv() it waits for completes (bpo-42130). The receive loop
then keeps running after cancel().

This existed before and is in the receive loop, which this PR leaves unchanged. The cancellation
tests therefore cancel while the client is idle.

coderabbitai · 2026-10-23T23:04:49Z

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

chipkent

Assessment: not ready to merge yet (medium risk).
The PR is strong and thoughtfully tested, but I found two correctness/design issues that are adversarially important for a client library API.

PR Summary

Intent is clear and valuable: make websocket subscription changes thread-safe, low-latency, and declarative.
Scope is large: ~9 files, +1360/-76, with major logic in massive/websocket/__init__.py and substantial new tests.
Test strategy is strong (race, latency, reconnection, wildcard, contract tests).
CI/workflow changes to include 3.12–3.14 and free-threaded interpreters are appropriate.

Blocking findings

Behavioral compatibility break in set_subscriptions: malformed resend entries are no longer ignored, they now raise ValueError even though malformed topics in desired are documented as “logged and skipped.”
In set_subscriptions, desired is parsed/sanitized via _parse_subscriptions, but resend is taken raw and validated against sanitized topics.
Example: set_subscriptions(["T.AAPL"], resend=["bad-format"]) raises ValueError (unknown contains malformed resend). This is surprising relative to existing method style (subscribe/unsubscribe skip malformed input) and the docstring emphasis on skipping malformed topics.
Fix: parse/sanitize resend with the same parser (or explicitly document that malformed resend is an error and add tests for that strict contract).
Potential stale-state overwrite on reconnect path due to assigning self.scheduled_subs = set(self.subs) after connection close.
In _connect exception path (ConnectionClosedError), code takes lock and sets:
- self.scheduled_subs = set(self.subs)
- self.subs = set()
- self.schedule_resub = True
  This can discard pending user-intended changes that were scheduled but not yet reconciled at the moment of close (especially under heavy concurrent caller activity). Your _reconcile retry logic preserves on send failure, but this path can still reset desired state to last known subscribed state.
  Fix: preserve caller-desired state across reconnects (prefer existing scheduled_subs as source of truth, merged with any necessary resend semantics), rather than clobbering with subs.

Non-blocking but important risks

connect() now delegates to _connect() and toggles _loop/_wake_pending/_woken; structure is good, but there are subtle shared-state flags (_woken) not lock-protected. It’s likely okay due to event-loop-thread ownership, but add a short invariant comment (“_woken read/write only on loop thread except reset before loop attach”) to reduce maintenance hazards.
_recv 3.10 branch creates/cancels tasks per cycle (intended), but cancellation surfaces can be brittle across websockets versions. You have tests, which helps; still worth pinning behavior in a comment near the await asyncio.wait({recv}) drain.
Test additions include large concurrency loops; good for races but may be noisy on constrained runners. Consider marking the heaviest test as stress/slow if flakes appear.

Conventions / project fit

Docstrings and comments are generally clear and in-project style.
Naming and organization are consistent.
New docs entries (set_subscriptions, unsubscribe_all) are good.
pyproject dependency-group split (test vs dev) is clean and justified by CI install strategy.
Tests are readable and targeted, with strong scenario coverage.

Overall verdict

Good PR with solid engineering intent and unusually strong tests, but I would require fixing the two blockers above before merge.

Copilot

🟡 Changes recommended

Inconsistent topic normalization can prevent unsubscription, and disconnect hooks can deadlock the fake server.

2 open issues

What changed in this PR

This PR improves WebSocket subscription synchronization and responsiveness, adds declarative subscription methods, and preserves desired topics across reconnects. It also includes the stacked free-threaded testing infrastructure.

Changes:

Synchronizes subscription updates and wakes the receive loop promptly.
Adds set_subscriptions() and resubscribe(), with corrected wildcard matching.
Adds local concurrency tests and expands CI coverage, including free-threaded Python.

File	Description
`test_websocket/test_subscriptions.py`	Tests subscription changes, wake-ups, wildcards, and reconnects.
`test_websocket/test_fake_server.py`	Tests fake-server behavior and concurrent clients.
`test_websocket/fake_server.py`	Adds a scriptable local WebSocket server.
`test_rest/test_snapshots.py`	Re-registers mocks for a repeated request.
`test_rest/test_dividends.py`	Re-registers mocks for a repeated request.
`test_rest/test_concurrency.py`	Tests a REST client shared across threads.
`test_rest/base.py`	Adds mock cleanup and reusable registration.
`pyproject.toml`	Separates test dependencies from development tools.
`poetry.lock`	Updates dependency groups and resolution.
`massive/websocket/__init__.py`	Implements synchronized subscription reconciliation and wake-ups.
`docs/source/WebSocket.rst`	Documents the two new subscription methods.
`.github/workflows/test.yml`	Expands Python coverage and adds free-threaded jobs.

_{Review effort: Balanced}

chipkent · 2026-10-24T04:52:08Z

+        parsed = self._parse_subscriptions(subscriptions)
+        with self._subs_lock:
+            for s, topic, sym in parsed:
+                self.scheduled_subs.discard(s)


Fixed by removing the inconsistency at its source. On master no method strips stored topics; the stripping came from this PR's parser, and d611483 removes it. The constructor, subscribe() and unsubscribe() now all store and match topics exactly as given, as on master, so your example subscriptions=[" T.AAPL "] followed by unsubscribe(" T.AAPL ") removes the topic. The constructor no longer parses or collapses its topics either, again matching master. Covered by test_topics_are_kept_as_given.

chipkent · 2026-10-24T04:52:09Z

+                if isinstance(msg, dict):
+                    await self._handle(conn, msg)
+                if self.on_message is not None:
+                    await self.on_message(conn, msg)


Fixed in 3adbdc3, which now carries the fake server. Each connection records its handler task, and close() and drop() skip waiting for closed when called from that task. The handler sets closed itself once the hook returns, so neither the hook nor FakeServer.stop() blocks. Regression tests: test_hook_can_close_the_connection and test_hook_can_drop_the_connection in test_websocket/test_fake_server.py. Both time out without the fix and pass with it.

BaseTest turned pook on and registered every mock in setUpClass but never turned it off. Mocks piled up across classes, and pook kept intercepting urllib3 after the REST mock tests had run. That breaks any later test that makes real HTTP requests: pook rejects them with PookNoMatches because it has no mock for them. The next commit adds such a test, which talks to a local HTTP server, and it fails in the full suite without this fix. The leak also hid a problem. Two tests repeat a request, and each mock answers only once, so their second request was served by a copy that an earlier class had registered. Both fail when run on their own. Turn pook off in tearDownClass. Let a test class name the mocks that must answer more than once; those are registered with persist. Other mocks still answer once, because the pagination tests depend on it: a persistent first-page mock would also match the request for the next page.

16 threads share one RESTClient and make 200 calls each against a local HTTP server that echoes every request back. The threads start together behind a barrier, and the server records how many requests it had in flight at once, so the test checks that they really overlapped. Every response is compared with the request that produced it: the path, the exact query string, the Authorization header, and the edge headers. Odd calls send per-request edge headers and even calls send none, so a header leaking from one thread into another call shows up as a mismatch. After the run, the client's default headers must be unchanged: same keys, same values, same order. The calls with headers go through RESTClient._get, which every endpoint method uses. Passing `options` to an endpoint method also adds it to the query string, and that would get in the way of checking the exact query.

Replace the deprecated [tool.poetry.dev-dependencies] section with a `test` group (pook, all the unit tests need) and a `dev` group (linters, type stubs, Sphinx, orjson for the examples). `poetry install` still installs both, so local and CI workflows are unchanged, while `poetry install --only main,test` installs just what `make test` needs. Versions are unchanged. Re-locking also merged a stale duplicate entry for typing-extensions (4.12.2 below Python 3.15, 4.16.0 from 3.15) into 4.16.0, a transitive dependency of the dev tools.

Add a job that runs `make test` on 3.13t and 3.14t with PYTHON_GIL=0, and extend the regular matrix to 3.12, 3.13 and 3.14. The free-threaded job installs only the main and test dependency groups: the dev group is not needed to run the tests, and orjson in it does not support free-threaded Python.

The WebSocket tests use mock_server, which listens on a fixed port, keeps module-level state and cannot script failures, so they cannot check what the client sends or how it behaves when the server misbehaves. FakeServer listens on a free port and records, per connection, every message the client sent and the topics it is subscribed to. It speaks the Massive protocol by default and can script the reply to auth (for example a max_connections status followed by a close), refuse handshakes, close or drop connections, and run a hook on each client message, which may itself close or drop the connection. Its waits are driven by changes to its state, and its state is only touched from its own event loop.

subscribe(), unsubscribe() and unsubscribe_all() changed scheduled_subs and schedule_resub on the caller's thread while connect() read them on the event loop's thread, with no synchronization. A change made between the loop computing the frames and clearing schedule_resub was lost, so the server could end up with a different set than the one requested. A lock now guards both attributes, both where the callers change them and where connect() resets them after an abnormal close; the loop copies the scheduled set and clears the flag under it, and sends the frames from the copy. The methods replace scheduled_subs with a new set instead of changing it in place, so another thread can read or iterate it without the lock. It remains a set, and the frames are sent in the same order as before.

connect() applied scheduled subscription changes only when its receive loop came round again, after a message or the one-second recv timeout, so on a quiet feed a subscribe() took up to a second to reach the server. Each connection now gets one sender task, started after authentication. It waits on an asyncio.Event and, when woken, sends the frames that move the server from subs to the latest scheduled set; it starts by sending the whole set, so a reconnect re-sends the desired state. Being the only sender, it never interleaves frames, and a burst of changes is merged into one round. subscribe(), unsubscribe() and unsubscribe_all() change the scheduled set under the lock and then, if the change is the first since the sender last read it, set the event through call_soon_threadsafe, which never blocks the caller. The sender sends on the socket it was started for, not self.websocket. When the connection ends, the sender is cancelled and awaited before subs is reset. An error other than ConnectionClosed ends the sender and is raised by connect() at the next turn of the receive loop; a closed connection is left to recv() to report, after the messages already received are read. The receive loop still looks at schedule_resub at the top of each iteration, now to wake the sender, so code that changes scheduled_subs itself and sets the flag keeps working. Otherwise the receive loop is unchanged. Sending while another coroutine waits in recv() on the same connection is safe in websockets 14.0 to 15.0.1 (websockets/asyncio/connection.py): recv() only reads from the message assembler (14.0 L270, 14.2 L294, 15.0.1 L303), whose only concurrency check rejects a second concurrent reader (asyncio/messages.py, Assembler.get: 14.0 L142, 14.2 L141, 15.0.1 L150), and send() waits only for a fragmented send in progress (14.0 L427, 14.2 L451, 15.0.1 L469) before writing through send_context() (14.0 L434, 14.2 L458, 15.0.1 L476).

subscribe("T.*") and unsubscribe("T.*") discard the other subscribed topics on the channel, but they matched with t.startswith(channel), so subscribe("A.*") also dropped AM.AAPL and TX topics went with T.*. The test also matched the wildcard itself, so subscribing again to an already subscribed T.* unsubscribed T.*. The collapse now matches the channel exactly and keeps the wildcard being subscribed. As before, it applies only to subscribed topics, and topics are stored as given.

There was no way to replace the subscriptions in one step: a caller that knows the full set it wants had to compute the difference itself and call unsubscribe() and subscribe(), which another thread or the event loop could interleave with. set_subscriptions() atomically replaces the scheduled set; the client then sends at most one unsubscribe message for the removed topics followed by at most one subscribe message for the added ones. Removals go first so that a client at the server's subscription limit can swap topics; an update that also carries subscribe() or unsubscribe() changes made since the last one sends removals first too. Updates made only through subscribe() and unsubscribe() keep their order, additions first. The method is new, so no existing behaviour changes.

The client sends a subscribe message only for topics missing from subs, so there was no way to send one again for a topic it already counts as subscribed, e.g. when the server never acknowledged it. resubscribe() adds such topics to the next subscribe message, together with any topics being added, and raises ValueError for a topic that is not subscribed or scheduled. The method is new, so no existing behaviour changes.

After an abnormal close, connect() replaced scheduled_subs with subs, the topics the old connection had applied. A subscribe() or unsubscribe() made before the close but not yet sent, or made during the outage, was lost, and an unsubscribed topic could come back. scheduled_subs is now the authoritative desired set and is left alone: only subs is reset, and the new connection subscribes to the scheduled set. Likewise, each connect() call starts with empty subs, so calling it again after close() or after cancelling it subscribes to the scheduled set instead of assuming the old connection's topics are still subscribed.

chipkent · 2026-10-24T13:55:35Z

Superseded by smaller PRs split by content: #18, #20, #21, #22. Closing this preview; the replacements carry the same changes, opt-in where behaviour changes.

chipkent force-pushed the fix/ws-subscriptions branch 2 times, most recently from 946c4b1 to ac207b9 Compare October 24, 2026 02:39

chipkent changed the title ~~[CI only] fix/ws-subscriptions~~ fix: thread-safe, prompt, and declarative WebSocket subscriptions Oct 24, 2026

chipkent force-pushed the fix/ws-subscriptions branch 2 times, most recently from 4cacf63 to e907f1e Compare October 24, 2026 02:58

chipkent commented Oct 24, 2026

View reviewed changes

chipkent force-pushed the fix/ws-subscriptions branch 2 times, most recently from 7b87471 to d8e5255 Compare October 24, 2026 03:37

chipkent requested a balanced review from Copilot October 24, 2026 03:48

Copilot started reviewing on behalf of chipkent October 24, 2026 03:48

View session

Copilot AI reviewed Oct 24, 2026

View reviewed changes

chipkent force-pushed the fix/ws-subscriptions branch from d8e5255 to ad354e0 Compare October 24, 2026 04:02

chipkent force-pushed the fix/ws-subscriptions branch from ad354e0 to edef902 Compare October 24, 2026 04:20

chipkent added 10 commits October 23, 2026 22:21

chipkent force-pushed the fix/ws-subscriptions branch from edef902 to adfb99b Compare October 24, 2026 04:47

chipkent changed the title ~~fix: thread-safe, prompt, and declarative WebSocket subscriptions~~ fix: keep WebSocket subscription changes in sync and send them promptly Oct 24, 2026

chipkent closed this Oct 24, 2026

Navigation Menu

Repository navigation

fix: keep WebSocket subscription changes in sync and send them promptly - #2

fix: keep WebSocket subscription changes in sync and send them promptly#2
chipkent wants to merge 11 commits into
masterfrom
fix/ws-subscriptions

chipkent commented Oct 23, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Oct 23, 2026

Draft PR not reviewed

Uh oh!

chipkent left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

chipkent Oct 24, 2026

Uh oh!

chipkent Oct 24, 2026

Uh oh!

chipkent commented Oct 24, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

chipkent commented Oct 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Compatibility

Bug fixes

1. Subscription changes made while connect() runs can be lost (#<issue: F1a>)

2. Subscription changes are sent up to 1 s late (#<issue: F1a>)

3. Wildcard topics match by string prefix (#<issue: F1b>)

4. Changes made around a reconnect are lost

New methods (#<issue: F1c>)

set_subscriptions(desired)

resubscribe(*topics)

Tests

Fake server

test_websocket/test_subscriptions.py

Throughput

Test runs

Known limitation

Uh oh!

coderabbitai Bot commented Oct 23, 2026

Draft PR not reviewed

Uh oh!

chipkent left a comment

Choose a reason for hiding this comment

PR Summary

Blocking findings

Non-blocking but important risks

Conventions / project fit

Overall verdict

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

chipkent Oct 24, 2026

Choose a reason for hiding this comment

Uh oh!

chipkent Oct 24, 2026

Choose a reason for hiding this comment

Uh oh!

chipkent commented Oct 24, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

chipkent commented Oct 23, 2026 •

edited

Loading

1. Subscription changes made while `connect()` runs can be lost (#<issue: F1a>)

`set_subscriptions(desired)`

`resubscribe(*topics)`

`test_websocket/test_subscriptions.py`