feat: add WebSocketClient.request_close()

docs/source/WebSocket.rst

@@ -39,6 +39,11 @@ Close
 ============================
 .. automethod:: massive.WebSocketClient.close
 
+============================
+Request close
+============================
+.. automethod:: massive.WebSocketClient.request_close
+
 .. _Stocks getting started: https://massive.com/docs/stocks/ws_getting-started
 .. _Options getting started: https://massive.com/docs/options/ws_getting-started
 .. _Forex getting started: https://massive.com/docs/forex/ws_getting-started

massive/websocket/__init__.py

@@ -7,6 +7,7 @@
 import asyncio
 import random
 import ssl
+import threading
 import time
 from .models import *
 from websockets.asyncio.client import connect, ClientConnection, process_exception
@@ -113,6 +114,12 @@ def __init__(
             self.json = custom_json
         else:
             self.json = json
+        # set by request_close(), and by close() while connect() runs;
+        # connect() returns instead of reconnecting
+        self._close_requested = threading.Event()
+        # the event loop connect() runs on, while it runs
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+        self._closing: Optional[asyncio.Future] = None
 
     # https://websockets.readthedocs.io/en/stable/reference/client.html#opening-a-connection
     async def connect(
@@ -132,6 +139,24 @@ async def connect(
         :param kwargs: Passed to websockets' connect(), e.g. process_exception.
         :raises AuthError: If invalid API key is supplied.
         """
+        self._loop = asyncio.get_running_loop()
+        try:
+            await self._connect(processor, close_timeout, **kwargs)
+        finally:
+            self._loop = None
+            self._close_requested.clear()
+
+    async def _connect(
+        self,
+        processor: Union[
+            Callable[[List[WebSocketMessage]], Awaitable],
+            Callable[[Union[str, bytes]], Awaitable],
+        ],
+        close_timeout: int,
+        **kwargs,
+    ):
+        if self._close_requested.is_set():
+            return
         # a new connection has no subscriptions, whatever an earlier call left
         self.subs = set()
         self.schedule_resub = True
@@ -161,6 +186,9 @@ def on_connect_error(exc: Exception) -> Optional[Exception]:
             process_exception=on_connect_error,
             **kwargs,
         ):
+            if self._close_requested.is_set():
+                await s.close()  # requested while connecting
+                return
             self.websocket = s
             connected_at = None
             try:
@@ -230,6 +258,8 @@ def on_connect_error(exc: Exception) -> Optional[Exception]:
                 last_exc = e
                 logger.debug("connection closed (ERR): %s", e)
                 self._flush_server_errors(force=True)
+            if self._close_requested.is_set():
+                return
             reconnects += 1
             # keep scheduled_subs, which may hold changes not sent yet,
             # and subscribe to all of it on the next connection
@@ -250,6 +280,8 @@ def on_connect_error(exc: Exception) -> Optional[Exception]:
                 delay = self._reconnect_delay(failures)
                 logger.debug("reconnecting in %.1f s", delay)
                 await asyncio.sleep(delay)
+                if self._close_requested.is_set():
+                    return  # requested during the wait
 
         if last_exc is not None:
             raise last_exc
@@ -398,13 +430,42 @@ def unsubscribe_all(self):
 
     async def close(self):
         """
-        Close the websocket connection.
+        Close the websocket connection. If ``connect()`` is running, it
+        returns instead of reconnecting, also during a wait before
+        reconnecting.
         """
         logger.debug("closing")
+        if self._loop is not None:
+            self._close_requested.set()  # connect() clears it when it returns
         self._flush_server_errors(force=True)
 
         if self.websocket:
             await self.websocket.close()
             self.websocket = None
         else:
             logger.warning("no websocket open to close")
+
+    def request_close(self):
+        """
+        Close the websocket connection and make ``connect()`` return instead
+        of reconnecting. Unlike ``close()``, this is safe to call from any
+        thread, e.g. to stop ``run()``, and it does not wait.
+
+        An attempt to connect, or a wait before reconnecting, is not
+        interrupted: ``connect()`` returns once it ends. If ``connect()`` is
+        not running, the next call returns at once.
+        """
+        logger.debug("close requested")
+        self._close_requested.set()
+        loop = self._loop
+        if loop is None:
+            return
+        try:
+            loop.call_soon_threadsafe(self._close_soon)
+        except RuntimeError:
+            pass  # the loop is closed, so connect() has returned
+
+    def _close_soon(self):
+        # on connect()'s loop; recv() then raises and connect() returns
+        if self.websocket is not None:
+            self._closing = asyncio.ensure_future(self.websocket.close())

test_websocket/test_request_close.py

@@ -0,0 +1,82 @@
+from massive import WebSocketClient
+from base_ws import BaseTest
+from mock_server import inbox, port
+import asyncio
+import threading
+from unittest import mock
+
+
+class StopInSleep:
+    # asyncio, except that sleep() awaits stop() and returns at once
+    def __init__(self, stop):
+        self.stop = stop
+
+    def __getattr__(self, name):
+        return getattr(asyncio, name)
+
+    async def sleep(self, delay):
+        await self.stop()
+
+
+class RequestCloseTest(BaseTest):
+    def client(self, *subscriptions, **kwargs):
+        return WebSocketClient(
+            api_key="",
+            feed=f"localhost:{port}",
+            secure=False,
+            subscriptions=list(subscriptions),
+            **kwargs,
+        )
+
+    async def test_stops_run_from_another_thread(self):
+        c = self.client("T.AAPL", reconnect_on_close=True)
+        thread = threading.Thread(target=c.run, args=(lambda msgs: None,))
+        thread.start()
+        await self.until(lambda: inbox and len(inbox[0]) == 2)
+        c.request_close()
+        await asyncio.to_thread(thread.join, 5)
+        self.assertFalse(thread.is_alive())
+        self.assertEqual(len(inbox), 1)
+
+    async def test_stops_reconnecting(self):
+        # DROP.X makes the server drop the connection
+        c = self.client("DROP.X", max_reconnects=None)
+
+        async def processor(msgs):
+            c.request_close()
+
+        await asyncio.wait_for(c.connect(processor), 5)
+        self.assertEqual(len(inbox), 1)
+
+    async def test_before_connect(self):
+        c = self.client("T.AAPL")
+        c.request_close()
+        await asyncio.wait_for(c.connect(lambda msgs: asyncio.sleep(0)), 5)
+        self.assertEqual(inbox, [])
+        # the request is used up
+        task = asyncio.create_task(c.connect(lambda msgs: asyncio.sleep(0)))
+        await self.until(lambda: inbox and len(inbox[0]) == 2)
+        await c.close()
+        await task
+
+    async def test_close_before_connect_is_unchanged(self):
+        c = self.client("T.AAPL")
+        await c.close()  # no connection yet: only logs a warning
+        task = asyncio.create_task(c.connect(lambda msgs: asyncio.sleep(0)))
+        await self.until(lambda: inbox and len(inbox[0]) == 2)
+        await c.close()
+        await task
+
+    async def stop_during_wait(self, stop):
+        # DROP.X makes the server drop the connection; stop() runs during
+        # the wait before reconnecting
+        c = self.client("DROP.X", max_reconnects=None, reconnect_backoff=1)
+        with mock.patch("massive.websocket.asyncio", StopInSleep(lambda: stop(c))):
+            await asyncio.wait_for(c.connect(lambda msgs: asyncio.sleep(0)), 5)
+        self.assertEqual(len(inbox), 1)
+
+    async def test_stops_a_wait_before_reconnecting(self):
+        await self.stop_during_wait(lambda c: asyncio.to_thread(c.request_close))
+
+    async def test_close_stops_a_wait_before_reconnecting(self):
+        await self.stop_during_wait(lambda c: c.close())