From 0c0a16c7220ae5cce9b2a00ced84b90ef5afcc07 Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Sat, 24 Oct 2026 00:16:21 -0600 Subject: [PATCH] feat: send WebSocket subscription changes without waiting for recv connect() sent subscription changes only when its receive loop came round again, after the next message or the one-second recv timeout. On a quiet feed, a subscribe() took up to a second to reach the server. Each connection now has one task that sends subscription changes. It starts after authentication by sending the whole scheduled set, then waits on an asyncio.Event. subscribe(), unsubscribe() and unsubscribe_all() set the event through call_soon_threadsafe(), which never blocks the caller. Being the only sender, the task never interleaves messages, and it merges a burst of changes into one subscribe and one unsubscribe message. The receive loop still checks schedule_resub on every turn, now to wake the task, so code that changes scheduled_subs itself and sets the flag keeps working. An error other than a closed connection ends the task and is raised by connect(), as before. The task is stopped when the connection ends. The receive path is otherwise unchanged. Sending while another coroutine waits in recv() on the same connection is safe in websockets 14.0 through 17.3, the current release (websockets/asyncio/connection.py). recv() only reads from the message assembler (14.0 L270, 14.2 L294, 15.0.1 L303, 16.1.1 L305, 17.3 L298), whose only concurrency check rejects a second concurrent reader (asyncio/messages.py, Assembler.get: 14.0 L142, 14.2 L141, 15.0.1 L150, 16.1.1 L149, 17.3 L148). send() waits only for a fragmented send in progress (fragmented_send_waiter, named send_in_progress in 16.1.1 and 17.3: 14.0 L427, 14.2 L451, 15.0.1 L469, 16.1.1 L472, 17.3 L466) before writing through send_context() (14.0 L434, 14.2 L458, 15.0.1 L476, 16.1.1 L479, 17.3 L473). --- massive/websocket/__init__.py | 149 ++++++++++++++++++++++++++-------- test_websocket/test_sender.py | 61 ++++++++++++++ 2 files changed, 175 insertions(+), 35 deletions(-) create mode 100644 test_websocket/test_sender.py diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py index 3bf043be..eea7a70c 100644 --- a/massive/websocket/__init__.py +++ b/massive/websocket/__init__.py @@ -1,7 +1,18 @@ import os from collections import Counter from enum import Enum -from typing import Optional, Union, List, Set, Callable, Awaitable, Any, TypeVar, cast +from typing import ( + Optional, + Union, + List, + Set, + Callable, + Awaitable, + Any, + Tuple, + TypeVar, + cast, +) import logging import json import functools @@ -29,11 +40,14 @@ def _changes_subscriptions(method: _F) -> _F: # runs a WebSocketClient method that changes scheduled_subs under the - # client's lock, so it is safe to call from any thread + # client's lock, so it is safe to call from any thread, then wakes the + # task that sends the change @functools.wraps(method) def wrapper(self, *args, **kwargs): with self._subs_lock: - return method(self, *args, **kwargs) + result = method(self, *args, **kwargs) + self._notify() + return result return cast(_F, wrapper) @@ -97,6 +111,9 @@ def __init__( # unsubscribe() and unsubscribe_all() change on the caller's thread # while connect() reads them self._subs_lock = threading.Lock() + # 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 self._server_errors: Counter = Counter() self._server_errors_since = 0.0 self.ca_certs = ca_certs @@ -154,39 +171,39 @@ async def connect( # e.g. max_connections, which the server sends before closing self._log_server_statuses(auth_msg_parsed) self._flush_server_errors(force=True) - while True: - if self.schedule_resub: - # a change made from here on sets schedule_resub again - with self._subs_lock: - desired = set(self.scheduled_subs) - self.schedule_resub = False - logger.debug("reconciling: %s %s", self.subs, desired) - new_subs = desired.difference(self.subs) - await self._subscribe(new_subs) - old_subs = self.subs.difference(desired) - await self._unsubscribe(old_subs) - self.subs = desired - - try: - cmsg: Union[List[WebSocketMessage], Union[str, bytes]] = ( - await asyncio.wait_for(s.recv(), timeout=1) - ) - except asyncio.TimeoutError: + # subscription changes go out from a task of their own, + # starting with the whole scheduled set + sender, wake = self._start_sender() + try: + while True: + if sender.done(): + self._check_sender(sender) + if self.schedule_resub: + # e.g. set by code that changed scheduled_subs itself + wake.set() + + try: + cmsg: Union[List[WebSocketMessage], Union[str, bytes]] = ( + await asyncio.wait_for(s.recv(), timeout=1) + ) + except asyncio.TimeoutError: + self._flush_server_errors() + continue + + if not self.raw: + # we know cmsg is Data + msgJson = self.json.loads(cmsg) # type: ignore + for m in msgJson: + if m["ev"] == "status": + logger.debug("status: %s", m["message"]) + self._log_server_statuses(msgJson) + cmsg = parse(msgJson, logger, self.market) self._flush_server_errors() - continue - - if not self.raw: - # we know cmsg is Data - msgJson = self.json.loads(cmsg) # type: ignore - for m in msgJson: - if m["ev"] == "status": - logger.debug("status: %s", m["message"]) - self._log_server_statuses(msgJson) - cmsg = parse(msgJson, logger, self.market) - self._flush_server_errors() - - if len(cmsg) > 0: - await processor(cmsg) # type: ignore + + if len(cmsg) > 0: + await processor(cmsg) # type: ignore + finally: + await self._stop_sender(sender) except ConnectionClosedOK as e: logger.debug("connection closed (OK): %s", e) # the server sends a status (e.g. slow_consumer) before closing @@ -283,6 +300,68 @@ async def _unsubscribe(self, topics: Union[List[str], Set[str]]): self.json.dumps({"action": "unsubscribe", "params": subs}) ) + def _start_sender(self) -> Tuple["asyncio.Task[None]", asyncio.Event]: + """Start the task that sends subscription changes; see _notify.""" + wake = asyncio.Event() + wake.set() # first, the whole scheduled set + self._wake = (asyncio.get_running_loop(), wake) + return asyncio.ensure_future(self._send_changes(wake)), wake + + async def _send_changes(self, wake: asyncio.Event): + # The only sender of subscription messages, so messages for two + # changes never interleave. Sending while connect() waits in recv() is + # safe: in websockets, only two concurrent recv() calls conflict. + while True: + await wake.wait() + wake.clear() + await self._reconcile() + + def _check_sender(self, sender: "asyncio.Task[None]"): + # raise what ended the sender, unless the connection closed: then + # recv() raises too, once the messages already received are read + if sender.cancelled(): + return + exc = sender.exception() + if exc is not None and not isinstance( + exc, (ConnectionClosedOK, ConnectionClosedError) + ): + raise exc + + async def _stop_sender(self, sender: "asyncio.Task[None]"): + self._wake = None + sender.cancel() + await asyncio.wait({sender}) # does not raise + self._check_sender(sender) + + def _notify(self): + """Wake the sender, if a connection is open. Safe from any thread.""" + waker = self._wake + if waker is None: + return + loop, wake = waker + try: + loop.call_soon_threadsafe(wake.set) + except RuntimeError: + pass # the loop is closed + + async def _reconcile(self): + """Send the changes that move the server from subs to scheduled_subs.""" + with self._subs_lock: + if not self.schedule_resub: + return + desired = set(self.scheduled_subs) + self.schedule_resub = False + logger.debug("reconciling: %s %s", self.subs, desired) + try: + await self._subscribe(desired.difference(self.subs)) + await self._unsubscribe(self.subs.difference(desired)) + except BaseException: + # e.g. the connection closed: send it all on the next connection + with self._subs_lock: + self.schedule_resub = True + raise + self.subs = desired + @staticmethod def _parse_subscription(s: str): s = s.strip() diff --git a/test_websocket/test_sender.py b/test_websocket/test_sender.py new file mode 100644 index 00000000..71cb2898 --- /dev/null +++ b/test_websocket/test_sender.py @@ -0,0 +1,61 @@ +from massive import WebSocketClient +from base_ws import BaseTest +from mock_server import inbox, port +from unittest import mock +import asyncio + + +class NoRecvTimeout: + # asyncio, except that connect() waits in recv() until a message arrives + def __getattr__(self, name): + return getattr(asyncio, name) + + @staticmethod + def wait_for(aw, timeout): + return asyncio.wait_for(aw, None) + + +def sent(conn, action): + # the topics in each `action` message one connection received + return [set(m["params"].split(",")) for m in conn if m.get("action") == action] + + +class SenderTest(BaseTest): + async def connected(self, *subscriptions): + # a client whose receive loop has handled every message so far + c = WebSocketClient( + api_key="", + feed=f"localhost:{port}", + secure=False, + subscriptions=list(subscriptions), + ) + received = [] + + async def processor(msgs): + received.append(msgs) + + task = asyncio.create_task(c.connect(processor)) + # each subscribed topic gets two trade messages + await self.until(lambda: len(received) == 2 * len(subscriptions)) + return c, task + + async def test_change_is_sent_while_waiting_for_a_message(self): + with mock.patch("massive.websocket.asyncio", NoRecvTimeout()): + c, task = await self.connected("T.AAPL") + await asyncio.to_thread(c.subscribe, "T.MSFT") + await self.until(lambda: len(sent(inbox[0], "subscribe")) == 2) + await asyncio.to_thread(c.unsubscribe, "T.AAPL") + await self.until(lambda: sent(inbox[0], "unsubscribe")) + self.assertEqual(sent(inbox[0], "subscribe"), [{"T.AAPL"}, {"T.MSFT"}]) + self.assertEqual(sent(inbox[0], "unsubscribe"), [{"T.AAPL"}]) + await c.close() + await task + + async def test_scheduled_subs_can_still_be_changed_directly(self): + c, task = await self.connected("T.AAPL") + c.scheduled_subs.add("T.MSFT") + c.schedule_resub = True + await self.until(lambda: len(sent(inbox[0], "subscribe")) == 2) + self.assertEqual(sent(inbox[0], "subscribe")[1], {"T.MSFT"}) + await c.close() + await task