From 6c8a38d0e27f90815d220901e369f92cc50ac97b Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Sat, 24 Oct 2026 00:18:01 -0600 Subject: [PATCH] feat: add WebSocketClient.set_subscriptions() and resubscribe() Two things a long-running client needs were not possible: - Replacing one set of topics with another, removals first. The server enforces a per-connection subscription limit (it answers "Subscription limit reached for feed", which the client logs since #1039). The client always sends the subscribe message before the unsubscribe message, so a swap made at the limit asks for more topics than the limit allows before the unsubscribe arrives. - Sending a subscribe message again for a topic the server never acknowledged. subscribe() sends only topics missing from subs, so a topic already sent can never be sent twice. set_subscriptions(desired) replaces the scheduled topics in one step. The next update then sends one unsubscribe message for the removed topics, followed by one subscribe message for the added ones. resubscribe(*topics) adds topics to the next subscribe message. A topic that is neither subscribed nor scheduled raises ValueError, and then nothing is sent. subscribe(), unsubscribe() and unsubscribe_all() keep their order: subscribe first, then unsubscribe. --- docs/source/WebSocket.rst | 10 ++++ massive/websocket/__init__.py | 58 +++++++++++++++++++++++- test_websocket/test_set_subscriptions.py | 55 ++++++++++++++++++++++ 3 files changed, 121 insertions(+), 2 deletions(-) create mode 100644 test_websocket/test_set_subscriptions.py diff --git a/docs/source/WebSocket.rst b/docs/source/WebSocket.rst index c624d9d2..f557525b 100644 --- a/docs/source/WebSocket.rst +++ b/docs/source/WebSocket.rst @@ -34,6 +34,16 @@ Unsubscribe ============================ .. automethod:: massive.WebSocketClient.unsubscribe +============================ +Set subscriptions +============================ +.. automethod:: massive.WebSocketClient.set_subscriptions + +============================ +Resubscribe +============================ +.. automethod:: massive.WebSocketClient.resubscribe + ============================ Close ============================ diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py index eea7a70c..dadd68e0 100644 --- a/massive/websocket/__init__.py +++ b/massive/websocket/__init__.py @@ -12,6 +12,7 @@ Tuple, TypeVar, cast, + Iterable, ) import logging import json @@ -111,6 +112,11 @@ def __init__( # unsubscribe() and unsubscribe_all() change on the caller's thread # while connect() reads them self._subs_lock = threading.Lock() + # set_subscriptions() was called since the last update was sent, so + # send removals first; guarded by _subs_lock + self._removals_first = False + # topics resubscribe() asked to send again; guarded by _subs_lock + self._resend: Set[str] = set() # while a connection is open: its loop, and the event that wakes its # sender task (see _start_sender) self._wake: Optional[Tuple[asyncio.AbstractEventLoop, asyncio.Event]] = None @@ -350,11 +356,19 @@ async def _reconcile(self): if not self.schedule_resub: return desired = set(self.scheduled_subs) + removals_first, self._removals_first = self._removals_first, False + resend, self._resend = self._resend, set() self.schedule_resub = False logger.debug("reconciling: %s %s", self.subs, desired) + added = desired.difference(self.subs) | resend.intersection(desired) + removed = self.subs.difference(desired) try: - await self._subscribe(desired.difference(self.subs)) - await self._unsubscribe(self.subs.difference(desired)) + if removals_first: + await self._unsubscribe(removed) + await self._subscribe(added) + else: + await self._subscribe(added) + await self._unsubscribe(removed) except BaseException: # e.g. the connection closed: send it all on the next connection with self._subs_lock: @@ -423,6 +437,46 @@ def unsubscribe_all(self): self.scheduled_subs = set() self.schedule_resub = True + @_changes_subscriptions + def set_subscriptions(self, desired: Iterable[str]): + """ + Replace all subscriptions with ``desired``, in one step. Safe to call + from any thread. + + The client sends one unsubscribe message for the topics no longer + desired, then one subscribe message for the new ones, together with + any changes from ``subscribe()`` or ``unsubscribe()`` not sent yet. + Removals go first, so a client at the server's subscription limit + can swap topics; data that a removed and an added topic both cover + (e.g. ``T.AAPL`` replaced by ``T.*``) is not delivered in between. + Wildcards such as ``T.*`` are ordinary topics here. + + :param desired: Every topic to be subscribed to afterwards. Malformed + topics are logged and skipped. + """ + topics = {s for s in desired if self._parse_subscription(s)[0] is not None} + self.scheduled_subs = topics + self._removals_first = True + self.schedule_resub = True + + @_changes_subscriptions + def resubscribe(self, *topics: str): + """ + Send a subscribe message again for topics already subscribed, e.g. + ones the server never acknowledged. Safe to call from any thread. + + :param topics: Topics to send again (args). Malformed topics are + logged and skipped. + :raises ValueError: If a topic is not subscribed or scheduled to be; + nothing is sent then. + """ + again = {s for s in topics if self._parse_subscription(s)[0] is not None} + unknown = again.difference(self.scheduled_subs) + if unknown: + raise ValueError(f"not subscribed: {sorted(unknown)}") + self._resend |= again + self.schedule_resub = True + async def close(self): """ Close the websocket connection. diff --git a/test_websocket/test_set_subscriptions.py b/test_websocket/test_set_subscriptions.py new file mode 100644 index 00000000..3d4cbc21 --- /dev/null +++ b/test_websocket/test_set_subscriptions.py @@ -0,0 +1,55 @@ +from massive import WebSocketClient +from base_ws import BaseTest +from mock_server import inbox, port +import asyncio + + +def sent(conn): + # the subscribe and unsubscribe messages one connection received + return [ + (m["action"], set(m["params"].split(","))) + for m in conn + if m.get("action") in ("subscribe", "unsubscribe") + ] + + +class SetSubscriptionsTest(BaseTest): + async def connected(self, *subscriptions): + c = WebSocketClient( + api_key="", + feed=f"localhost:{port}", + secure=False, + subscriptions=list(subscriptions), + ) + task = asyncio.create_task(c.connect(lambda msgs: asyncio.sleep(0))) + await self.until(lambda: inbox and sent(inbox[0])) + return c, task + + async def test_set_subscriptions_sends_removals_first(self): + c, task = await self.connected("T.AAPL", "T.MSFT") + c.set_subscriptions(["T.AAPL", "T.NVDA"]) + await self.until(lambda: len(sent(inbox[0])) == 3) + self.assertEqual( + sent(inbox[0])[1:], + [("unsubscribe", {"T.MSFT"}), ("subscribe", {"T.NVDA"})], + ) + await c.close() + await task + + async def test_resubscribe_sends_again(self): + c, task = await self.connected("T.AAPL") + c.resubscribe("T.AAPL") + await self.until(lambda: len(sent(inbox[0])) == 2) + self.assertEqual(sent(inbox[0])[1], ("subscribe", {"T.AAPL"})) + await c.close() + await task + + async def test_resubscribe_unknown_topic_sends_nothing(self): + c, task = await self.connected("T.AAPL") + with self.assertRaises(ValueError): + c.resubscribe("T.AAPL", "T.MSFT") + c.subscribe("T.NVDA") + await self.until(lambda: len(sent(inbox[0])) == 2) + self.assertEqual(sent(inbox[0])[1], ("subscribe", {"T.NVDA"})) + await c.close() + await task