diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py index 0f2d9610..3bf043be 100644 --- a/massive/websocket/__init__.py +++ b/massive/websocket/__init__.py @@ -1,11 +1,13 @@ import os from collections import Counter from enum import Enum -from typing import Optional, Union, List, Set, Callable, Awaitable, Any +from typing import Optional, Union, List, Set, Callable, Awaitable, Any, TypeVar, cast import logging import json +import functools import asyncio import ssl +import threading import time from .models import * from websockets.asyncio.client import connect, ClientConnection @@ -22,6 +24,19 @@ # "max_connections", "slow_consumer", "force_disconnect") is logged as a warning. OK_STATUSES = {"connected", "auth_success", "success"} +_F = TypeVar("_F", bound=Callable[..., Any]) + + +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 + @functools.wraps(method) + def wrapper(self, *args, **kwargs): + with self._subs_lock: + return method(self, *args, **kwargs) + + return cast(_F, wrapper) + class WebSocketClient: def __init__( @@ -78,6 +93,10 @@ def __init__( subscriptions = [] self.scheduled_subs: Set[str] = set(subscriptions) self.schedule_resub = True + # guards scheduled_subs and schedule_resub, which subscribe(), + # unsubscribe() and unsubscribe_all() change on the caller's thread + # while connect() reads them + self._subs_lock = threading.Lock() self._server_errors: Counter = Counter() self._server_errors_since = 0.0 self.ca_certs = ca_certs @@ -137,16 +156,16 @@ async def connect( self._flush_server_errors(force=True) while True: if self.schedule_resub: - logger.debug( - "reconciling: %s %s", self.subs, self.scheduled_subs - ) - new_subs = self.scheduled_subs.difference(self.subs) + # 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(self.scheduled_subs) + old_subs = self.subs.difference(desired) await self._unsubscribe(old_subs) - self.subs = self.scheduled_subs - self.subs = set(self.scheduled_subs) - self.schedule_resub = False + self.subs = desired try: cmsg: Union[List[WebSocketMessage], Union[str, bytes]] = ( @@ -274,9 +293,10 @@ def _parse_subscription(s: str): return split + @_changes_subscriptions def subscribe(self, *subscriptions: str): """ - Subscribe to given subscriptions. + Subscribe to given subscriptions. Safe to call from any thread. :param subscriptions: Subscriptions (args) """ @@ -294,9 +314,10 @@ def subscribe(self, *subscriptions: str): self.schedule_resub = True + @_changes_subscriptions def unsubscribe(self, *subscriptions: str): """ - Unsubscribe from given subscriptions. + Unsubscribe from given subscriptions. Safe to call from any thread. :param subscriptions: Subscriptions (args) """ @@ -315,9 +336,10 @@ def unsubscribe(self, *subscriptions: str): self.schedule_resub = True + @_changes_subscriptions def unsubscribe_all(self): """ - Unsubscribe from all subscriptions. + Unsubscribe from all subscriptions. Safe to call from any thread. """ self.scheduled_subs = set() self.schedule_resub = True diff --git a/test_websocket/test_threads.py b/test_websocket/test_threads.py new file mode 100644 index 00000000..0f3bc215 --- /dev/null +++ b/test_websocket/test_threads.py @@ -0,0 +1,38 @@ +from massive import WebSocketClient +from base_ws import BaseTest +from mock_server import inbox, port +import asyncio + + +def streamed(conn): + # the topics one connection was asked to stream + topics = set() + for m in conn: + if m.get("action") == "subscribe": + topics |= set(m["params"].split(",")) + elif m.get("action") == "unsubscribe": + topics -= set(m["params"].split(",")) + return topics + + +class ThreadTest(BaseTest): + async def test_change_made_while_sending_is_sent(self): + c = WebSocketClient( + api_key="", feed=f"localhost:{port}", secure=False, subscriptions=["T.AAPL"] + ) + send = c._subscribe + changed = False + + async def send_and_change(topics): + nonlocal changed + await send(topics) + if not changed: + changed = True + # another thread's change, made while the loop is sending + c.subscribe("T.LATE") + + c._subscribe = send_and_change + task = asyncio.create_task(c.connect(lambda msgs: asyncio.sleep(0))) + await self.until(lambda: inbox and "T.LATE" in streamed(inbox[0])) + await c.close() + await task