fix: make WebSocket subscribe and unsubscribe thread-safe

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

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