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