From e818c9cfacba5dde9310f49d5bd8104c2efa12b3 Mon Sep 17 00:00:00 2001 From: Chip Kent <5250374+chipkent@users.noreply.github.com> Date: Sat, 24 Oct 2026 00:24:22 -0600 Subject: [PATCH] feat: add WebSocketClient.request_close() close() is a coroutine that must run on connect()'s event loop. run() creates that loop itself, so another thread cannot stop a client started with run(): asyncio.run(client.close()) from another thread fails with "attached to a different loop". request_close() can be called from any thread and does not wait. It sets a flag and, while connect() runs, closes the connection through loop.call_soon_threadsafe(). connect() checks the flag before it reconnects and after a wait before reconnecting, so it returns instead, also with reconnect_on_close. An attempt to connect, or a wait before reconnecting, is not interrupted: connect() returns once it ends, without opening another connection. If connect() is not running, the next call returns at once. close() called while connect() runs now sets the same flag. Before, a close() during the wait before a reconnect (reconnect_backoff > 0) or during a connect or reconnect handshake found no open connection to close, and connect() went on to connect. A close() before connect() starts is unchanged: the next connect() proceeds. With the default reconnect_backoff of 0 there is no wait, so the only difference is a close() made during a connect or reconnect handshake. --- docs/source/WebSocket.rst | 5 ++ massive/websocket/__init__.py | 63 ++++++++++++++++++++- test_websocket/test_request_close.py | 82 ++++++++++++++++++++++++++++ 3 files changed, 149 insertions(+), 1 deletion(-) create mode 100644 test_websocket/test_request_close.py diff --git a/docs/source/WebSocket.rst b/docs/source/WebSocket.rst index c624d9d2..e4a6b412 100644 --- a/docs/source/WebSocket.rst +++ b/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 diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py index 5d2ebdbe..e6703de4 100644 --- a/massive/websocket/__init__.py +++ b/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,9 +430,13 @@ 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: @@ -408,3 +444,28 @@ async def close(self): 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()) diff --git a/test_websocket/test_request_close.py b/test_websocket/test_request_close.py new file mode 100644 index 00000000..4e895bd5 --- /dev/null +++ b/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())