From ce828dd29447bf715c888ff1370e7684bce145f7 Mon Sep 17 00:00:00 2001 From: Shawn Snyder Date: Fri, 28 Aug 2026 13:09:55 -0500 Subject: [PATCH 1/7] connection resiliency --- intriniorealtime/equities_client.py | 108 ++++++++++++++++++++++------ intriniorealtime/options_client.py | 97 ++++++++++++++++++------- 2 files changed, 158 insertions(+), 47 deletions(-) diff --git a/intriniorealtime/equities_client.py b/intriniorealtime/equities_client.py index 7e3b30e..adc502a 100644 --- a/intriniorealtime/equities_client.py +++ b/intriniorealtime/equities_client.py @@ -11,6 +11,7 @@ from typing import Optional, Dict, Any SELF_HEAL_BACKOFFS = [10, 30, 60, 300, 600] +CONNECT_TIMEOUT_SECONDS = 30 _EMPTY_STRING = "" _NAN = float("NAN") REALTIME = "REALTIME" @@ -170,6 +171,7 @@ def __init__(self, configuration: Dict[str, Any], on_trade: Optional[callable], self.joined_channels = set() self.last_queue_warning_time = 0 self.last_self_heal_backoff = -1 + self._stop_event = threading.Event() self.quote_handler.start() def auth_url(self) -> str: @@ -227,33 +229,43 @@ def do_backoff(self): self.last_self_heal_backoff += 1 i = min(self.last_self_heal_backoff, len(SELF_HEAL_BACKOFFS) - 1) backoff = SELF_HEAL_BACKOFFS[i] - time.sleep(backoff) + self._stop_event.wait(timeout=backoff) def connect(self): - connected = False - while not connected: + self._stop_event.clear() + if self.quote_receiver is not None and self.quote_receiver.is_alive() and self.quote_receiver.enabled: + if self.ws: + try: + self.ws.close() + except Exception: + pass + return + + while not self._stop_event.is_set(): try: self.logger.info("Connecting...") self.ready = False self.joined_channels = set() - - if self.ws: - self.ws.close() - time.sleep(3) - self.refresh_token() self.refresh_websocket() - connected = True + return except Exception as e: self.logger.error(f"Cannot connect: {repr(e)}") self.do_backoff() def disconnect(self): + self._stop_event.set() self.ready = False self.joined_channels = set() + if self.quote_receiver: + self.quote_receiver.enabled = False + if self.ws: - self.ws.close() + try: + self.ws.close() + except Exception: + pass time.sleep(1) def refresh_token(self): @@ -270,6 +282,8 @@ def refresh_token(self): self.logger.info("Authentication successful!") def refresh_websocket(self): + if self.quote_receiver is not None and self.quote_receiver.is_alive(): + return self.quote_receiver = EquitiesQuoteReceiver(self) self.quote_receiver.start() @@ -368,18 +382,63 @@ def __init__(self, client): self.continuation_lock: threading.Lock = threading.Lock() def run(self): - self.client.ws = websocket.WebSocketApp( - self.client.websocket_url(), - header={HEADER_MESSAGE_FORMAT_KEY: HEADER_MESSAGE_FORMAT_VALUE, HEADER_CLIENT_INFORMATION_KEY: HEADER_CLIENT_INFORMATION_VALUE}, - on_open=self.on_open, - on_close=self.on_close, - on_message=self.on_message, - on_cont_message=self.on_cont_message, - on_error=self.on_error - ) - - self.client.logger.debug("QuoteReceiver ready") - self.client.ws.run_forever(skip_utf8_validation=True) # skip_utf8_validation for more performance + while self.enabled and not self.client._stop_event.is_set(): + handshake_event = threading.Event() + app = websocket.WebSocketApp( + self.client.websocket_url(), + header={HEADER_MESSAGE_FORMAT_KEY: HEADER_MESSAGE_FORMAT_VALUE, HEADER_CLIENT_INFORMATION_KEY: HEADER_CLIENT_INFORMATION_VALUE}, + on_open=self.on_open, + on_close=self.on_close, + on_message=self.on_message, + on_cont_message=self.on_cont_message, + on_error=self.on_error + ) + self.client.ws = app + original_on_open = app.on_open + + def on_open_with_timeout(ws, *args): + handshake_event.set() + if original_on_open: + original_on_open(ws, *args) + + app.on_open = on_open_with_timeout + + def watchdog(): + if not handshake_event.wait(CONNECT_TIMEOUT_SECONDS): + if not self.enabled or self.client._stop_event.is_set(): + return + self.client.logger.warning(f"Websocket connect timed out after {CONNECT_TIMEOUT_SECONDS}s") + try: + app.close() + except Exception: + pass + try: + if app.sock: + app.sock.close() + except Exception: + pass + + self.client.logger.debug("QuoteReceiver ready") + threading.Thread(target=watchdog, daemon=True).start() + try: + app.run_forever(skip_utf8_validation=True) # skip_utf8_validation for more performance + except Exception as e: + self.client.logger.error(f"Websocket ERROR: {repr(e)}") + + self.client.ready = False + if not self.enabled or self.client._stop_event.is_set(): + break + + self.client.logger.info("Websocket closed, reconnecting...") + self.client.do_backoff() + if not self.enabled or self.client._stop_event.is_set(): + break + try: + self.client.refresh_token() + except Exception as e: + self.client.logger.error(f"Cannot connect: {repr(e)}") + continue + self.client.logger.debug("QuoteReceiver exiting") def on_open(self, ws): @@ -392,7 +451,10 @@ def on_close(self, ws, code, message): def on_error(self, ws, error, *args): try: self.client.logger.error(f"Websocket ERROR: {error}") - self.client.connect() + try: + ws.close() + except Exception: + pass except Exception as e: self.client.logger.error(f"Error in on_error handler: {repr(e)}; {repr(error)}") raise e diff --git a/intriniorealtime/options_client.py b/intriniorealtime/options_client.py index 3815b70..5055357 100644 --- a/intriniorealtime/options_client.py +++ b/intriniorealtime/options_client.py @@ -10,6 +10,7 @@ from enum import IntEnum, unique _SELF_HEAL_BACKOFFS = [10, 30, 60, 300, 600] +_CONNECT_TIMEOUT_SECONDS = 30 _EMPTY_STRING = "" _OPTIONS_TRADE_MESSAGE_SIZE = 72 # 61 used + 11 pad _OPTIONS_QUOTE_MESSAGE_SIZE = 52 # 48 used + 4 pad @@ -33,13 +34,18 @@ def log(message: str): def do_backoff(fn: Callable[[None], bool]): i: int = 0 - backoff: int = _SELF_HEAL_BACKOFFS[i] - success: bool = fn() - while (not success): - time.sleep(backoff) + while not _stopFlag.is_set(): + try: + if fn(): + return + except Exception as e: + _log.warning("Websocket - Attempt failed: {0}".format(e)) + if _stopFlag.is_set(): + return + backoff: int = _SELF_HEAL_BACKOFFS[i] i = min(i + 1, len(_SELF_HEAL_BACKOFFS) - 1) - backoff = _SELF_HEAL_BACKOFFS[i] - success = fn() + if _stopFlag.wait(timeout=backoff): + return @unique class Providers(IntEnum): @@ -300,8 +306,25 @@ def __init__(self, self.__is_reconnecting: bool = False self.__last_reset: float = time.time() self.isReady: bool = False + self.__handshake_event: threading.Event = threading.Event() + self.__backoff_index: int = 0 + + def __close_leftover_sock(self): + try: + self.close() + except Exception: + pass + try: + sock = getattr(self, "sock", None) + if sock: + sock.close() + self.sock = None + except Exception: + pass def __on_open(self, ws): + self.__handshake_event.set() + self.__backoff_index = 0 _log.info("Websocket - Connected") self.__wsLock.acquire() try: @@ -325,26 +348,11 @@ def __on_open(self, ws): _log.info("Websocket - Joining channel: {0}".format(symbol)) self.send_binary(message) - def __try_reconnect(self) -> bool: - _log.info("Websocket - Reconnecting...") - if self.isReady: - return True - else: - with self.__wsLock: - self.__is_reconnecting = True - token: str = self.__get_token(None) - super().url = self.__get_url(token) - self.start() - return False - def __on_close(self, ws, closeStatusCode, closeMsg): + _log.info("Websocket - Closed - {0}: {1}".format(closeStatusCode, closeMsg)) self.__wsLock.acquire() try: - if (not self.__is_reconnecting): - _log.info("Websocket - Closed - {0}: {1}".format(closeStatusCode, closeMsg)) - self.isReady = False - if (not _stopFlag.is_set()): - do_backoff(self.__try_reconnect) + self.isReady = False finally: self.__wsLock.release() @@ -388,7 +396,48 @@ def __on_data(self, ws, data, code, is_last): #continueFlag - If 0, the data con _log.error("Error received: {0}".format(data)) def start(self): - super().run_forever(skip_utf8_validation=True) + while not _stopFlag.is_set(): + self.__handshake_event.clear() + self.__wsLock.acquire() + try: + self.isReady = False + finally: + self.__wsLock.release() + self.__close_leftover_sock() + + def watchdog(): + if not self.__handshake_event.wait(_CONNECT_TIMEOUT_SECONDS): + if _stopFlag.is_set(): + return + _log.warning("Websocket - Connect timed out after {0}s".format(_CONNECT_TIMEOUT_SECONDS)) + self.__close_leftover_sock() + + threading.Thread(target=watchdog, daemon=True).start() + try: + super().run_forever(skip_utf8_validation=True) + except Exception as e: + _log.warning("Websocket - Attempt failed: {0}".format(e)) + + self.__wsLock.acquire() + try: + self.isReady = False + finally: + self.__wsLock.release() + + if _stopFlag.is_set(): + return + + _log.info("Websocket - Reconnecting...") + backoff: int = _SELF_HEAL_BACKOFFS[self.__backoff_index] + self.__backoff_index = min(self.__backoff_index + 1, len(_SELF_HEAL_BACKOFFS) - 1) + if _stopFlag.wait(timeout=backoff): + return + try: + token: str = self.__get_token(None) + if token: + super().url = self.__get_url(token) + except Exception as e: + _log.warning("Websocket - Reconnect attempt failed: {0}".format(e)) # super().run_forever(ping_interval = 5, ping_timeout = 2, skip_utf8_validation = True) def stop(self): From 0adc60c4ad3c389f4b995042203475341a8f78ea Mon Sep 17 00:00:00 2001 From: Shawn Snyder Date: Mon, 21 Sep 2026 23:47:38 -0500 Subject: [PATCH 2/7] further improvements to connection resiliency --- intriniorealtime/_websocket.py | 69 ++++ intriniorealtime/equities_client.py | 119 ++++-- intriniorealtime/options_client.py | 288 ++++++++++----- tests/__init__.py | 0 tests/test_equities_client.py | 372 +++++++++++++++++++ tests/test_options_client.py | 549 ++++++++++++++++++++++++++++ tests/test_websocket_abort.py | 102 ++++++ 7 files changed, 1373 insertions(+), 126 deletions(-) create mode 100644 intriniorealtime/_websocket.py create mode 100644 tests/__init__.py create mode 100644 tests/test_equities_client.py create mode 100644 tests/test_options_client.py create mode 100644 tests/test_websocket_abort.py diff --git a/intriniorealtime/_websocket.py b/intriniorealtime/_websocket.py new file mode 100644 index 0000000..82c0403 --- /dev/null +++ b/intriniorealtime/_websocket.py @@ -0,0 +1,69 @@ +import socket +import threading +from typing import Optional + + +def abort_websocket_app(app) -> None: + if app is None: + return + try: + app.keep_running = False + except Exception: + pass + + ws = getattr(app, "sock", None) + raw = getattr(ws, "sock", None) if ws is not None else None + + if raw is not None: + try: + raw.shutdown(socket.SHUT_RDWR) + except Exception: + pass + try: + raw.close() + except Exception: + pass + + try: + if ws is not None: + ws.sock = None + ws.connected = False + except Exception: + pass + + try: + app.sock = None + except Exception: + pass + + try: + close = getattr(app, "close", None) + if callable(close): + close() + except Exception: + pass + + +def should_abort_handshake( + handshake_event: threading.Event, + stop_flag: Optional[threading.Event], + generation: int, + current_generation, + timeout_seconds: float, +) -> bool: + if handshake_event.wait(timeout_seconds): + return False + current = current_generation() if callable(current_generation) else current_generation + if generation != current: + return False + return True + + +def clear_socket_timeout(app) -> None: + sock = getattr(app, "sock", None) + if sock is None: + return + try: + sock.settimeout(None) + except Exception: + pass diff --git a/intriniorealtime/equities_client.py b/intriniorealtime/equities_client.py index adc502a..2fcdbae 100644 --- a/intriniorealtime/equities_client.py +++ b/intriniorealtime/equities_client.py @@ -10,8 +10,11 @@ from enum import IntEnum, unique from typing import Optional, Dict, Any +from ._websocket import abort_websocket_app, clear_socket_timeout, should_abort_handshake + SELF_HEAL_BACKOFFS = [10, 30, 60, 300, 600] CONNECT_TIMEOUT_SECONDS = 30 +STABLE_SESSION_SECONDS = 1.0 _EMPTY_STRING = "" _NAN = float("NAN") REALTIME = "REALTIME" @@ -234,19 +237,17 @@ def do_backoff(self): def connect(self): self._stop_event.clear() if self.quote_receiver is not None and self.quote_receiver.is_alive() and self.quote_receiver.enabled: - if self.ws: - try: - self.ws.close() - except Exception: - pass return + self.last_self_heal_backoff = -1 while not self._stop_event.is_set(): try: self.logger.info("Connecting...") self.ready = False self.joined_channels = set() self.refresh_token() + if self._stop_event.is_set(): + return self.refresh_websocket() return except Exception as e: @@ -261,19 +262,25 @@ def disconnect(self): if self.quote_receiver: self.quote_receiver.enabled = False - if self.ws: - try: - self.ws.close() - except Exception: - pass - time.sleep(1) + abort_websocket_app(self.ws) + + receiver = self.quote_receiver + if receiver is not None: + receiver.join(timeout=CONNECT_TIMEOUT_SECONDS + 1) + if receiver.is_alive(): + abort_websocket_app(self.ws) + else: + self.quote_receiver = None + self.ws = None + else: + self.ws = None def refresh_token(self): headers = {HEADER_CLIENT_INFORMATION_KEY: HEADER_CLIENT_INFORMATION_VALUE} if self.api_key: - response = requests.get(self.auth_url(), headers=headers) + response = requests.get(self.auth_url(), headers=headers, timeout=CONNECT_TIMEOUT_SECONDS) else: - response = requests.get(self.auth_url(), auth=(self.username, self.password), headers=headers) + response = requests.get(self.auth_url(), auth=(self.username, self.password), headers=headers, timeout=CONNECT_TIMEOUT_SECONDS) if response.status_code != 200: raise RuntimeError("Auth failed") @@ -282,14 +289,24 @@ def refresh_token(self): self.logger.info("Authentication successful!") def refresh_websocket(self): - if self.quote_receiver is not None and self.quote_receiver.is_alive(): - return + receiver = self.quote_receiver + if receiver is not None: + if receiver.is_alive(): + if receiver.enabled: + return + abort_websocket_app(self.ws) + receiver.join(timeout=CONNECT_TIMEOUT_SECONDS + 1) + if receiver.is_alive(): + abort_websocket_app(self.ws) + raise RuntimeError("Previous quote receiver did not exit") + self.quote_receiver = None + self.ws = None self.quote_receiver = EquitiesQuoteReceiver(self) self.quote_receiver.start() def on_connect(self): self.ready = True - self.last_self_heal_backoff = -1 + self.joined_channels = set() self.refresh_channels() def on_queue_full(self): @@ -381,8 +398,20 @@ def __init__(self, client): self.continuation_queue = queue.Queue(100) self.continuation_lock: threading.Lock = threading.Lock() + def _clear_continuation_state(self): + with self.continuation_lock: + while True: + try: + self.continuation_queue.get_nowait() + except queue.Empty: + break + def run(self): + websocket.setdefaulttimeout(CONNECT_TIMEOUT_SECONDS) + generation = 0 while self.enabled and not self.client._stop_event.is_set(): + generation += 1 + self._clear_continuation_state() handshake_event = threading.Event() app = websocket.WebSocketApp( self.client.websocket_url(), @@ -396,30 +425,36 @@ def run(self): self.client.ws = app original_on_open = app.on_open + session_opened = False + session_opened_at = 0.0 + def on_open_with_timeout(ws, *args): + nonlocal session_opened, session_opened_at + session_opened = True + session_opened_at = time.time() handshake_event.set() + clear_socket_timeout(app) if original_on_open: original_on_open(ws, *args) app.on_open = on_open_with_timeout - def watchdog(): - if not handshake_event.wait(CONNECT_TIMEOUT_SECONDS): - if not self.enabled or self.client._stop_event.is_set(): - return + def watchdog(this_generation=generation, handshake_event=handshake_event): + if should_abort_handshake( + handshake_event, + self.client._stop_event, + this_generation, + lambda: generation, + CONNECT_TIMEOUT_SECONDS, + ): self.client.logger.warning(f"Websocket connect timed out after {CONNECT_TIMEOUT_SECONDS}s") - try: - app.close() - except Exception: - pass - try: - if app.sock: - app.sock.close() - except Exception: - pass + abort_websocket_app(app) self.client.logger.debug("QuoteReceiver ready") threading.Thread(target=watchdog, daemon=True).start() + if not self.enabled or self.client._stop_event.is_set(): + abort_websocket_app(app) + break try: app.run_forever(skip_utf8_validation=True) # skip_utf8_validation for more performance except Exception as e: @@ -430,14 +465,24 @@ def watchdog(): break self.client.logger.info("Websocket closed, reconnecting...") - self.client.do_backoff() - if not self.enabled or self.client._stop_event.is_set(): - break - try: - self.client.refresh_token() - except Exception as e: - self.client.logger.error(f"Cannot connect: {repr(e)}") - continue + skip_wait = ( + session_opened + and (time.time() - session_opened_at) >= STABLE_SESSION_SECONDS + ) + if skip_wait: + self.client.last_self_heal_backoff = -1 + while self.enabled and not self.client._stop_event.is_set(): + if skip_wait: + skip_wait = False + else: + self.client.do_backoff() + if not self.enabled or self.client._stop_event.is_set(): + break + try: + self.client.refresh_token() + break + except Exception as e: + self.client.logger.error(f"Cannot connect: {repr(e)}") self.client.logger.debug("QuoteReceiver exiting") diff --git a/intriniorealtime/options_client.py b/intriniorealtime/options_client.py index 5055357..17c9832 100644 --- a/intriniorealtime/options_client.py +++ b/intriniorealtime/options_client.py @@ -9,8 +9,11 @@ from collections.abc import Callable from enum import IntEnum, unique +from ._websocket import abort_websocket_app, clear_socket_timeout, should_abort_handshake + _SELF_HEAL_BACKOFFS = [10, 30, 60, 300, 600] _CONNECT_TIMEOUT_SECONDS = 30 +_STABLE_SESSION_SECONDS = 1.0 _EMPTY_STRING = "" _OPTIONS_TRADE_MESSAGE_SIZE = 72 # 61 used + 11 pad _OPTIONS_QUOTE_MESSAGE_SIZE = 52 # 48 used + 4 pad @@ -32,19 +35,20 @@ def log(message: str): _log.info(message) -def do_backoff(fn: Callable[[None], bool]): +def do_backoff(fn: Callable[[], bool], stop_flag: threading.Event = None): + flag: threading.Event = stop_flag if stop_flag is not None else _stopFlag i: int = 0 - while not _stopFlag.is_set(): + while not flag.is_set(): try: if fn(): return except Exception as e: _log.warning("Websocket - Attempt failed: {0}".format(e)) - if _stopFlag.is_set(): + if flag.is_set(): return backoff: int = _SELF_HEAL_BACKOFFS[i] i = min(i + 1, len(_SELF_HEAL_BACKOFFS) - 1) - if _stopFlag.wait(timeout=backoff): + if flag.wait(timeout=backoff): return @unique @@ -281,22 +285,23 @@ def __init__(self, ws_url: str, ws_lock: threading.Lock, worker_threads: list[threading.Thread], - get_channels: Callable[[None], set[tuple[str, bool]]], - get_token: Callable[[None], str], + get_channels: Callable[[], set], + get_token: Callable[[], str], get_url: Callable[[str], str], use_on_trade: bool, use_on_quote: bool, use_on_refresh: bool, use_on_ua: bool, - data_queue: queue.Queue): + data_queue: queue.Queue, + stop_flag: threading.Event): super().__init__(ws_url, on_open=self.__on_open, on_close=self.__on_close, on_data=self.__on_data, on_error=self.__on_error) self.__wsLock: threading.Lock = ws_lock self.__continuation_queue = queue.Queue(100) self.__continuation_lock: threading.Lock = threading.Lock() self.__currently_continuing: bool = False self.__worker_threads: list[threading.Thread] = worker_threads - self.__get_channels: Callable[[None], set[tuple[str, bool]]] = get_channels - self.__get_token: Callable[[None], str] = get_token + self.__get_channels: Callable[[], set] = get_channels + self.__get_token: Callable[[], str] = get_token self.__get_url: Callable[[str], str] = get_url self.__use_on_trade: bool = use_on_trade self.__use_on_quote: bool = use_on_quote @@ -308,23 +313,28 @@ def __init__(self, self.isReady: bool = False self.__handshake_event: threading.Event = threading.Event() self.__backoff_index: int = 0 + self.__connect_generation: int = 0 + self.__session_opened: bool = False + self.__session_opened_at: float = 0.0 + self.__stop_flag: threading.Event = stop_flag def __close_leftover_sock(self): - try: - self.close() - except Exception: - pass - try: - sock = getattr(self, "sock", None) - if sock: - sock.close() - self.sock = None - except Exception: - pass + abort_websocket_app(self) + + def __clear_continuation_state(self): + with self.__continuation_lock: + self.__currently_continuing = False + while True: + try: + self.__continuation_queue.get_nowait() + except queue.Empty: + break def __on_open(self, ws): self.__handshake_event.set() - self.__backoff_index = 0 + self.__session_opened = True + self.__session_opened_at = time.time() + clear_socket_timeout(self) _log.info("Websocket - Connected") self.__wsLock.acquire() try: @@ -396,8 +406,14 @@ def __on_data(self, ws, data, code, is_last): #continueFlag - If 0, the data con _log.error("Error received: {0}".format(data)) def start(self): - while not _stopFlag.is_set(): - self.__handshake_event.clear() + websocket.setdefaulttimeout(_CONNECT_TIMEOUT_SECONDS) + while not self.__stop_flag.is_set(): + handshake_event = threading.Event() + self.__handshake_event = handshake_event + self.__connect_generation += 1 + generation = self.__connect_generation + self.__session_opened = False + self.__clear_continuation_state() self.__wsLock.acquire() try: self.isReady = False @@ -405,16 +421,23 @@ def start(self): self.__wsLock.release() self.__close_leftover_sock() - def watchdog(): - if not self.__handshake_event.wait(_CONNECT_TIMEOUT_SECONDS): - if _stopFlag.is_set(): - return + def watchdog(this_generation=generation, handshake_event=handshake_event): + if should_abort_handshake( + handshake_event, + self.__stop_flag, + this_generation, + lambda: self.__connect_generation, + _CONNECT_TIMEOUT_SECONDS, + ): _log.warning("Websocket - Connect timed out after {0}s".format(_CONNECT_TIMEOUT_SECONDS)) self.__close_leftover_sock() threading.Thread(target=watchdog, daemon=True).start() + if self.__stop_flag.is_set(): + abort_websocket_app(self) + return try: - super().run_forever(skip_utf8_validation=True) + self.run_forever(skip_utf8_validation=True) except Exception as e: _log.warning("Websocket - Attempt failed: {0}".format(e)) @@ -424,24 +447,40 @@ def watchdog(): finally: self.__wsLock.release() - if _stopFlag.is_set(): + if self.__stop_flag.is_set(): return _log.info("Websocket - Reconnecting...") - backoff: int = _SELF_HEAL_BACKOFFS[self.__backoff_index] - self.__backoff_index = min(self.__backoff_index + 1, len(_SELF_HEAL_BACKOFFS) - 1) - if _stopFlag.wait(timeout=backoff): + skip_wait: bool = ( + self.__session_opened + and (time.time() - self.__session_opened_at) >= _STABLE_SESSION_SECONDS + ) + if skip_wait: + self.__backoff_index = 0 + token_ready: bool = False + while not self.__stop_flag.is_set(): + if skip_wait: + skip_wait = False + else: + backoff: int = _SELF_HEAL_BACKOFFS[self.__backoff_index] + self.__backoff_index = min(self.__backoff_index + 1, len(_SELF_HEAL_BACKOFFS) - 1) + if self.__stop_flag.wait(timeout=backoff): + return + try: + token: str = self.__get_token() + if token: + self.url = self.__get_url(token) + token_ready = True + break + _log.warning("Websocket - Reconnect attempt failed: no token") + except Exception as e: + _log.warning("Websocket - Reconnect attempt failed: {0}".format(e)) + if not token_ready: return - try: - token: str = self.__get_token(None) - if token: - super().url = self.__get_url(token) - except Exception as e: - _log.warning("Websocket - Reconnect attempt failed: {0}".format(e)) # super().run_forever(ping_interval = 5, ping_timeout = 2, skip_utf8_validation = True) def stop(self): - super().close() + abort_websocket_app(self) def send(self, message: str): super().send(message, websocket.ABNF.OPCODE_TEXT) @@ -549,14 +588,16 @@ def _thread_fn(index: int, data: queue.Queue, on_trade: Callable[[OptionsTrade], None], on_quote: Callable[[OptionsQuote], None] = None, on_refresh: Callable[[OptionsRefresh], None] = None, - on_unusual_activity: Callable[[OptionsUnusualActivity], None] = None): + on_unusual_activity: Callable[[OptionsUnusualActivity], None] = None, + stop_flag: threading.Event = None): + flag: threading.Event = stop_flag if stop_flag is not None else _stopFlag _log.debug("Starting worker thread {0}".format(index)) datum: bytes = None count: int = 0 start_index: int = 1 msg_type: int = 0 message: bytes = None - while not _stopFlag.is_set(): + while not flag.is_set(): try: datum = data.get(True, 1.0) count = datum[0] @@ -731,16 +772,26 @@ def __init__(self, config: Config, on_trade: Callable[[OptionsTrade], None], on_ self.__data: queue.Queue = queue.Queue() self.__t_lock: threading.Lock = threading.Lock() self.__ws_lock: threading.Lock = threading.Lock() - self.__worker_threads: list[threading.Thread] = [threading.Thread( + self.__on_trade = on_trade + self.__on_quote = on_quote + self.__on_refresh = on_refresh + self.__on_unusual_activity = on_unusual_activity + self.__num_threads: int = config.num_threads + self._stop_event: threading.Event = threading.Event() + self._stop_generation: int = 0 + self.__worker_threads: list[threading.Thread] = self.__make_worker_threads() + self.__socket_thread: threading.Thread = None + self.__is_started: bool = False + _log.setLevel(config.log_level) + + def __make_worker_threads(self) -> list[threading.Thread]: + return [threading.Thread( group=None, target=_thread_fn, - args=(i, self.__data, on_trade, on_quote, on_refresh, on_unusual_activity), + args=(i, self.__data, self.__on_trade, self.__on_quote, self.__on_refresh, self.__on_unusual_activity, self._stop_event), kwargs={}, daemon=True - ) for i in range(config.num_threads)] - self.__socket_thread: threading.Thread = None - self.__is_started: bool = False - _log.setLevel(config.log_level) + ) for i in range(self.__num_threads)] def __all_ready(self) -> bool: self.__ws_lock.acquire() @@ -799,7 +850,7 @@ def __get_token(self) -> str: self.__t_lock.acquire() try: if ((time.time() - self.__token[1]) > (60 * 60 * 24)): # 60sec/min * 60min/hr * 24hrs = 1 day - do_backoff(self.__try_set_token) + do_backoff(self.__try_set_token, self._stop_event) return self.__token[0] finally: self.__t_lock.release() @@ -816,22 +867,28 @@ def __join(self, symbol: str): message[0] = 74 # join code message[1] = _get_option_mask(self.__use_on_trade, self.__use_on_quote, self.__use_on_refresh, self.__use_on_unusual_activity) message[2:] = symbol_bytes - if self.__webSocket.isReady: + ws = self.__webSocket + if ws is not None and ws.isReady: _log.info("Websocket - Joining channel: {0}".format(transformed_symbol)) - self.__webSocket.send_binary(message) + ws.send_binary(message) - def __leave(self, symbol: str): + def __leave(self, symbol: str, forget: bool = True): transformed_symbol: str = _transform_contract_to_new(symbol) - if transformed_symbol in self.__channels: - self.__channels.remove(transformed_symbol) - symbol_bytes = bytes(transformed_symbol, 'utf-8') - message: bytes = bytearray(len(symbol_bytes) + 2) - message[0] = 76 # leave code - message[1] = _get_option_mask(self.__use_on_trade, self.__use_on_quote, self.__use_on_refresh, self.__use_on_unusual_activity) - message[2:] = symbol_bytes - if self.__webSocket.isReady: + if transformed_symbol not in self.__channels: + return + symbol_bytes = bytes(transformed_symbol, 'utf-8') + message: bytes = bytearray(len(symbol_bytes) + 2) + message[0] = 76 # leave code + message[1] = _get_option_mask(self.__use_on_trade, self.__use_on_quote, self.__use_on_refresh, self.__use_on_unusual_activity) + message[2:] = symbol_bytes + try: + ws = self.__webSocket + if ws is not None and ws.isReady: _log.info("Websocket - Leaving channel: {0}".format(transformed_symbol)) - self.__webSocket.send_binary(message) + ws.send_binary(message) + finally: + if forget: + self.__channels.discard(transformed_symbol) def join(self, *symbols): if self.__is_started: @@ -864,6 +921,9 @@ def leave_firehose(self): self.__leave("$FIREHOSE") def __socket_start_fn(self, token: str): + if not token: + _log.error("Websocket - Missing token") + return _log.info("Websocket - Connecting...") ws_url: str = self.__get_web_socket_url(token) self.__webSocket = _WebSocket(ws_url, @@ -876,45 +936,95 @@ def __socket_start_fn(self, token: str): self.__use_on_quote, self.__use_on_refresh, self.__use_on_unusual_activity, - self.__data) + self.__data, + self._stop_event) self.__webSocket.start() def start(self): if (not (self.__use_on_trade or self.__use_on_quote or self.__use_on_refresh or self.__use_on_unusual_activity)): raise ValueError("You must set at least one callback method before starting client") - token: str = self.__get_token() - self.__ws_lock.acquire() - try: - self.__socket_thread = threading.Thread( - group=None, - target=self.__socket_start_fn, - args=(token,), - kwargs={}, - daemon=True - ) - finally: - self.__ws_lock.release() - self.__socket_thread.start() - self.__is_started = True + if (self.__is_started + and self.__socket_thread is not None + and self.__socket_thread.is_alive() + and not self._stop_event.is_set()): + return + stop_generation = self._stop_generation + while True: + leftover = self.__socket_thread + if leftover is not None and leftover.is_alive(): + abort_websocket_app(self.__webSocket) + leftover.join(timeout=_CONNECT_TIMEOUT_SECONDS + 1) + if leftover.is_alive(): + if self._stop_generation != stop_generation: + return + continue + if self._stop_generation != stop_generation: + return + self._stop_event = threading.Event() + self.__worker_threads = self.__make_worker_threads() + token: str = self.__get_token() + if self._stop_event.is_set(): + return + if not token: + if self._stop_event.wait(timeout=_SELF_HEAL_BACKOFFS[0]): + return + continue + self.__ws_lock.acquire() + try: + self.__socket_thread = threading.Thread( + group=None, + target=self.__socket_start_fn, + args=(token,), + kwargs={}, + daemon=True + ) + finally: + self.__ws_lock.release() + self.__socket_thread.start() + self.__is_started = True + return def stop(self): _log.info("Stopping...") - if len(self.__channels) > 0: - self.leave() - time.sleep(1.0) - self.__ws_lock.acquire() + self._stop_generation += 1 + self._stop_event.set() try: - self.__webSocket.isReady = False + ws = self.__webSocket + if ws is not None: + try: + sock = getattr(ws, "sock", None) + if sock is not None: + sock.settimeout(1) + except Exception: + pass + if len(self.__channels) > 0: + channels: set[str] = self.__channels.copy() + for symbol in channels: + try: + self.__leave(symbol, forget=False) + except Exception as e: + _log.warning("Websocket - Leave failed: {0}".format(e)) + break + time.sleep(1.0) finally: - self.__ws_lock.release() - _stopFlag.set() - self.__webSocket.stop() - for i in range(len(self.__worker_threads)): - self.__worker_threads[i].join() - _log.debug("Worker thread {0} joined".format(i)) - self.__socket_thread.join() - _log.debug("Socket thread joined") - _log.info("Stopped") + self.__ws_lock.acquire() + try: + if self.__webSocket is not None: + self.__webSocket.isReady = False + finally: + self.__ws_lock.release() + if self.__webSocket is not None: + self.__webSocket.stop() + for i in range(len(self.__worker_threads)): + worker = self.__worker_threads[i] + if worker.ident is not None: + worker.join(timeout=_CONNECT_TIMEOUT_SECONDS + 1) + _log.debug("Worker thread {0} joined".format(i)) + if self.__socket_thread is not None and self.__socket_thread.ident is not None: + self.__socket_thread.join(timeout=_CONNECT_TIMEOUT_SECONDS + 1) + _log.debug("Socket thread joined") + self.__is_started = False + _log.info("Stopped") def get_stats(self) -> tuple[int, int, int]: return _dataMsgCount, _txtMsgCount, self.__data.qsize() diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_equities_client.py b/tests/test_equities_client.py new file mode 100644 index 0000000..bc038ed --- /dev/null +++ b/tests/test_equities_client.py @@ -0,0 +1,372 @@ +import logging +import time +import unittest +from unittest.mock import Mock, patch + +from intriniorealtime.equities_client import ( + CONNECT_TIMEOUT_SECONDS, + EquitiesQuoteHandler, + EquitiesQuoteReceiver, + IntrinioRealtimeEquitiesClient, +) + + +def _make_client(): + with patch.object(EquitiesQuoteHandler, "start"): + return IntrinioRealtimeEquitiesClient( + {"api_key": "test-key", "provider": "IEX"}, + lambda *args: None, + lambda *args: None, + ) + + +class EquitiesReconnectTests(unittest.TestCase): + def setUp(self): + logging.disable(logging.CRITICAL) + + def tearDown(self): + logging.disable(logging.NOTSET) + + def test_on_connect_rejoins_channels_after_reconnect(self): + client = _make_client() + client.channels = {"AAPL"} + client.joined_channels = {"AAPL"} + client.ready = False + client.ws = Mock() + + client.on_connect() + + self.assertTrue(client.ready) + self.assertEqual(client.joined_channels, {"AAPL"}) + client.ws.send.assert_called() + + def test_on_connect_does_not_skip_join_when_joined_channels_are_stale(self): + client = _make_client() + client.channels = {"AAPL", "MSFT"} + client.joined_channels = {"AAPL", "MSFT"} + client.ready = True + client.ws = Mock() + + client.refresh_channels() + client.ws.send.assert_not_called() + + client.on_connect() + self.assertEqual(client.ws.send.call_count, 2) + + def test_watchdog_aborts_app_when_handshake_times_out(self): + client = _make_client() + client.token = "tok" + receiver = EquitiesQuoteReceiver(client) + aborted = [] + + def fake_run_forever(app, *args, **kwargs): + time.sleep(0.12) + receiver.enabled = False + client._stop_event.set() + + with patch("intriniorealtime.equities_client.abort_websocket_app", side_effect=lambda app: aborted.append(app)), \ + patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("websocket.setdefaulttimeout"), \ + patch("websocket.WebSocketApp.run_forever", fake_run_forever): + receiver.run() + + self.assertTrue(aborted) + + def test_run_sets_connect_timeout_and_clears_it_on_open(self): + client = _make_client() + client.token = "tok" + receiver = EquitiesQuoteReceiver(client) + sock = Mock() + + def fake_run_forever(app, *args, **kwargs): + app.sock = sock + app.on_open(app) + receiver.enabled = False + client._stop_event.set() + + with patch("websocket.setdefaulttimeout") as set_timeout, \ + patch("websocket.WebSocketApp.run_forever", fake_run_forever): + receiver.run() + + set_timeout.assert_called_with(CONNECT_TIMEOUT_SECONDS) + sock.settimeout.assert_called_with(None) + + def test_refresh_token_failure_backs_off_and_skips_socket_until_auth_succeeds(self): + client = _make_client() + client.token = "stale" + client.do_backoff = Mock() + attempts = {"n": 0} + + def refresh(): + attempts["n"] += 1 + if attempts["n"] < 3: + raise RuntimeError("Auth failed") + + client.refresh_token = Mock(side_effect=refresh) + receiver = EquitiesQuoteReceiver(client) + run_forever_calls = [] + + def fake_run_forever(app, *args, **kwargs): + run_forever_calls.append(1) + if len(run_forever_calls) >= 2: + receiver.enabled = False + client._stop_event.set() + + with patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("websocket.WebSocketApp.run_forever", fake_run_forever), \ + patch("intriniorealtime.equities_client.abort_websocket_app"): + receiver.run() + + self.assertEqual(len(run_forever_calls), 2) + self.assertEqual(client.refresh_token.call_count, 3) + self.assertGreaterEqual(client.do_backoff.call_count, 3) + + def test_disconnect_then_connect_starts_a_new_receiver(self): + client = _make_client() + client.refresh_token = Mock() + old = Mock() + old.is_alive.return_value = True + old.enabled = True + + def join_and_exit(timeout=None): + old.is_alive.return_value = False + + old.join.side_effect = join_and_exit + client.quote_receiver = old + client.ws = Mock() + + with patch("intriniorealtime.equities_client.abort_websocket_app"): + client.disconnect() + + self.assertFalse(old.enabled) + old.join.assert_called() + self.assertIsNone(client.quote_receiver) + self.assertTrue(client._stop_event.is_set()) + + new_receiver = Mock() + with patch("intriniorealtime.equities_client.EquitiesQuoteReceiver", return_value=new_receiver) as ctor: + client.connect() + + ctor.assert_called_once() + new_receiver.start.assert_called_once() + self.assertIs(client.quote_receiver, new_receiver) + self.assertFalse(client._stop_event.is_set()) + + def test_refresh_websocket_replaces_disabled_live_receiver(self): + client = _make_client() + old = Mock() + old.is_alive.return_value = True + old.enabled = False + + def join_and_exit(timeout=None): + old.is_alive.return_value = False + + old.join.side_effect = join_and_exit + client.quote_receiver = old + new_receiver = Mock() + + with patch("intriniorealtime.equities_client.EquitiesQuoteReceiver", return_value=new_receiver): + client.refresh_websocket() + + old.join.assert_called() + new_receiver.start.assert_called_once() + self.assertIs(client.quote_receiver, new_receiver) + + def test_connect_on_healthy_session_does_not_close_socket(self): + client = _make_client() + client.ready = True + client.refresh_token = Mock() + receiver = Mock() + receiver.is_alive.return_value = True + receiver.enabled = True + client.quote_receiver = receiver + ws = Mock() + client.ws = ws + + client.connect() + + ws.close.assert_not_called() + client.refresh_token.assert_not_called() + self.assertIs(client.quote_receiver, receiver) + self.assertTrue(client.ready) + self.assertTrue(receiver.enabled) + + def test_connect_while_receiver_is_starting_does_not_close_socket(self): + client = _make_client() + client.ready = False + client.refresh_token = Mock() + receiver = Mock() + receiver.is_alive.return_value = True + receiver.enabled = True + client.quote_receiver = receiver + ws = Mock() + client.ws = ws + + client.connect() + + ws.close.assert_not_called() + client.refresh_token.assert_not_called() + self.assertIs(client.quote_receiver, receiver) + + def test_connect_resets_backoff_on_new_attempt(self): + client = _make_client() + client.last_self_heal_backoff = 4 + client.refresh_token = Mock() + new_receiver = Mock() + with patch("intriniorealtime.equities_client.EquitiesQuoteReceiver", return_value=new_receiver): + client.connect() + self.assertEqual(client.last_self_heal_backoff, -1) + new_receiver.start.assert_called_once() + + def test_connect_healthy_noop_does_not_reset_backoff(self): + client = _make_client() + client.last_self_heal_backoff = 4 + receiver = Mock() + receiver.is_alive.return_value = True + receiver.enabled = True + client.quote_receiver = receiver + client.connect() + self.assertEqual(client.last_self_heal_backoff, 4) + + def test_receiver_clears_continuation_queue_on_each_connect_attempt(self): + client = _make_client() + client.token = "tok" + client.do_backoff = Mock() + client.refresh_token = Mock() + receiver = EquitiesQuoteReceiver(client) + receiver.continuation_queue.put(b"partial-from-prior-session") + seen_empty = [] + run_count = {"n": 0} + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + seen_empty.append(receiver.continuation_queue.empty()) + if run_count["n"] == 1: + receiver.continuation_queue.put(b"leftover-mid-fragment") + return + receiver.enabled = False + client._stop_event.set() + + with patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("websocket.WebSocketApp.run_forever", fake_run_forever), \ + patch("intriniorealtime.equities_client.abort_websocket_app"): + receiver.run() + + self.assertEqual(run_count["n"], 2) + self.assertEqual(seen_empty, [True, True]) + + def test_refresh_token_passes_timeout(self): + client = _make_client() + with patch("intriniorealtime.equities_client.requests.get") as get: + get.return_value.status_code = 200 + get.return_value.text = "tok" + client.refresh_token() + self.assertEqual(get.call_args.kwargs["timeout"], CONNECT_TIMEOUT_SECONDS) + + def test_connect_does_not_start_receiver_if_stopped_during_auth(self): + client = _make_client() + + def refresh(): + client._stop_event.set() + + client.refresh_token = Mock(side_effect=refresh) + with patch("intriniorealtime.equities_client.EquitiesQuoteReceiver") as ctor: + client.connect() + ctor.assert_not_called() + + def test_disconnect_keeps_receiver_if_join_times_out(self): + client = _make_client() + old = Mock() + old.is_alive.return_value = True + old.enabled = True + old.join = Mock() + client.quote_receiver = old + client.ws = Mock() + + ws = client.ws + with patch("intriniorealtime.equities_client.abort_websocket_app") as abort: + client.disconnect() + + self.assertFalse(old.enabled) + old.join.assert_called() + self.assertIs(client.quote_receiver, old) + self.assertIs(client.ws, ws) + self.assertGreaterEqual(abort.call_count, 2) + + def test_refresh_websocket_does_not_replace_receiver_still_alive_after_join(self): + client = _make_client() + old = Mock() + old.is_alive.return_value = True + old.enabled = False + old.join = Mock() + client.quote_receiver = old + client.ws = Mock() + + with patch("intriniorealtime.equities_client.abort_websocket_app") as abort, \ + patch("intriniorealtime.equities_client.EquitiesQuoteReceiver") as ctor: + with self.assertRaises(RuntimeError): + client.refresh_websocket() + + ctor.assert_not_called() + self.assertIs(client.quote_receiver, old) + self.assertGreaterEqual(abort.call_count, 2) + + def test_open_then_immediate_drop_backs_off(self): + client = _make_client() + client.token = "tok" + client.do_backoff = Mock() + client.refresh_token = Mock() + receiver = EquitiesQuoteReceiver(client) + run_count = {"n": 0} + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + if run_count["n"] == 1: + app.on_open(app) + return + receiver.enabled = False + client._stop_event.set() + + with patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("intriniorealtime.equities_client.STABLE_SESSION_SECONDS", 10), \ + patch("websocket.WebSocketApp.run_forever", fake_run_forever), \ + patch("intriniorealtime.equities_client.abort_websocket_app"): + receiver.run() + + self.assertEqual(run_count["n"], 2) + client.do_backoff.assert_called() + + def test_stable_session_drop_skips_first_backoff(self): + client = _make_client() + client.token = "tok" + client.do_backoff = Mock() + client.refresh_token = Mock() + receiver = EquitiesQuoteReceiver(client) + run_count = {"n": 0} + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + if run_count["n"] == 1: + app.on_open(app) + time.sleep(0.06) + return + receiver.enabled = False + client._stop_event.set() + + with patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("intriniorealtime.equities_client.STABLE_SESSION_SECONDS", 0.05), \ + patch("websocket.WebSocketApp.run_forever", fake_run_forever), \ + patch("intriniorealtime.equities_client.abort_websocket_app"): + receiver.run() + + self.assertEqual(run_count["n"], 2) + self.assertEqual(client.refresh_token.call_count, 1) + client.do_backoff.assert_not_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_options_client.py b/tests/test_options_client.py new file mode 100644 index 0000000..5e8b9cd --- /dev/null +++ b/tests/test_options_client.py @@ -0,0 +1,549 @@ +import logging +import queue +import threading +import time +import unittest +from unittest.mock import Mock, patch + +from intriniorealtime.options_client import ( + Config, + IntrinioRealtimeOptionsClient, + Providers, + _CONNECT_TIMEOUT_SECONDS, + _WebSocket, +) + + +def _make_ws(stop_flag, get_token=None, get_url=None, channels=None): + return _WebSocket( + "ws://example/old", + threading.Lock(), + [], + lambda: channels if channels is not None else set(), + get_token or (lambda: "tok"), + get_url or (lambda token: "ws://example/" + token), + True, + False, + False, + False, + queue.Queue(), + stop_flag, + ) + + +def _make_client(): + config = Config( + api_key="test-key", + provider=Providers.MANUAL, + num_threads=1, + manual_ip_address="127.0.0.1", + ) + return IntrinioRealtimeOptionsClient(config, on_trade=lambda trade: None) + + +class OptionsReconnectTests(unittest.TestCase): + def setUp(self): + logging.disable(logging.CRITICAL) + + def tearDown(self): + logging.disable(logging.NOTSET) + + def test_stale_watchdog_does_not_abort_later_attempt(self): + stop = threading.Event() + second_opened = threading.Event() + aborts_after_second_open = [] + run_count = {"n": 0} + + ws = _make_ws(stop) + + def track_abort(app): + if app is ws and second_opened.is_set(): + aborts_after_second_open.append(1) + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + if run_count["n"] == 1: + return + app.on_open(app) + second_opened.set() + time.sleep(0.2) + stop.set() + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0]), \ + patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.options_client.abort_websocket_app", side_effect=track_abort), \ + patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): + ws.start() + + self.assertGreaterEqual(run_count["n"], 2) + self.assertEqual(aborts_after_second_open, []) + + def test_reconnect_refreshes_token_without_extra_argument(self): + stop = threading.Event() + + def real_get_token(): + return "tok2" + + get_token = Mock(side_effect=lambda *args, **kwargs: real_get_token(*args, **kwargs)) + get_url = Mock(return_value="ws://example/new") + run_count = {"n": 0} + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + if run_count["n"] >= 2: + stop.set() + + ws = _make_ws(stop, get_token=get_token, get_url=get_url) + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0]), \ + patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.options_client.abort_websocket_app"), \ + patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): + ws.start() + + get_token.assert_called_with() + get_url.assert_called_with("tok2") + self.assertEqual(ws.url, "ws://example/new") + + def test_reconnect_does_not_run_forever_with_stale_url_when_token_refresh_fails(self): + stop = threading.Event() + token_calls = {"n": 0} + + def get_token(): + token_calls["n"] += 1 + if token_calls["n"] < 2: + raise RuntimeError("token failed") + return "tok2" + + urls = [] + run_count = {"n": 0} + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + urls.append(app.url) + if run_count["n"] >= 2: + stop.set() + + ws = _make_ws(stop, get_token=get_token, get_url=lambda token: "ws://example/" + token) + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0]), \ + patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.options_client.abort_websocket_app"), \ + patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): + ws.start() + + self.assertEqual(run_count["n"], 2) + self.assertEqual(urls[0], "ws://example/old") + self.assertEqual(urls[1], "ws://example/tok2") + self.assertEqual(token_calls["n"], 2) + + def test_start_after_stop_clears_stop_event_and_connects_again(self): + client = _make_client() + start_calls = [] + + def fake_ws_start(self): + start_calls.append(1) + self.isReady = True + client._stop_event.wait(5) + + with patch("intriniorealtime.options_client.requests.get") as get, \ + patch("intriniorealtime.options_client.time.sleep"), \ + patch.object(_WebSocket, "start", fake_ws_start): + get.return_value.status_code = 200 + get.return_value.text = "tok" + + client.start() + deadline = time.time() + 2 + while not start_calls and time.time() < deadline: + time.sleep(0.01) + self.assertEqual(start_calls, [1]) + + client.stop() + self.assertTrue(client._stop_event.is_set()) + + client.start() + deadline = time.time() + 2 + while len(start_calls) < 2 and time.time() < deadline: + time.sleep(0.01) + self.assertEqual(len(start_calls), 2) + + client.stop() + + def test_websocket_start_is_noop_while_stopped_and_runs_after_clear(self): + stop = threading.Event() + stop.set() + ws = _make_ws(stop) + with patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever") as run_forever, \ + patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("websocket.setdefaulttimeout"): + ws.start() + run_forever.assert_not_called() + + stop.clear() + calls = [] + + def fake_run_forever(app, *args, **kwargs): + calls.append(1) + stop.set() + + with patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever), \ + patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.options_client.abort_websocket_app"): + ws.start() + self.assertEqual(calls, [1]) + + def test_clears_continuation_state_on_each_connect_attempt(self): + stop = threading.Event() + ws = _make_ws(stop) + ws._WebSocket__currently_continuing = True + ws._WebSocket__continuation_queue.put(b"partial-from-prior-session") + seen = [] + run_count = {"n": 0} + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + seen.append(( + app._WebSocket__currently_continuing, + app._WebSocket__continuation_queue.empty(), + )) + if run_count["n"] == 1: + app._WebSocket__currently_continuing = True + app._WebSocket__continuation_queue.put(b"leftover-mid-fragment") + return + stop.set() + + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0]), \ + patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.options_client.abort_websocket_app"), \ + patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): + ws.start() + + self.assertEqual(run_count["n"], 2) + self.assertEqual(seen, [(False, True), (False, True)]) + + def test_stop_preserves_channels_and_clears_started(self): + client = _make_client() + client._IntrinioRealtimeOptionsClient__channels = {"AAPL"} + client._IntrinioRealtimeOptionsClient__is_started = True + ws = Mock() + ws.isReady = True + client._IntrinioRealtimeOptionsClient__webSocket = ws + client._IntrinioRealtimeOptionsClient__socket_thread = None + + with patch("intriniorealtime.options_client.time.sleep"): + client.stop() + + self.assertEqual(client._IntrinioRealtimeOptionsClient__channels, {"AAPL"}) + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) + ws.send_binary.assert_called() + leave = ws.send_binary.call_args[0][0] + self.assertEqual(leave[0], 76) + + def test_join_after_stop_does_not_wait_for_ready(self): + client = _make_client() + client._IntrinioRealtimeOptionsClient__is_started = True + client._IntrinioRealtimeOptionsClient__webSocket = None + with patch("intriniorealtime.options_client.time.sleep"): + client.stop() + + with patch.object(client, "_IntrinioRealtimeOptionsClient__all_ready") as ready: + client.join("MSFT") + ready.assert_not_called() + self.assertIn("MSFT", client._IntrinioRealtimeOptionsClient__channels) + + def test_leave_drops_channel_if_send_raises(self): + client = _make_client() + client._IntrinioRealtimeOptionsClient__channels = {"AAPL"} + ws = Mock() + ws.isReady = True + ws.send_binary.side_effect = RuntimeError("connection closed") + client._IntrinioRealtimeOptionsClient__webSocket = ws + + with self.assertRaises(RuntimeError): + client.leave("AAPL") + + self.assertNotIn("AAPL", client._IntrinioRealtimeOptionsClient__channels) + + def test_stop_joins_socket_thread_with_timeout(self): + client = _make_client() + thread = Mock() + thread.ident = 1 + client._IntrinioRealtimeOptionsClient__socket_thread = thread + client._IntrinioRealtimeOptionsClient__webSocket = None + with patch("intriniorealtime.options_client.time.sleep"): + client.stop() + thread.join.assert_called_with(timeout=_CONNECT_TIMEOUT_SECONDS + 1) + + def test_open_then_drop_twice_backs_off_on_second_reconnect(self): + stop = threading.Event() + waits = [] + original_wait = threading.Event.wait + + def tracking_wait(self, timeout=None): + waits.append(timeout) + return original_wait(self, timeout) + + run_count = {"n": 0} + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + app.on_open(app) + if run_count["n"] >= 3: + stop.set() + + ws = _make_ws(stop) + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("intriniorealtime.options_client._STABLE_SESSION_SECONDS", 10), \ + patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0.01]), \ + patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.options_client.abort_websocket_app"), \ + patch("threading.Event.wait", tracking_wait), \ + patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): + ws.start() + + self.assertEqual(run_count["n"], 3) + self.assertIn(0.01, waits) + + def test_stable_session_drop_skips_first_backoff(self): + stop = threading.Event() + waits = [] + original_wait = threading.Event.wait + + def tracking_wait(self, timeout=None): + waits.append(timeout) + return original_wait(self, timeout) + + run_count = {"n": 0} + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + if run_count["n"] == 1: + app.on_open(app) + time.sleep(0.06) + return + stop.set() + + ws = _make_ws(stop) + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("intriniorealtime.options_client._STABLE_SESSION_SECONDS", 0.05), \ + patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [10]), \ + patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.options_client.abort_websocket_app"), \ + patch("threading.Event.wait", tracking_wait), \ + patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): + ws.start() + + self.assertEqual(run_count["n"], 2) + self.assertNotIn(10, waits) + + def test_failed_handshake_still_backs_off_before_reconnect(self): + stop = threading.Event() + waits = [] + original_wait = threading.Event.wait + + def tracking_wait(self, timeout=None): + waits.append(timeout) + return original_wait(self, timeout) + + run_count = {"n": 0} + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + if run_count["n"] == 1: + return + stop.set() + + ws = _make_ws(stop) + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ + patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0.01]), \ + patch("websocket.setdefaulttimeout"), \ + patch("intriniorealtime.options_client.abort_websocket_app"), \ + patch("threading.Event.wait", tracking_wait), \ + patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): + ws.start() + + self.assertEqual(run_count["n"], 2) + self.assertIn(0.01, waits) + + def test_stop_sets_flag_before_sleep(self): + client = _make_client() + saw_stop = {} + + def fake_sleep(seconds): + saw_stop["set"] = client._stop_event.is_set() + + with patch("intriniorealtime.options_client.time.sleep", side_effect=fake_sleep): + client.stop() + self.assertTrue(saw_stop["set"]) + + def test_start_does_not_spawn_thread_when_stopped_during_auth(self): + client = _make_client() + + def get_token(): + client._stop_event.set() + return None + + with patch.object(client, "_IntrinioRealtimeOptionsClient__get_token", side_effect=get_token), \ + patch.object(threading.Thread, "start") as thread_start: + client.start() + + thread_start.assert_not_called() + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) + + def test_start_retries_until_token_then_connects(self): + client = _make_client() + tokens = {"n": 0} + start_calls = [] + + def get_token(): + tokens["n"] += 1 + if tokens["n"] < 3: + return None + return "tok" + + def fake_ws_start(self): + start_calls.append(1) + + with patch.object(client, "_IntrinioRealtimeOptionsClient__get_token", side_effect=get_token), \ + patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0]), \ + patch.object(_WebSocket, "start", fake_ws_start): + client.start() + + self.assertEqual(tokens["n"], 3) + self.assertEqual(start_calls, [1]) + self.assertTrue(client._IntrinioRealtimeOptionsClient__is_started) + + def test_start_on_healthy_session_is_noop(self): + client = _make_client() + thread = Mock() + thread.is_alive.return_value = True + client._IntrinioRealtimeOptionsClient__socket_thread = thread + client._IntrinioRealtimeOptionsClient__is_started = True + client._stop_event.clear() + + with patch.object(client, "_IntrinioRealtimeOptionsClient__get_token") as get_token: + client.start() + + get_token.assert_not_called() + + def test_start_waits_for_previous_socket_thread_then_starts(self): + client = _make_client() + leftover = Mock() + leftover.ident = 1 + leftover.is_alive.side_effect = [True, True, False] + leftover.join = Mock() + client._IntrinioRealtimeOptionsClient__socket_thread = leftover + client._IntrinioRealtimeOptionsClient__is_started = False + client._stop_event.set() + start_calls = [] + + def fake_ws_start(self): + start_calls.append(1) + + with patch("intriniorealtime.options_client.requests.get") as get, \ + patch.object(_WebSocket, "start", fake_ws_start): + get.return_value.status_code = 200 + get.return_value.text = "tok" + client.start() + + leftover.join.assert_called() + self.assertEqual(start_calls, [1]) + + def test_stop_completes_teardown_if_leave_raises(self): + client = _make_client() + client._IntrinioRealtimeOptionsClient__channels = {"AAPL"} + client._IntrinioRealtimeOptionsClient__is_started = True + ws = Mock() + ws.isReady = True + ws.sock = Mock() + ws.send_binary.side_effect = RuntimeError("connection closed") + client._IntrinioRealtimeOptionsClient__webSocket = ws + thread = Mock() + thread.ident = 1 + client._IntrinioRealtimeOptionsClient__socket_thread = thread + + with patch("intriniorealtime.options_client.time.sleep"): + client.stop() + + ws.sock.settimeout.assert_called_with(1) + ws.stop.assert_called() + thread.join.assert_called() + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) + + def test_start_uses_new_stop_event_so_old_workers_stay_stopped(self): + client = _make_client() + old_event = client._stop_event + old_event.set() + old_worker = Mock() + old_worker.ident = 1 + client._IntrinioRealtimeOptionsClient__worker_threads = [old_worker] + client._IntrinioRealtimeOptionsClient__is_started = False + start_calls = [] + + def fake_ws_start(self): + start_calls.append(1) + + with patch("intriniorealtime.options_client.requests.get") as get, \ + patch.object(_WebSocket, "start", fake_ws_start): + get.return_value.status_code = 200 + get.return_value.text = "tok" + client.start() + + self.assertIsNot(client._stop_event, old_event) + self.assertTrue(old_event.is_set()) + self.assertFalse(client._stop_event.is_set()) + self.assertNotIn(old_worker, client._IntrinioRealtimeOptionsClient__worker_threads) + self.assertEqual(start_calls, [1]) + + def test_start_aborts_leftover_websocket_while_joining(self): + client = _make_client() + leftover = Mock() + leftover.ident = 1 + leftover.is_alive.side_effect = [True, True, False] + leftover.join = Mock() + ws = Mock() + client._IntrinioRealtimeOptionsClient__socket_thread = leftover + client._IntrinioRealtimeOptionsClient__webSocket = ws + client._IntrinioRealtimeOptionsClient__is_started = False + client._stop_event.set() + start_calls = [] + + def fake_ws_start(self): + start_calls.append(1) + + with patch("intriniorealtime.options_client.abort_websocket_app") as abort, \ + patch("intriniorealtime.options_client.requests.get") as get, \ + patch.object(_WebSocket, "start", fake_ws_start): + get.return_value.status_code = 200 + get.return_value.text = "tok" + client.start() + + abort.assert_called() + leftover.join.assert_called() + self.assertEqual(start_calls, [1]) + + def test_start_returns_if_stop_during_leftover_join(self): + client = _make_client() + leftover = Mock() + leftover.ident = 1 + leftover.is_alive.return_value = True + + def join(timeout=None): + client._stop_generation += 1 + client._stop_event.set() + + leftover.join.side_effect = join + client._IntrinioRealtimeOptionsClient__socket_thread = leftover + client._IntrinioRealtimeOptionsClient__is_started = False + client._stop_event.set() + + with patch("intriniorealtime.options_client.abort_websocket_app"), \ + patch.object(client, "_IntrinioRealtimeOptionsClient__get_token") as get_token, \ + patch.object(threading.Thread, "start") as thread_start: + client.start() + + get_token.assert_not_called() + thread_start.assert_not_called() + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) diff --git a/tests/test_websocket_abort.py b/tests/test_websocket_abort.py new file mode 100644 index 0000000..c7ecf3a --- /dev/null +++ b/tests/test_websocket_abort.py @@ -0,0 +1,102 @@ +import socket +import threading +import time +import unittest +from unittest.mock import Mock, call + +from intriniorealtime._websocket import abort_websocket_app, should_abort_handshake + + +class FakeWebSocket: + def __init__(self, raw, connected=False): + self.sock = raw + self.connected = connected + + def close(self, **kwargs): + if not self.connected: + return + self.shutdown() + + def shutdown(self): + if self.sock: + self.sock.close() + self.sock = None + self.connected = False + + +class FakeApp: + def __init__(self, ws): + self.sock = ws + self.keep_running = True + + def close(self, **kwargs): + self.keep_running = False + if self.sock: + self.sock.close() + self.sock = None + + +class AbortWebsocketAppTests(unittest.TestCase): + def test_shutdowns_raw_socket_when_websocket_is_not_connected(self): + raw = Mock() + ws = FakeWebSocket(raw, connected=False) + app = FakeApp(ws) + + abort_websocket_app(app) + + self.assertFalse(app.keep_running) + raw.shutdown.assert_called_with(socket.SHUT_RDWR) + raw.close.assert_called() + self.assertEqual(raw.mock_calls[0], call.shutdown(socket.SHUT_RDWR)) + self.assertEqual(raw.mock_calls[1], call.close()) + self.assertIsNone(app.sock) + self.assertIsNone(ws.sock) + self.assertFalse(ws.connected) + + def test_noop_on_none(self): + abort_websocket_app(None) + + def test_sets_keep_running_false_without_sock(self): + app = FakeApp(None) + abort_websocket_app(app) + self.assertFalse(app.keep_running) + + +class ShouldAbortHandshakeTests(unittest.TestCase): + def test_true_on_timeout_for_current_generation(self): + handshake = threading.Event() + stop = threading.Event() + self.assertTrue(should_abort_handshake(handshake, stop, 1, 1, 0.01)) + + def test_false_when_handshake_completes(self): + handshake = threading.Event() + handshake.set() + stop = threading.Event() + self.assertFalse(should_abort_handshake(handshake, stop, 1, 1, 0.01)) + + def test_false_for_stale_generation(self): + handshake = threading.Event() + stop = threading.Event() + self.assertFalse(should_abort_handshake(handshake, stop, 1, 2, 0.01)) + + def test_true_when_stopped_and_handshake_timed_out(self): + handshake = threading.Event() + stop = threading.Event() + stop.set() + self.assertTrue(should_abort_handshake(handshake, stop, 1, 1, 0.01)) + + def test_callable_current_generation_is_read_after_wait(self): + handshake = threading.Event() + stop = threading.Event() + state = {"g": 1} + + def bump(): + time.sleep(0.02) + state["g"] = 2 + + threading.Thread(target=bump, daemon=True).start() + self.assertFalse(should_abort_handshake(handshake, stop, 1, lambda: state["g"], 0.05)) + + +if __name__ == "__main__": + unittest.main() From 39b6d5557a4ed462dc15491bb7a824a0fa874681 Mon Sep 17 00:00:00 2001 From: Shawn Snyder Date: Wed, 23 Sep 2026 00:49:29 -0500 Subject: [PATCH 3/7] version bump --- intriniorealtime/equities_client.py | 2 +- setup.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/intriniorealtime/equities_client.py b/intriniorealtime/equities_client.py index 2fcdbae..8f86c38 100644 --- a/intriniorealtime/equities_client.py +++ b/intriniorealtime/equities_client.py @@ -37,7 +37,7 @@ HEADER_MESSAGE_FORMAT_KEY = "UseNewEquitiesFormat" HEADER_MESSAGE_FORMAT_VALUE = "v2" HEADER_CLIENT_INFORMATION_KEY = "Client-Information" -HEADER_CLIENT_INFORMATION_VALUE = "IntrinioPythonSDKv6.3.0" +HEADER_CLIENT_INFORMATION_VALUE = "IntrinioPythonSDKv6.4.0" class EquitiesQuote: diff --git a/setup.py b/setup.py index 5b6d081..8a7d21d 100644 --- a/setup.py +++ b/setup.py @@ -7,7 +7,7 @@ def readme(): setup( name = 'intriniorealtime', packages = ['intriniorealtime'], - version = '6.3.0', + version = '6.4.0', author = 'Intrinio Python SDK for Real-Time Stock Prices', author_email = 'success@intrinio.com', url = 'https://intrinio.com', @@ -16,7 +16,7 @@ def readme(): long_description_content_type = 'text/markdown', install_requires = ['requests>=2.26.0','websocket-client>=1.2.1','wsaccel>=0.6.3', 'intrinio-sdk>=6.26.0'], python_requires = '~=3.10', - download_url = 'https://github.com/intrinio/intrinio-realtime-python-sdk/archive/v6.3.0.tar.gz', + download_url = 'https://github.com/intrinio/intrinio-realtime-python-sdk/archive/v6.4.0.tar.gz', keywords = ['realtime','stock prices','intrinio','stock market','stock data','financial'], classifiers = [ 'Intended Audience :: Financial and Insurance Industry', From d0399326b002e0e59a3fa9b2086f3da9478d651c Mon Sep 17 00:00:00 2001 From: Shawn Snyder Date: Wed, 23 Sep 2026 01:16:28 -0500 Subject: [PATCH 4/7] more cases covered --- intriniorealtime/_websocket.py | 35 +++++ intriniorealtime/equities_client.py | 18 +-- intriniorealtime/options_client.py | 193 +++++++++++++++++----------- tests/test_equities_client.py | 14 +- tests/test_options_client.py | 28 ++++ tests/test_websocket_abort.py | 35 ++++- 6 files changed, 238 insertions(+), 85 deletions(-) diff --git a/intriniorealtime/_websocket.py b/intriniorealtime/_websocket.py index 82c0403..02549e3 100644 --- a/intriniorealtime/_websocket.py +++ b/intriniorealtime/_websocket.py @@ -2,6 +2,11 @@ import threading from typing import Optional +import websocket + + +_websocket_timeout_lock = threading.Lock() + def abort_websocket_app(app) -> None: if app is None: @@ -67,3 +72,33 @@ def clear_socket_timeout(app) -> None: sock.settimeout(None) except Exception: pass + + +def run_forever_with_connect_timeout(app, timeout_seconds: float, **kwargs): + """Run an app with a connect timeout without leaving a global timeout behind.""" + _websocket_timeout_lock.acquire() + previous_timeout = websocket.getdefaulttimeout() + original_on_open = app.on_open + restored = False + + def restore_timeout() -> None: + nonlocal restored + if restored: + return + websocket.setdefaulttimeout(previous_timeout) + restored = True + _websocket_timeout_lock.release() + + def on_open(ws, *args): + clear_socket_timeout(app) + restore_timeout() + if original_on_open: + original_on_open(ws, *args) + + app.on_open = on_open + websocket.setdefaulttimeout(timeout_seconds) + try: + return app.run_forever(**kwargs) + finally: + restore_timeout() + app.on_open = original_on_open diff --git a/intriniorealtime/equities_client.py b/intriniorealtime/equities_client.py index 8f86c38..d251dbb 100644 --- a/intriniorealtime/equities_client.py +++ b/intriniorealtime/equities_client.py @@ -10,7 +10,7 @@ from enum import IntEnum, unique from typing import Optional, Dict, Any -from ._websocket import abort_websocket_app, clear_socket_timeout, should_abort_handshake +from ._websocket import abort_websocket_app, run_forever_with_connect_timeout, should_abort_handshake SELF_HEAL_BACKOFFS = [10, 30, 60, 300, 600] CONNECT_TIMEOUT_SECONDS = 30 @@ -407,7 +407,6 @@ def _clear_continuation_state(self): break def run(self): - websocket.setdefaulttimeout(CONNECT_TIMEOUT_SECONDS) generation = 0 while self.enabled and not self.client._stop_event.is_set(): generation += 1 @@ -433,7 +432,6 @@ def on_open_with_timeout(ws, *args): session_opened = True session_opened_at = time.time() handshake_event.set() - clear_socket_timeout(app) if original_on_open: original_on_open(ws, *args) @@ -456,7 +454,11 @@ def watchdog(this_generation=generation, handshake_event=handshake_event): abort_websocket_app(app) break try: - app.run_forever(skip_utf8_validation=True) # skip_utf8_validation for more performance + run_forever_with_connect_timeout( + app, + CONNECT_TIMEOUT_SECONDS, + skip_utf8_validation=True, + ) # skip_utf8_validation for more performance except Exception as e: self.client.logger.error(f"Websocket ERROR: {repr(e)}") @@ -511,10 +513,10 @@ def stitch(self): if full is None: full = partial else: - full = full.join(partial) + full = full + partial return full - def on_cont_message(self, partial_message, is_last): # The 3rd argument is continue flag. if 0, the data continue + def on_cont_message(self, ws, partial_message, is_last): # The 3rd argument is continue flag. if 0, the data continue try: if DEBUGGING: # This is here for performance reasons so we don't use slow reflection on every message. if isinstance(partial_message, str): @@ -531,7 +533,7 @@ def on_cont_message(self, partial_message, is_last): # The 3rd argument is conti else: self.continuation_queue.put(partial_message) full_message = self.stitch() - self.on_message(self.client.ws, full_message) + self.on_message(ws, full_message) finally: self.continuation_lock.release() except queue.Full: @@ -658,4 +660,4 @@ def run(self): items_in_message = message[0] start_index = 1 for i in range(0, items_in_message): - start_index = self.parse_message(message, start_index, backlog_len) \ No newline at end of file + start_index = self.parse_message(message, start_index, backlog_len) diff --git a/intriniorealtime/options_client.py b/intriniorealtime/options_client.py index 17c9832..bbf5f80 100644 --- a/intriniorealtime/options_client.py +++ b/intriniorealtime/options_client.py @@ -9,7 +9,7 @@ from collections.abc import Callable from enum import IntEnum, unique -from ._websocket import abort_websocket_app, clear_socket_timeout, should_abort_handshake +from ._websocket import abort_websocket_app, run_forever_with_connect_timeout, should_abort_handshake _SELF_HEAL_BACKOFFS = [10, 30, 60, 300, 600] _CONNECT_TIMEOUT_SECONDS = 30 @@ -334,7 +334,6 @@ def __on_open(self, ws): self.__handshake_event.set() self.__session_opened = True self.__session_opened_at = time.time() - clear_socket_timeout(self) _log.info("Websocket - Connected") self.__wsLock.acquire() try: @@ -376,7 +375,7 @@ def __stitch(self): if full is None: full = partial else: - full = full.join(partial) + full = full + partial return full def __on_data(self, ws, data, code, is_last): #continueFlag - If 0, the data continues @@ -406,7 +405,6 @@ def __on_data(self, ws, data, code, is_last): #continueFlag - If 0, the data con _log.error("Error received: {0}".format(data)) def start(self): - websocket.setdefaulttimeout(_CONNECT_TIMEOUT_SECONDS) while not self.__stop_flag.is_set(): handshake_event = threading.Event() self.__handshake_event = handshake_event @@ -437,7 +435,11 @@ def watchdog(this_generation=generation, handshake_event=handshake_event): abort_websocket_app(self) return try: - self.run_forever(skip_utf8_validation=True) + run_forever_with_connect_timeout( + self, + _CONNECT_TIMEOUT_SECONDS, + skip_utf8_validation=True, + ) except Exception as e: _log.warning("Websocket - Attempt failed: {0}".format(e)) @@ -772,6 +774,7 @@ def __init__(self, config: Config, on_trade: Callable[[OptionsTrade], None], on_ self.__data: queue.Queue = queue.Queue() self.__t_lock: threading.Lock = threading.Lock() self.__ws_lock: threading.Lock = threading.Lock() + self.__lifecycle_lock: threading.RLock = threading.RLock() self.__on_trade = on_trade self.__on_quote = on_quote self.__on_refresh = on_refresh @@ -782,6 +785,8 @@ def __init__(self, config: Config, on_trade: Callable[[OptionsTrade], None], on_ self.__worker_threads: list[threading.Thread] = self.__make_worker_threads() self.__socket_thread: threading.Thread = None self.__is_started: bool = False + self.__is_starting: bool = False + self.__is_stopping: bool = False _log.setLevel(config.log_level) def __make_worker_threads(self) -> list[threading.Thread]: @@ -920,76 +925,108 @@ def leave_firehose(self): if "$FIREHOSE" in self.__channels: self.__leave("$FIREHOSE") - def __socket_start_fn(self, token: str): + def __socket_start_fn(self, token: str, stop_event: threading.Event, + worker_threads: list[threading.Thread], data_queue: queue.Queue): if not token: _log.error("Websocket - Missing token") return _log.info("Websocket - Connecting...") ws_url: str = self.__get_web_socket_url(token) - self.__webSocket = _WebSocket(ws_url, - self.__ws_lock, - self.__worker_threads, - self.__get_channels, - self.__get_token, - self.__get_web_socket_url, - self.__use_on_trade, - self.__use_on_quote, - self.__use_on_refresh, - self.__use_on_unusual_activity, - self.__data, - self._stop_event) - self.__webSocket.start() + if stop_event.is_set(): + return + web_socket = _WebSocket(ws_url, + self.__ws_lock, + worker_threads, + self.__get_channels, + self.__get_token, + self.__get_web_socket_url, + self.__use_on_trade, + self.__use_on_quote, + self.__use_on_refresh, + self.__use_on_unusual_activity, + data_queue, + stop_event) + self.__webSocket = web_socket + web_socket.start() def start(self): if (not (self.__use_on_trade or self.__use_on_quote or self.__use_on_refresh or self.__use_on_unusual_activity)): raise ValueError("You must set at least one callback method before starting client") - if (self.__is_started - and self.__socket_thread is not None - and self.__socket_thread.is_alive() - and not self._stop_event.is_set()): - return - stop_generation = self._stop_generation - while True: - leftover = self.__socket_thread - if leftover is not None and leftover.is_alive(): - abort_websocket_app(self.__webSocket) - leftover.join(timeout=_CONNECT_TIMEOUT_SECONDS + 1) - if leftover.is_alive(): + with self.__lifecycle_lock: + if (self.__is_started + and self.__socket_thread is not None + and self.__socket_thread.is_alive() + and not self._stop_event.is_set()): + return + if self.__is_starting: + return + if self.__is_stopping: + return + self.__is_starting = True + stop_generation = self._stop_generation + try: + while True: + leftover = self.__socket_thread + if leftover is not None and leftover.is_alive(): + abort_websocket_app(self.__webSocket) + leftover.join(timeout=_CONNECT_TIMEOUT_SECONDS + 1) + if leftover.is_alive(): + with self.__lifecycle_lock: + if self._stop_generation != stop_generation: + return + continue + with self.__lifecycle_lock: if self._stop_generation != stop_generation: return + self._stop_event = threading.Event() + stop_event = self._stop_event + self.__data = queue.Queue() + data_queue = self.__data + self.__worker_threads = self.__make_worker_threads() + worker_threads = self.__worker_threads + token: str = self.__get_token() + if stop_event.is_set(): + return + if not token: + if stop_event.wait(timeout=_SELF_HEAL_BACKOFFS[0]): + return continue - if self._stop_generation != stop_generation: - return - self._stop_event = threading.Event() - self.__worker_threads = self.__make_worker_threads() - token: str = self.__get_token() - if self._stop_event.is_set(): + with self.__lifecycle_lock: + if (self._stop_generation != stop_generation + or self._stop_event is not stop_event + or stop_event.is_set()): + return + socket_thread = threading.Thread( + group=None, + target=self.__socket_start_fn, + args=(token, stop_event, worker_threads, data_queue), + kwargs={}, + daemon=True + ) + self.__socket_thread = socket_thread + self.__is_started = True + try: + socket_thread.start() + except Exception: + self.__is_started = False + raise return - if not token: - if self._stop_event.wait(timeout=_SELF_HEAL_BACKOFFS[0]): - return - continue - self.__ws_lock.acquire() - try: - self.__socket_thread = threading.Thread( - group=None, - target=self.__socket_start_fn, - args=(token,), - kwargs={}, - daemon=True - ) - finally: - self.__ws_lock.release() - self.__socket_thread.start() - self.__is_started = True - return + finally: + with self.__lifecycle_lock: + self.__is_starting = False def stop(self): - _log.info("Stopping...") - self._stop_generation += 1 - self._stop_event.set() - try: + with self.__lifecycle_lock: + if self.__is_stopping: + return + self.__is_stopping = True + _log.info("Stopping...") + self._stop_generation += 1 + self._stop_event.set() ws = self.__webSocket + worker_threads = self.__worker_threads + socket_thread = self.__socket_thread + try: if ws is not None: try: sock = getattr(ws, "sock", None) @@ -1005,26 +1042,32 @@ def stop(self): except Exception as e: _log.warning("Websocket - Leave failed: {0}".format(e)) break + if ws is not None: + ws.stop() time.sleep(1.0) finally: - self.__ws_lock.acquire() try: - if self.__webSocket is not None: - self.__webSocket.isReady = False + self.__ws_lock.acquire() + try: + if ws is not None: + ws.isReady = False + finally: + self.__ws_lock.release() + if ws is not None: + ws.stop() + for i in range(len(worker_threads)): + worker = worker_threads[i] + if worker.ident is not None: + worker.join(timeout=_CONNECT_TIMEOUT_SECONDS + 1) + _log.debug("Worker thread {0} joined".format(i)) + if socket_thread is not None and socket_thread.ident is not None: + socket_thread.join(timeout=_CONNECT_TIMEOUT_SECONDS + 1) + _log.debug("Socket thread joined") finally: - self.__ws_lock.release() - if self.__webSocket is not None: - self.__webSocket.stop() - for i in range(len(self.__worker_threads)): - worker = self.__worker_threads[i] - if worker.ident is not None: - worker.join(timeout=_CONNECT_TIMEOUT_SECONDS + 1) - _log.debug("Worker thread {0} joined".format(i)) - if self.__socket_thread is not None and self.__socket_thread.ident is not None: - self.__socket_thread.join(timeout=_CONNECT_TIMEOUT_SECONDS + 1) - _log.debug("Socket thread joined") - self.__is_started = False - _log.info("Stopped") + with self.__lifecycle_lock: + self.__is_started = False + self.__is_stopping = False + _log.info("Stopped") def get_stats(self) -> tuple[int, int, int]: return _dataMsgCount, _txtMsgCount, self.__data.qsize() diff --git a/tests/test_equities_client.py b/tests/test_equities_client.py index bc038ed..3122da9 100644 --- a/tests/test_equities_client.py +++ b/tests/test_equities_client.py @@ -88,7 +88,8 @@ def fake_run_forever(app, *args, **kwargs): patch("websocket.WebSocketApp.run_forever", fake_run_forever): receiver.run() - set_timeout.assert_called_with(CONNECT_TIMEOUT_SECONDS) + self.assertIn(((CONNECT_TIMEOUT_SECONDS,), {}), set_timeout.call_args_list) + set_timeout.assert_called_with(None) sock.settimeout.assert_called_with(None) def test_refresh_token_failure_backs_off_and_skips_socket_until_auth_succeeds(self): @@ -257,6 +258,17 @@ def fake_run_forever(app, *args, **kwargs): self.assertEqual(run_count["n"], 2) self.assertEqual(seen_empty, [True, True]) + def test_fragmented_message_is_reassembled_and_delivered(self): + client = _make_client() + receiver = EquitiesQuoteReceiver(client) + ws = Mock() + receiver.on_message = Mock() + + receiver.on_cont_message(ws, b"abc", 0) + receiver.on_cont_message(ws, b"def", 1) + + receiver.on_message.assert_called_once_with(ws, b"abcdef") + def test_refresh_token_passes_timeout(self): client = _make_client() with patch("intriniorealtime.equities_client.requests.get") as get: diff --git a/tests/test_options_client.py b/tests/test_options_client.py index 5e8b9cd..a373454 100644 --- a/tests/test_options_client.py +++ b/tests/test_options_client.py @@ -415,6 +415,34 @@ def fake_ws_start(self): self.assertEqual(start_calls, [1]) self.assertTrue(client._IntrinioRealtimeOptionsClient__is_started) + def test_stop_during_socket_thread_start_does_not_publish_started_state(self): + client = _make_client() + + def stop_before_thread_runs(): + client.stop() + + with patch.object(client, "_IntrinioRealtimeOptionsClient__get_token", return_value="tok"), \ + patch.object(threading.Thread, "start", side_effect=stop_before_thread_runs), \ + patch("intriniorealtime.options_client.time.sleep"): + client.start() + + self.assertTrue(client._stop_event.is_set()) + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) + + def test_start_replaces_data_queue_from_previous_session(self): + client = _make_client() + old_queue = client._IntrinioRealtimeOptionsClient__data + old_queue.put(b"stale") + + with patch.object(client, "_IntrinioRealtimeOptionsClient__get_token", return_value="tok"), \ + patch.object(_WebSocket, "start"): + client.start() + + new_queue = client._IntrinioRealtimeOptionsClient__data + self.assertIsNot(new_queue, old_queue) + self.assertTrue(new_queue.empty()) + self.assertEqual(old_queue.get_nowait(), b"stale") + def test_start_on_healthy_session_is_noop(self): client = _make_client() thread = Mock() diff --git a/tests/test_websocket_abort.py b/tests/test_websocket_abort.py index c7ecf3a..f38bc16 100644 --- a/tests/test_websocket_abort.py +++ b/tests/test_websocket_abort.py @@ -4,7 +4,13 @@ import unittest from unittest.mock import Mock, call -from intriniorealtime._websocket import abort_websocket_app, should_abort_handshake +import websocket + +from intriniorealtime._websocket import ( + abort_websocket_app, + run_forever_with_connect_timeout, + should_abort_handshake, +) class FakeWebSocket: @@ -62,6 +68,33 @@ def test_sets_keep_running_false_without_sock(self): self.assertFalse(app.keep_running) +class ConnectTimeoutTests(unittest.TestCase): + def test_connect_timeout_is_restored_as_soon_as_socket_opens(self): + original_timeout = websocket.getdefaulttimeout() + websocket.setdefaulttimeout(17) + observed = [] + sock = Mock() + app = Mock() + app.sock = sock + + def on_open(ws): + observed.append(websocket.getdefaulttimeout()) + + def run_forever(**kwargs): + self.assertEqual(websocket.getdefaulttimeout(), 3) + app.on_open(app) + + app.on_open = on_open + app.run_forever.side_effect = run_forever + try: + run_forever_with_connect_timeout(app, 3) + self.assertEqual(observed, [17]) + self.assertEqual(websocket.getdefaulttimeout(), 17) + sock.settimeout.assert_called_once_with(None) + finally: + websocket.setdefaulttimeout(original_timeout) + + class ShouldAbortHandshakeTests(unittest.TestCase): def test_true_on_timeout_for_current_generation(self): handshake = threading.Event() From bd9269634e8c7e58fcb549309b8ea3bf51e1c6b9 Mon Sep 17 00:00:00 2001 From: Shawn Snyder Date: Thu, 24 Sep 2026 17:19:03 -0500 Subject: [PATCH 5/7] more --- intriniorealtime/_websocket.py | 104 +----- intriniorealtime/equities_client.py | 427 +++++++++++++-------- intriniorealtime/options_client.py | 532 +++++++++++++++++---------- tests/test_equities_client.py | 502 +++++++++++++++++++++++-- tests/test_options_client.py | 552 +++++++++++++++++++++++++--- tests/test_websocket_abort.py | 120 ++---- 6 files changed, 1620 insertions(+), 617 deletions(-) diff --git a/intriniorealtime/_websocket.py b/intriniorealtime/_websocket.py index 02549e3..a5f71ab 100644 --- a/intriniorealtime/_websocket.py +++ b/intriniorealtime/_websocket.py @@ -1,50 +1,14 @@ -import socket import threading +import time from typing import Optional -import websocket - -_websocket_timeout_lock = threading.Lock() - - -def abort_websocket_app(app) -> None: +def close_websocket_app(app) -> None: + """Best-effort close using websocket-client's public API only.""" if app is None: return try: - app.keep_running = False - except Exception: - pass - - ws = getattr(app, "sock", None) - raw = getattr(ws, "sock", None) if ws is not None else None - - if raw is not None: - try: - raw.shutdown(socket.SHUT_RDWR) - except Exception: - pass - try: - raw.close() - except Exception: - pass - - try: - if ws is not None: - ws.sock = None - ws.connected = False - except Exception: - pass - - try: - app.sock = None - except Exception: - pass - - try: - close = getattr(app, "close", None) - if callable(close): - close() + app.close() except Exception: pass @@ -54,51 +18,17 @@ def should_abort_handshake( stop_flag: Optional[threading.Event], generation: int, current_generation, - timeout_seconds: float, + deadline: float, ) -> bool: - if handshake_event.wait(timeout_seconds): - return False - current = current_generation() if callable(current_generation) else current_generation - if generation != current: - return False - return True - - -def clear_socket_timeout(app) -> None: - sock = getattr(app, "sock", None) - if sock is None: - return - try: - sock.settimeout(None) - except Exception: - pass - - -def run_forever_with_connect_timeout(app, timeout_seconds: float, **kwargs): - """Run an app with a connect timeout without leaving a global timeout behind.""" - _websocket_timeout_lock.acquire() - previous_timeout = websocket.getdefaulttimeout() - original_on_open = app.on_open - restored = False - - def restore_timeout() -> None: - nonlocal restored - if restored: - return - websocket.setdefaulttimeout(previous_timeout) - restored = True - _websocket_timeout_lock.release() - - def on_open(ws, *args): - clear_socket_timeout(app) - restore_timeout() - if original_on_open: - original_on_open(ws, *args) - - app.on_open = on_open - websocket.setdefaulttimeout(timeout_seconds) - try: - return app.run_forever(**kwargs) - finally: - restore_timeout() - app.on_open = original_on_open + while True: + if handshake_event.is_set(): + return False + if stop_flag is not None and stop_flag.is_set(): + return False + current = current_generation() if callable(current_generation) else current_generation + if generation != current: + return False + remaining = deadline - time.monotonic() + if remaining <= 0: + return True + handshake_event.wait(min(remaining, 0.1)) diff --git a/intriniorealtime/equities_client.py b/intriniorealtime/equities_client.py index d251dbb..530186a 100644 --- a/intriniorealtime/equities_client.py +++ b/intriniorealtime/equities_client.py @@ -10,7 +10,7 @@ from enum import IntEnum, unique from typing import Optional, Dict, Any -from ._websocket import abort_websocket_app, run_forever_with_connect_timeout, should_abort_handshake +from ._websocket import close_websocket_app, should_abort_handshake SELF_HEAL_BACKOFFS = [10, 30, 60, 300, 600] CONNECT_TIMEOUT_SECONDS = 30 @@ -134,9 +134,10 @@ def __init__(self, configuration: Dict[str, Any], on_trade: Optional[callable], self.logger.addHandler(log_handler) if 'max_queue_size' in configuration: - self.quotes = queue.Queue(maxsize=configuration['max_queue_size']) + self.max_queue_size = configuration['max_queue_size'] else: - self.quotes = queue.Queue(maxsize=MAX_QUEUE_SIZE) + self.max_queue_size = MAX_QUEUE_SIZE + self.quotes = queue.Queue(maxsize=self.max_queue_size) if self.api_key: if not self.valid_api_key(self.api_key): @@ -175,6 +176,12 @@ def __init__(self, configuration: Dict[str, Any], on_trade: Optional[callable], self.last_queue_warning_time = 0 self.last_self_heal_backoff = -1 self._stop_event = threading.Event() + self._connect_lock = threading.Lock() + self._disconnect_lock = threading.Lock() + self._lifecycle_lock = threading.RLock() + self._channels_lock = threading.RLock() + self._lifecycle_generation = 0 + self._disconnecting = False self.quote_handler.start() def auth_url(self) -> str: @@ -235,45 +242,105 @@ def do_backoff(self): self._stop_event.wait(timeout=backoff) def connect(self): - self._stop_event.clear() - if self.quote_receiver is not None and self.quote_receiver.is_alive() and self.quote_receiver.enabled: - return - - self.last_self_heal_backoff = -1 - while not self._stop_event.is_set(): - try: - self.logger.info("Connecting...") - self.ready = False - self.joined_channels = set() - self.refresh_token() - if self._stop_event.is_set(): + with self._lifecycle_lock: + if self._disconnecting: + raise RuntimeError("Cannot connect while client is disconnecting") + requested_generation = self._lifecycle_generation + with self._connect_lock: + with self._lifecycle_lock: + if self._disconnecting: + raise RuntimeError("Cannot connect while client is disconnecting") + if requested_generation != self._lifecycle_generation: return - self.refresh_websocket() + self._stop_event.clear() + if self.quote_receiver is not None and self.quote_receiver.is_alive() and self.quote_receiver.enabled: return - except Exception as e: - self.logger.error(f"Cannot connect: {repr(e)}") - self.do_backoff() + + self.last_self_heal_backoff = -1 + while not self._stop_event.is_set(): + try: + self.logger.info("Connecting...") + with self._channels_lock: + self.ready = False + self.joined_channels = set() + self.refresh_quote_handler() + self.refresh_token() + with self._lifecycle_lock: + if (requested_generation != self._lifecycle_generation + or self._stop_event.is_set()): + return + self.refresh_websocket(requested_generation) + return + except Exception as e: + self.logger.error(f"Cannot connect: {repr(e)}") + self.do_backoff() def disconnect(self): - self._stop_event.set() - self.ready = False - self.joined_channels = set() + with self._disconnect_lock: + with self._lifecycle_lock: + self._disconnecting = True + self._lifecycle_generation += 1 + self._stop_event.set() + handler = self.quote_handler + receiver = self.quote_receiver + websocket_app = self.ws + if receiver is not None: + receiver.enabled = False + try: + if handler is not None: + handler.stop() - if self.quote_receiver: - self.quote_receiver.enabled = False + close_websocket_app(websocket_app) - abort_websocket_app(self.ws) + with self._channels_lock: + self.ready = False + self.joined_channels = set() - receiver = self.quote_receiver - if receiver is not None: - receiver.join(timeout=CONNECT_TIMEOUT_SECONDS + 1) - if receiver.is_alive(): - abort_websocket_app(self.ws) - else: - self.quote_receiver = None - self.ws = None - else: - self.ws = None + if receiver is not None: + receiver.join(timeout=CONNECT_TIMEOUT_SECONDS + 1) + if receiver.is_alive(): + close_websocket_app(websocket_app) + else: + finished_websocket = None + with self._lifecycle_lock: + if self.quote_receiver is receiver: + finished_websocket = self.ws + self.quote_receiver = None + self.ws = None + close_websocket_app(finished_websocket) + else: + with self._lifecycle_lock: + if self.ws is websocket_app: + self.ws = None + + if ( + handler is not None + and handler.ident is not None + and handler is not threading.current_thread() + ): + handler.join(timeout=CONNECT_TIMEOUT_SECONDS + 1) + if handler is not None and not handler.is_alive(): + with self._lifecycle_lock: + if self.quote_handler is handler: + self.quote_handler = None + finally: + with self._lifecycle_lock: + self._disconnecting = False + + def refresh_quote_handler(self): + handler = self.quote_handler + if handler is not None: + if handler.is_alive() and not handler.stopped: + return + current_thread = threading.current_thread() + replacing_current_stopped_handler = handler is current_thread and handler.stopped + if handler.ident is not None and handler is not current_thread: + handler.join(timeout=CONNECT_TIMEOUT_SECONDS + 1) + if handler.is_alive() and not replacing_current_stopped_handler: + raise RuntimeError("Previous quote handler did not exit") + self.quotes = queue.Queue(maxsize=self.max_queue_size) + self.quote_handler = EquitiesQuoteHandler(self, self.bypass_parsing) + self.quote_handler.start() def refresh_token(self): headers = {HEADER_CLIENT_INFORMATION_KEY: HEADER_CLIENT_INFORMATION_VALUE} @@ -288,25 +355,48 @@ def refresh_token(self): self.token = response.text self.logger.info("Authentication successful!") - def refresh_websocket(self): - receiver = self.quote_receiver + def refresh_websocket(self, expected_generation=None): + with self._lifecycle_lock: + if (expected_generation is not None + and expected_generation != self._lifecycle_generation): + return + receiver = self.quote_receiver if receiver is not None: if receiver.is_alive(): if receiver.enabled: return - abort_websocket_app(self.ws) + close_websocket_app(self.ws) receiver.join(timeout=CONNECT_TIMEOUT_SECONDS + 1) if receiver.is_alive(): - abort_websocket_app(self.ws) + close_websocket_app(self.ws) raise RuntimeError("Previous quote receiver did not exit") - self.quote_receiver = None - self.ws = None - self.quote_receiver = EquitiesQuoteReceiver(self) - self.quote_receiver.start() - - def on_connect(self): - self.ready = True - self.joined_channels = set() + with self._lifecycle_lock: + if self.quote_receiver is receiver: + self.quote_receiver = None + self.ws = None + new_receiver = EquitiesQuoteReceiver(self) + with self._lifecycle_lock: + if ((expected_generation is not None + and expected_generation != self._lifecycle_generation) + or self._stop_event.is_set()): + return + self.quote_receiver = new_receiver + try: + new_receiver.start() + except Exception: + if self.quote_receiver is new_receiver: + self.quote_receiver = None + raise + + def on_connect(self, ws=None, receiver=None): + with self._lifecycle_lock: + if (self._stop_event.is_set() + or (receiver is not None and self.quote_receiver is not receiver) + or (ws is not None and self.ws is not ws)): + return + with self._channels_lock: + self.ready = True + self.joined_channels = set() self.refresh_channels() def on_queue_full(self): @@ -318,42 +408,47 @@ def join(self, channels: list[str]): if isinstance(channels, str): channels = [channels] - self.channels = self.channels | set(channels) - self.refresh_channels() + with self._channels_lock: + self.channels = self.channels | set(channels) + self.refresh_channels() def leave(self, channels: list[str]): if isinstance(channels, str): channels = [channels] - self.channels = self.channels - set(channels) - self.refresh_channels() + with self._channels_lock: + self.channels = self.channels - set(channels) + self.refresh_channels() def leave_all(self): - self.channels = set() - self.refresh_channels() + with self._channels_lock: + self.channels = set() + self.refresh_channels() def refresh_channels(self): - if self.ready != True: - return - - # Join new channels - new_channels = self.channels - self.joined_channels - self.logger.debug(f"New channels: {new_channels}") - for channel in new_channels: - msg = self.join_binary_message(channel) - self.ws.send(msg, websocket.ABNF.OPCODE_BINARY) - self.logger.info(f"Joined channel {channel}") - - # Leave old channels - old_channels = self.joined_channels - self.channels - self.logger.debug(f"Old channels: {old_channels}") - for channel in old_channels: - msg = self.leave_binary_message(channel) - self.ws.send(msg, websocket.ABNF.OPCODE_BINARY) - self.logger.info(f"Left channel {channel}") - - self.joined_channels = self.channels.copy() - self.logger.debug(f"Current channels: {self.joined_channels}") + with self._channels_lock: + if self.ready != True or self._stop_event.is_set(): + return + + # Join new channels + new_channels = self.channels - self.joined_channels + self.logger.debug(f"New channels: {new_channels}") + for channel in new_channels: + msg = self.join_binary_message(channel) + self.ws.send(msg, websocket.ABNF.OPCODE_BINARY) + self.joined_channels.add(channel) + self.logger.info(f"Joined channel {channel}") + + # Leave old channels + old_channels = self.joined_channels - self.channels + self.logger.debug(f"Old channels: {old_channels}") + for channel in old_channels: + msg = self.leave_binary_message(channel) + self.ws.send(msg, websocket.ABNF.OPCODE_BINARY) + self.joined_channels.discard(channel) + self.logger.info(f"Left channel {channel}") + + self.logger.debug(f"Current channels: {self.joined_channels}") def join_binary_message(self, channel: str): if channel == "lobby": @@ -395,81 +490,103 @@ def __init__(self, client): self.daemon = True self.client = client self.enabled = True - self.continuation_queue = queue.Queue(100) - self.continuation_lock: threading.Lock = threading.Lock() - - def _clear_continuation_state(self): - with self.continuation_lock: - while True: - try: - self.continuation_queue.get_nowait() - except queue.Empty: - break def run(self): generation = 0 + generation_lock = threading.Lock() while self.enabled and not self.client._stop_event.is_set(): - generation += 1 - self._clear_continuation_state() + with generation_lock: + generation += 1 handshake_event = threading.Event() + attempt_cancelled = threading.Event() + connect_deadline = time.monotonic() + CONNECT_TIMEOUT_SECONDS + session = {"opened": False, "opened_at": 0.0} + app = None + + def on_open_with_timeout( + ws, + *args, + attempt_handshake=handshake_event, + cancelled=attempt_cancelled, + attempt_session=session, + attempt_generation=generation, + deadline=connect_deadline, + ): + rejected = False + with generation_lock: + if (cancelled.is_set() + or attempt_generation != generation + or not self.enabled + or self.client._stop_event.is_set() + or time.monotonic() >= deadline): + cancelled.set() + rejected = True + attempt_handshake.set() + if not rejected: + attempt_session["opened"] = True + attempt_session["opened_at"] = time.monotonic() + if rejected: + close_websocket_app(app) + return + self.on_open(ws) + + # Leave on_cont_message unset so websocket-client aggregates fragmented messages. app = websocket.WebSocketApp( self.client.websocket_url(), header={HEADER_MESSAGE_FORMAT_KEY: HEADER_MESSAGE_FORMAT_VALUE, HEADER_CLIENT_INFORMATION_KEY: HEADER_CLIENT_INFORMATION_VALUE}, - on_open=self.on_open, + on_open=on_open_with_timeout, on_close=self.on_close, on_message=self.on_message, - on_cont_message=self.on_cont_message, on_error=self.on_error ) - self.client.ws = app - original_on_open = app.on_open - - session_opened = False - session_opened_at = 0.0 - - def on_open_with_timeout(ws, *args): - nonlocal session_opened, session_opened_at - session_opened = True - session_opened_at = time.time() - handshake_event.set() - if original_on_open: - original_on_open(ws, *args) - - app.on_open = on_open_with_timeout - - def watchdog(this_generation=generation, handshake_event=handshake_event): + with self.client._lifecycle_lock: + if not self.enabled or self.client._stop_event.is_set(): + close_websocket_app(app) + break + self.client.ws = app + + def watchdog( + this_generation=generation, + handshake_event=handshake_event, + attempt_app=app, + cancelled=attempt_cancelled, + deadline=connect_deadline, + ): if should_abort_handshake( handshake_event, self.client._stop_event, this_generation, lambda: generation, - CONNECT_TIMEOUT_SECONDS, + deadline, ): - self.client.logger.warning(f"Websocket connect timed out after {CONNECT_TIMEOUT_SECONDS}s") - abort_websocket_app(app) + with generation_lock: + if this_generation != generation: + return + self.client.logger.warning(f"Websocket connect timed out after {CONNECT_TIMEOUT_SECONDS}s") + cancelled.set() + close_websocket_app(attempt_app) self.client.logger.debug("QuoteReceiver ready") threading.Thread(target=watchdog, daemon=True).start() if not self.enabled or self.client._stop_event.is_set(): - abort_websocket_app(app) + close_websocket_app(app) break try: - run_forever_with_connect_timeout( - app, - CONNECT_TIMEOUT_SECONDS, + app.run_forever( skip_utf8_validation=True, ) # skip_utf8_validation for more performance except Exception as e: self.client.logger.error(f"Websocket ERROR: {repr(e)}") - self.client.ready = False + with self.client._channels_lock: + self.client.ready = False if not self.enabled or self.client._stop_event.is_set(): break self.client.logger.info("Websocket closed, reconnecting...") skip_wait = ( - session_opened - and (time.time() - session_opened_at) >= STABLE_SESSION_SECONDS + session["opened"] + and (time.monotonic() - session["opened_at"]) >= STABLE_SESSION_SECONDS ) if skip_wait: self.client.last_self_heal_backoff = -1 @@ -490,7 +607,7 @@ def watchdog(this_generation=generation, handshake_event=handshake_event): def on_open(self, ws): self.client.logger.info("Websocket opened!") - self.client.on_connect() + self.client.on_connect(ws, self) def on_close(self, ws, code, message): self.client.logger.info("Websocket closed!") @@ -506,48 +623,6 @@ def on_error(self, ws, error, *args): self.client.logger.error(f"Error in on_error handler: {repr(e)}; {repr(error)}") raise e - def stitch(self): - full = None - while not self.continuation_queue.empty(): - partial = self.continuation_queue.get(True, 1) - if full is None: - full = partial - else: - full = full + partial - return full - - def on_cont_message(self, ws, partial_message, is_last): # The 3rd argument is continue flag. if 0, the data continue - try: - if DEBUGGING: # This is here for performance reasons so we don't use slow reflection on every message. - if isinstance(partial_message, str): - self.client.logger.debug(f"Received partial message (str): {partial_message.encode('utf-8').hex()}") - else: - if isinstance(partial_message, bytes): - self.client.logger.debug(f"Received partial message (hex): {partial_message.hex()}") - #self.client.logger.debug(f"Received partial message (hex): {partial_message.hex()}") - - self.continuation_lock.acquire() - try: - if is_last == 0: - self.continuation_queue.put(partial_message) - else: - self.continuation_queue.put(partial_message) - full_message = self.stitch() - self.on_message(ws, full_message) - finally: - self.continuation_lock.release() - except queue.Full: - self.client.on_queue_full() - except Exception as e: - hex_message = "" - if isinstance(partial_message, str): - hex_message = partial_message.encode('utf-8').hex() - else: - if isinstance(partial_message, bytes): - hex_message = partial_message.hex() - self.client.logger.error(f"Websocket on_message ERROR. Message as hex: {hex_message}; error: {repr(e)}") - raise e - def on_message(self, ws, message): try: if DEBUGGING: # This is here for performance reasons so we don't use slow reflection on every message. @@ -576,6 +651,7 @@ def __init__(self, client, bypass_parsing: bool): self.daemon = True self.client = client self.bypass_parsing = bypass_parsing + self._stop_event = threading.Event() self.subprovider_codes = { 0: NO_SUBPROVIDER, 1: CTA_A, @@ -588,6 +664,17 @@ def __init__(self, client, bypass_parsing: bool): 8: EQUITIES_EDGE } + @property + def stopped(self) -> bool: + return self._stop_event.is_set() + + def stop(self): + self._stop_event.set() + try: + self.client.quotes.put_nowait(None) + except queue.Full: + pass + def parse_quote(self, quote_bytes: bytes, start_index: int = 0) -> EquitiesQuote: buffer = memoryview(quote_bytes) symbol_length = buffer[start_index + 2] @@ -632,7 +719,7 @@ def parse_message(self, message_bytes: bytes, start_index: int, backlog_len: int if callable(self.client.on_trade): try: if self.bypass_parsing: - item = message_bytes[start_index:new_start_index - 1] + item = message_bytes[start_index:new_start_index] else: item = self.parse_trade(message_bytes, start_index) self.client.on_trade(item, backlog_len) @@ -642,7 +729,7 @@ def parse_message(self, message_bytes: bytes, start_index: int, backlog_len: int if callable(self.client.on_quote): try: if self.bypass_parsing: - item = message_bytes[start_index:new_start_index - 1] + item = message_bytes[start_index:new_start_index] else: item = self.parse_quote(message_bytes, start_index) self.client.on_quote(item, backlog_len) @@ -653,11 +740,27 @@ def parse_message(self, message_bytes: bytes, start_index: int, backlog_len: int def run(self): self.client.logger.debug("QuoteHandler ready") - while True: + while not self._stop_event.is_set(): message = self.client.quotes.get() + if message is None or self._stop_event.is_set(): + break backlog_len = self.client.quotes.qsize() - if message is not None and len(message) > 0 and len(message) >= message[0] * 24: #sanity check on length. Should be at least as long as the smallest message times the number of messages it says it has. + if not isinstance(message, (bytes, bytearray, memoryview)): + self.client.logger.error("Ignored non-binary websocket message") + continue + try: + if len(message) == 0 or len(message) < 1 + message[0] * 24: + raise ValueError("Equities message is shorter than its item count") items_in_message = message[0] start_index = 1 for i in range(0, items_in_message): + if self._stop_event.is_set(): + break + if start_index + 2 > len(message): + raise ValueError("Equities message item header is truncated") + message_length = message[start_index + 1] + if message_length < 24 or start_index + message_length > len(message): + raise ValueError("Equities message item length is invalid") start_index = self.parse_message(message, start_index, backlog_len) + except Exception as error: + self.client.logger.error(f"Ignored malformed equities message: {repr(error)}") diff --git a/intriniorealtime/options_client.py b/intriniorealtime/options_client.py index bbf5f80..a93c2ba 100644 --- a/intriniorealtime/options_client.py +++ b/intriniorealtime/options_client.py @@ -9,7 +9,7 @@ from collections.abc import Callable from enum import IntEnum, unique -from ._websocket import abort_websocket_app, run_forever_with_connect_timeout, should_abort_handshake +from ._websocket import close_websocket_app, should_abort_handshake _SELF_HEAL_BACKOFFS = [10, 30, 60, 300, 600] _CONNECT_TIMEOUT_SECONDS = 30 @@ -22,10 +22,6 @@ _NAN = float("NAN") _stopFlag: threading.Event = threading.Event() -_dataMsgLock: threading.Lock = threading.Lock() -_dataMsgCount: int = 0 -_txtMsgLock: threading.Lock = threading.Lock() -_txtMsgCount: int = 0 _logHandler: logging.Logger = logging.StreamHandler() _logHandler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')) _log: logging.Logger = logging.getLogger('intrinio_realtime_options') @@ -280,10 +276,16 @@ def _get_option_mask(use_on_trade: bool, use_on_quote: bool, use_on_refresh: boo mask |= 0b1000 return mask -class _WebSocket(websocket.WebSocketApp): +class _WebSocket: + on_cont_message = None + def __init__(self, ws_url: str, ws_lock: threading.Lock, + channels_lock: threading.RLock, + sent_channels: set[str], + stats: dict[str, int], + stats_lock: threading.Lock, worker_threads: list[threading.Thread], get_channels: Callable[[], set], get_token: Callable[[], str], @@ -294,11 +296,12 @@ def __init__(self, use_on_ua: bool, data_queue: queue.Queue, stop_flag: threading.Event): - super().__init__(ws_url, on_open=self.__on_open, on_close=self.__on_close, on_data=self.__on_data, on_error=self.__on_error) + self.url: str = ws_url self.__wsLock: threading.Lock = ws_lock - self.__continuation_queue = queue.Queue(100) - self.__continuation_lock: threading.Lock = threading.Lock() - self.__currently_continuing: bool = False + self.__channels_lock: threading.RLock = channels_lock + self.__sent_channels: set[str] = sent_channels + self.__stats: dict[str, int] = stats + self.__stats_lock: threading.Lock = stats_lock self.__worker_threads: list[threading.Thread] = worker_threads self.__get_channels: Callable[[], set] = get_channels self.__get_token: Callable[[], str] = get_token @@ -312,53 +315,84 @@ def __init__(self, self.__last_reset: float = time.time() self.isReady: bool = False self.__handshake_event: threading.Event = threading.Event() + self.__attempt_cancelled: threading.Event = threading.Event() + self.__connect_deadline: float = float("inf") self.__backoff_index: int = 0 self.__connect_generation: int = 0 + self.__connect_lock: threading.Lock = threading.Lock() + self.__app_lock: threading.Lock = threading.Lock() + self.__app = None self.__session_opened: bool = False self.__session_opened_at: float = 0.0 self.__stop_flag: threading.Event = stop_flag def __close_leftover_sock(self): - abort_websocket_app(self) - - def __clear_continuation_state(self): - with self.__continuation_lock: - self.__currently_continuing = False - while True: - try: - self.__continuation_queue.get_nowait() - except queue.Empty: - break + with self.__app_lock: + app = self.__app + close_websocket_app(app) def __on_open(self, ws): - self.__handshake_event.set() - self.__session_opened = True - self.__session_opened_at = time.time() + rejected = False + with self.__connect_lock: + if (self.__attempt_cancelled.is_set() + or self.__stop_flag.is_set() + or time.monotonic() >= self.__connect_deadline): + self.__attempt_cancelled.set() + rejected = True + else: + self.__session_opened = True + self.__session_opened_at = time.monotonic() + self.__handshake_event.set() + if rejected: + close_websocket_app(ws) + return _log.info("Websocket - Connected") - self.__wsLock.acquire() - try: - self.isReady = True - self.__is_reconnecting = False - for worker in self.__worker_threads: - if not worker.is_alive(): - worker.start() - finally: - self.__wsLock.release() - if self.__get_channels and callable(self.__get_channels): - channels: set[str] = self.__get_channels() - if channels and (len(channels) > 0): - for symbol in channels: - symbol_bytes = bytes(symbol, 'utf-8') - message: bytes = bytearray(len(symbol_bytes) + 2) - message[0] = 74 # join code - message[1] = _get_option_mask(self.__use_on_trade, self.__use_on_quote, self.__use_on_refresh, self.__use_on_ua) - message[2:] = symbol_bytes - if self.isReady: - _log.info("Websocket - Joining channel: {0}".format(symbol)) - self.send_binary(message) + rejected = False + with self.__channels_lock: + if self.__attempt_cancelled.is_set() or self.__stop_flag.is_set(): + self.__attempt_cancelled.set() + rejected = True + else: + self.__sent_channels.clear() + self.__wsLock.acquire() + try: + self.isReady = True + self.__is_reconnecting = False + for worker in self.__worker_threads: + if not worker.is_alive(): + worker.start() + finally: + self.__wsLock.release() + if self.__get_channels and callable(self.__get_channels): + channels: set[str] = self.__get_channels() + if channels and (len(channels) > 0): + for symbol in channels: + symbol_bytes = bytes(symbol, 'utf-8') + message: bytes = bytearray(len(symbol_bytes) + 2) + message[0] = 74 # join code + message[1] = _get_option_mask(self.__use_on_trade, self.__use_on_quote, self.__use_on_refresh, self.__use_on_ua) + message[2:] = symbol_bytes + if self.isReady: + _log.info("Websocket - Joining channel: {0}".format(symbol)) + try: + self.send_binary(message) + self.__sent_channels.add(symbol) + except Exception: + self.__wsLock.acquire() + try: + self.isReady = False + finally: + self.__wsLock.release() + self.__sent_channels.clear() + close_websocket_app(ws) + raise + if rejected: + close_websocket_app(ws) def __on_close(self, ws, closeStatusCode, closeMsg): _log.info("Websocket - Closed - {0}: {1}".format(closeStatusCode, closeMsg)) + with self.__channels_lock: + self.__sent_channels.clear() self.__wsLock.acquire() try: self.isReady = False @@ -368,50 +402,29 @@ def __on_close(self, ws, closeStatusCode, closeMsg): def __on_error(self, ws, error): _log.error("Websocket - Error - {0}".format(error)) - def __stitch(self): - full = None - while not self.__continuation_queue.empty(): - partial = self.__continuation_queue.get(True, 1) - if full is None: - full = partial - else: - full = full + partial - return full - - def __on_data(self, ws, data, code, is_last): #continueFlag - If 0, the data continues + def __on_data(self, ws, data, code, is_last): if code == websocket.ABNF.OPCODE_BINARY: - with _dataMsgLock: - global _dataMsgCount - _dataMsgCount += 1 - - if self.__currently_continuing or is_last == 0: # we're either in the middle of a continue, or we're starting a continue - with self.__continuation_lock: - if self.__currently_continuing or is_last == 0: #check lock check - #_log.info(f"Received partial message (hex): {data.hex()}") - self.__continuation_queue.put(data) - if is_last != 0: - full_message = self.__stitch() - self.__data_queue.put(full_message) - self.__currently_continuing = True if is_last == 0 else False # we're in the middle of a continue, but this is the last message, so remove flag - else: # We're not in the middle of a continue, and this isn't marked as multi-part, so this is a full message by itself. - self.__data_queue.put(data) - else: # We're not in the middle of a continue, and this isn't marked as multi-part, so this is a full message by itself. - self.__data_queue.put(data) + with self.__stats_lock: + self.__stats["data"] += 1 + self.__data_queue.put(data) else: _log.debug("Websocket - Message received") - with _txtMsgLock: - global _txtMsgCount - _txtMsgCount += 1 + with self.__stats_lock: + self.__stats["text"] += 1 _log.error("Error received: {0}".format(data)) def start(self): while not self.__stop_flag.is_set(): handshake_event = threading.Event() self.__handshake_event = handshake_event - self.__connect_generation += 1 - generation = self.__connect_generation + attempt_cancelled = threading.Event() + self.__attempt_cancelled = attempt_cancelled + connect_deadline = time.monotonic() + _CONNECT_TIMEOUT_SECONDS + self.__connect_deadline = connect_deadline + with self.__connect_lock: + self.__connect_generation += 1 + generation = self.__connect_generation self.__session_opened = False - self.__clear_continuation_state() self.__wsLock.acquire() try: self.isReady = False @@ -419,29 +432,53 @@ def start(self): self.__wsLock.release() self.__close_leftover_sock() - def watchdog(this_generation=generation, handshake_event=handshake_event): + # Leave on_cont_message unset so websocket-client aggregates fragmented messages. + app = websocket.WebSocketApp( + self.url, + on_open=self.__on_open, + on_close=self.__on_close, + on_data=self.__on_data, + on_error=self.__on_error, + ) + with self.__app_lock: + if self.__stop_flag.is_set(): + close_websocket_app(app) + return + self.__app = app + + def watchdog( + this_generation=generation, + handshake_event=handshake_event, + cancelled=attempt_cancelled, + deadline=connect_deadline, + attempt_app=app, + ): if should_abort_handshake( handshake_event, self.__stop_flag, this_generation, lambda: self.__connect_generation, - _CONNECT_TIMEOUT_SECONDS, + deadline, ): - _log.warning("Websocket - Connect timed out after {0}s".format(_CONNECT_TIMEOUT_SECONDS)) - self.__close_leftover_sock() + with self.__connect_lock: + if this_generation != self.__connect_generation: + return + _log.warning("Websocket - Connect timed out after {0}s".format(_CONNECT_TIMEOUT_SECONDS)) + cancelled.set() + close_websocket_app(attempt_app) threading.Thread(target=watchdog, daemon=True).start() if self.__stop_flag.is_set(): - abort_websocket_app(self) + close_websocket_app(app) return try: - run_forever_with_connect_timeout( - self, - _CONNECT_TIMEOUT_SECONDS, - skip_utf8_validation=True, - ) + app.run_forever(skip_utf8_validation=True) except Exception as e: _log.warning("Websocket - Attempt failed: {0}".format(e)) + finally: + with self.__app_lock: + if self.__app is app: + self.__app = None self.__wsLock.acquire() try: @@ -455,7 +492,7 @@ def watchdog(this_generation=generation, handshake_event=handshake_event): _log.info("Websocket - Reconnecting...") skip_wait: bool = ( self.__session_opened - and (time.time() - self.__session_opened_at) >= _STABLE_SESSION_SECONDS + and (time.monotonic() - self.__session_opened_at) >= _STABLE_SESSION_SECONDS ) if skip_wait: self.__backoff_index = 0 @@ -482,13 +519,32 @@ def watchdog(this_generation=generation, handshake_event=handshake_event): # super().run_forever(ping_interval = 5, ping_timeout = 2, skip_utf8_validation = True) def stop(self): - abort_websocket_app(self) + self.__close_leftover_sock() + + def close(self): + self.stop() + + # SDK-owned callback entry points; these do not expose or mutate + # websocket-client state. + def on_open(self, ws): + self.__on_open(ws) + + def on_data(self, ws, data, code, is_last): + self.__on_data(ws, data, code, is_last) def send(self, message: str): - super().send(message, websocket.ABNF.OPCODE_TEXT) + with self.__app_lock: + app = self.__app + if app is None: + raise RuntimeError("Websocket is not connected") + app.send(message, websocket.ABNF.OPCODE_TEXT) def send_binary(self, message: bytes): - super().send(message, websocket.ABNF.OPCODE_BINARY) + with self.__app_lock: + app = self.__app + if app is None: + raise RuntimeError("Websocket is not connected") + app.send(message, websocket.ABNF.OPCODE_BINARY) def reset(self): self.__last_reset = time.time() @@ -496,11 +552,13 @@ def reset(self): class Config: def __init__(self, api_key: str, provider: Providers, num_threads: int = 4, log_level: LogLevel = LogLevel.INFO, manual_ip_address: str = None, symbols: set[str] = None, delayed: bool = False): + if isinstance(num_threads, bool) or not isinstance(num_threads, int) or num_threads < 1: + raise ValueError("num_threads must be a positive integer") self.api_key: str = api_key self.provider: Providers = provider self.num_threads: int = num_threads self.manual_ip_address: str = manual_ip_address - self.symbols: list[str] = symbols + self.symbols: set[str] = symbols self.log_level: LogLevel = log_level self.delayed: bool = delayed @@ -605,6 +663,8 @@ def _thread_fn(index: int, data: queue.Queue, count = datum[0] start_index = 1 for _ in range(count): + if flag.is_set(): + break msg_type = datum[start_index + 22] if msg_type == 1: # Quote message: bytes = datum[start_index:(start_index + _OPTIONS_QUOTE_MESSAGE_SIZE)] @@ -713,7 +773,9 @@ def _thread_fn(index: int, data: queue.Queue, continue except Exception as e: _log.error(f"Worker thread {index} Exception {e}") - _log.error(f"\tCurrent count: {count}\r\n\tCurrent start_index: {start_index}\r\n\tCurrent msg_type: {msg_type}\r\n\tFull message (hex): {datum.hex()}\r\n\tScoped message: {message.hex()}") + datum_hex = datum.hex() if isinstance(datum, (bytes, bytearray)) else repr(datum) + message_hex = message.hex() if isinstance(message, (bytes, bytearray)) else repr(message) + _log.error(f"\tCurrent count: {count}\r\n\tCurrent start_index: {start_index}\r\n\tCurrent msg_type: {msg_type}\r\n\tFull message (hex): {datum_hex}\r\n\tScoped message: {message_hex}") continue _log.debug("Worker thread {0} stopped".format(index)) @@ -767,18 +829,26 @@ def __init__(self, config: Config, on_trade: Callable[[OptionsTrade], None], on_ self.__manualIP: str = config.manual_ip_address self.__token: tuple[str, float] = (None, 0.0) self.__webSocket: _WebSocket = None - if config.symbols and (isinstance(config.symbols, list)) and (len(config.symbols) > 0): + if config.symbols and isinstance(config.symbols, (list, set)): self.__channels: set[str] = set((_transform_contract_to_new(symbol)) for symbol in config.symbols) else: self.__channels: set[str] = set() self.__data: queue.Queue = queue.Queue() self.__t_lock: threading.Lock = threading.Lock() self.__ws_lock: threading.Lock = threading.Lock() + self.__channels_lock: threading.RLock = threading.RLock() + self.__sent_channels: set[str] = set() + self.__stats: dict[str, int] = {"data": 0, "text": 0} + self.__stats_lock: threading.Lock = threading.Lock() self.__lifecycle_lock: threading.RLock = threading.RLock() self.__on_trade = on_trade self.__on_quote = on_quote self.__on_refresh = on_refresh self.__on_unusual_activity = on_unusual_activity + if (isinstance(config.num_threads, bool) + or not isinstance(config.num_threads, int) + or config.num_threads < 1): + raise ValueError("num_threads must be a positive integer") self.__num_threads: int = config.num_threads self._stop_event: threading.Event = threading.Event() self._stop_generation: int = 0 @@ -787,6 +857,13 @@ def __init__(self, config: Config, on_trade: Callable[[OptionsTrade], None], on_ self.__is_started: bool = False self.__is_starting: bool = False self.__is_stopping: bool = False + self.__start_complete: threading.Event = threading.Event() + self.__start_complete.set() + self.__start_result = {"error": None} + self.__stopping_thread = None + self.__stop_complete: threading.Event = threading.Event() + self.__stop_complete.set() + self.__stop_result = {"error": None} _log.setLevel(config.log_level) def __make_worker_threads(self) -> list[threading.Thread]: @@ -860,61 +937,71 @@ def __get_token(self) -> str: finally: self.__t_lock.release() + def __refresh_token(self) -> str: + self.__t_lock.acquire() + try: + self.__token = (None, 0.0) + finally: + self.__t_lock.release() + return self.__get_token() + def __get_channels(self) -> set[str]: - return self.__channels + with self.__channels_lock: + return self.__channels.copy() def __join(self, symbol: str): transformed_symbol: str = _transform_contract_to_new(symbol) - if transformed_symbol not in self.__channels: - self.__channels.add(transformed_symbol) - symbol_bytes = bytes(transformed_symbol, 'utf-8') - message: bytes = bytearray(len(symbol_bytes)+2) - message[0] = 74 # join code - message[1] = _get_option_mask(self.__use_on_trade, self.__use_on_quote, self.__use_on_refresh, self.__use_on_unusual_activity) - message[2:] = symbol_bytes - ws = self.__webSocket - if ws is not None and ws.isReady: - _log.info("Websocket - Joining channel: {0}".format(transformed_symbol)) - ws.send_binary(message) + with self.__channels_lock: + if transformed_symbol not in self.__channels: + self.__channels.add(transformed_symbol) + if transformed_symbol not in self.__sent_channels: + symbol_bytes = bytes(transformed_symbol, 'utf-8') + message: bytes = bytearray(len(symbol_bytes)+2) + message[0] = 74 # join code + message[1] = _get_option_mask(self.__use_on_trade, self.__use_on_quote, self.__use_on_refresh, self.__use_on_unusual_activity) + message[2:] = symbol_bytes + ws = self.__webSocket + if ws is not None and ws.isReady: + _log.info("Websocket - Joining channel: {0}".format(transformed_symbol)) + ws.send_binary(message) + self.__sent_channels.add(transformed_symbol) def __leave(self, symbol: str, forget: bool = True): transformed_symbol: str = _transform_contract_to_new(symbol) - if transformed_symbol not in self.__channels: - return - symbol_bytes = bytes(transformed_symbol, 'utf-8') - message: bytes = bytearray(len(symbol_bytes) + 2) - message[0] = 76 # leave code - message[1] = _get_option_mask(self.__use_on_trade, self.__use_on_quote, self.__use_on_refresh, self.__use_on_unusual_activity) - message[2:] = symbol_bytes - try: - ws = self.__webSocket - if ws is not None and ws.isReady: - _log.info("Websocket - Leaving channel: {0}".format(transformed_symbol)) - ws.send_binary(message) - finally: - if forget: - self.__channels.discard(transformed_symbol) + with self.__channels_lock: + if transformed_symbol not in self.__channels: + return + symbol_bytes = bytes(transformed_symbol, 'utf-8') + message: bytes = bytearray(len(symbol_bytes) + 2) + message[0] = 76 # leave code + message[1] = _get_option_mask(self.__use_on_trade, self.__use_on_quote, self.__use_on_refresh, self.__use_on_unusual_activity) + message[2:] = symbol_bytes + try: + ws = self.__webSocket + if ws is not None and ws.isReady: + _log.info("Websocket - Leaving channel: {0}".format(transformed_symbol)) + ws.send_binary(message) + finally: + self.__sent_channels.discard(transformed_symbol) + if forget: + self.__channels.discard(transformed_symbol) def join(self, *symbols): - if self.__is_started: - while not self.__all_ready(): - time.sleep(1.0) for (symbol) in symbols: self.__join(symbol) def join_firehose(self): - if "$FIREHOSE" in self.__channels: - _log.warn("This client has already joined the firehose channel") - else: - if self.__is_started: - while not self.__all_ready(): - time.sleep(1.0) - self.__join("$FIREHOSE") + with self.__channels_lock: + if "$FIREHOSE" in self.__channels: + _log.warn("This client has already joined the firehose channel") + return + self.__join("$FIREHOSE") def leave(self, *symbols): if not symbols: _log.info("Leaving all channels") - channels: set[str] = self.__channels.copy() + with self.__channels_lock: + channels: set[str] = self.__channels.copy() for (symbol) in channels: self.__leave(symbol) symbol_set: set[str] = set(symbols) @@ -922,7 +1009,9 @@ def leave(self, *symbols): self.__leave(sym) def leave_firehose(self): - if "$FIREHOSE" in self.__channels: + with self.__channels_lock: + joined = "$FIREHOSE" in self.__channels + if joined: self.__leave("$FIREHOSE") def __socket_start_fn(self, token: str, stop_event: threading.Event, @@ -936,9 +1025,13 @@ def __socket_start_fn(self, token: str, stop_event: threading.Event, return web_socket = _WebSocket(ws_url, self.__ws_lock, + self.__channels_lock, + self.__sent_channels, + self.__stats, + self.__stats_lock, worker_threads, self.__get_channels, - self.__get_token, + self.__refresh_token, self.__get_web_socket_url, self.__use_on_trade, self.__use_on_quote, @@ -958,26 +1051,45 @@ def start(self): and self.__socket_thread.is_alive() and not self._stop_event.is_set()): return - if self.__is_starting: - return if self.__is_stopping: - return - self.__is_starting = True - stop_generation = self._stop_generation + raise RuntimeError("Cannot start while client is stopping") + if self.__is_starting: + start_complete = self.__start_complete + start_result = self.__start_result + wait_for_start = True + else: + wait_for_start = False + self.__is_starting = True + start_complete = threading.Event() + start_result = {"error": None} + self.__start_complete = start_complete + self.__start_result = start_result + stop_generation = self._stop_generation + if wait_for_start: + start_complete.wait() + if start_result["error"] is not None: + raise RuntimeError( + "Concurrent start failed: {0}".format(start_result["error"]) + ) from start_result["error"] + return + + failure = None try: while True: leftover = self.__socket_thread if leftover is not None and leftover.is_alive(): - abort_websocket_app(self.__webSocket) + close_websocket_app(self.__webSocket) leftover.join(timeout=_CONNECT_TIMEOUT_SECONDS + 1) if leftover.is_alive(): with self.__lifecycle_lock: if self._stop_generation != stop_generation: - return - continue + raise RuntimeError("Cannot start while client is stopping") + raise RuntimeError("Previous socket thread did not exit") + if any(worker.is_alive() for worker in self.__worker_threads): + raise RuntimeError("Previous worker threads did not exit") with self.__lifecycle_lock: if self._stop_generation != stop_generation: - return + raise RuntimeError("Cannot start while client is stopping") self._stop_event = threading.Event() stop_event = self._stop_event self.__data = queue.Queue() @@ -986,16 +1098,16 @@ def start(self): worker_threads = self.__worker_threads token: str = self.__get_token() if stop_event.is_set(): - return + raise RuntimeError("Cannot start while client is stopping") if not token: if stop_event.wait(timeout=_SELF_HEAL_BACKOFFS[0]): - return + raise RuntimeError("Cannot start while client is stopping") continue with self.__lifecycle_lock: if (self._stop_generation != stop_generation or self._stop_event is not stop_event or stop_event.is_set()): - return + raise RuntimeError("Cannot start while client is stopping") socket_thread = threading.Thread( group=None, target=self.__socket_start_fn, @@ -1010,43 +1122,73 @@ def start(self): except Exception: self.__is_started = False raise + if (self._stop_generation != stop_generation + or self._stop_event is not stop_event + or stop_event.is_set()): + self.__is_started = False + raise RuntimeError("Cannot start while client is stopping") return + except BaseException as error: + failure = error + raise finally: with self.__lifecycle_lock: self.__is_starting = False + start_result["error"] = failure + start_complete.set() def stop(self): + current_thread = threading.current_thread() with self.__lifecycle_lock: if self.__is_stopping: - return - self.__is_stopping = True - _log.info("Stopping...") - self._stop_generation += 1 - self._stop_event.set() - ws = self.__webSocket - worker_threads = self.__worker_threads - socket_thread = self.__socket_thread + if self.__stopping_thread is current_thread: + return + stop_complete = self.__stop_complete + stop_result = self.__stop_result + wait_for_stop = True + else: + wait_for_stop = False + self.__is_stopping = True + self.__stopping_thread = current_thread + stop_complete = threading.Event() + stop_result = {"error": None} + self.__stop_complete = stop_complete + self.__stop_result = stop_result + _log.info("Stopping...") + self._stop_generation += 1 + self._stop_event.set() + ws = self.__webSocket + worker_threads = self.__worker_threads + socket_thread = self.__socket_thread + shutdown_deadline = time.monotonic() + _CONNECT_TIMEOUT_SECONDS + 1 + if wait_for_stop: + stop_complete.wait() + if stop_result["error"] is not None: + raise RuntimeError(stop_result["error"]) + return + + candidate_threads = [ + worker for worker in worker_threads + if worker.ident is not None and worker is not current_thread + ] + if (socket_thread is not None + and socket_thread.ident is not None + and socket_thread is not current_thread): + candidate_threads.append(socket_thread) + surviving_threads = [] + failure = None try: - if ws is not None: - try: - sock = getattr(ws, "sock", None) - if sock is not None: - sock.settimeout(1) - except Exception: - pass - if len(self.__channels) > 0: - channels: set[str] = self.__channels.copy() - for symbol in channels: - try: - self.__leave(symbol, forget=False) - except Exception as e: - _log.warning("Websocket - Leave failed: {0}".format(e)) - break - if ws is not None: - ws.stop() - time.sleep(1.0) - finally: try: + if ws is not None: + self.__ws_lock.acquire() + try: + ws.isReady = False + finally: + self.__ws_lock.release() + ws.stop() + with self.__channels_lock: + self.__sent_channels.clear() + finally: self.__ws_lock.acquire() try: if ws is not None: @@ -1057,17 +1199,37 @@ def stop(self): ws.stop() for i in range(len(worker_threads)): worker = worker_threads[i] - if worker.ident is not None: - worker.join(timeout=_CONNECT_TIMEOUT_SECONDS + 1) - _log.debug("Worker thread {0} joined".format(i)) - if socket_thread is not None and socket_thread.ident is not None: - socket_thread.join(timeout=_CONNECT_TIMEOUT_SECONDS + 1) - _log.debug("Socket thread joined") - finally: - with self.__lifecycle_lock: - self.__is_started = False - self.__is_stopping = False + if worker.ident is not None and worker is not current_thread: + worker.join(timeout=max(0.0, shutdown_deadline - time.monotonic())) + if (socket_thread is not None + and socket_thread.ident is not None + and socket_thread is not current_thread): + socket_thread.join(timeout=max(0.0, shutdown_deadline - time.monotonic())) + + surviving_threads = [thread for thread in candidate_threads if thread.is_alive()] + except BaseException as error: + failure = error + finally: + # Recheck after all joins and cleanup so threads that exited while a + # later thread was being joined are not retained as false survivors. + surviving_threads = [thread for thread in candidate_threads if thread.is_alive()] + if failure is None and surviving_threads: + failure = RuntimeError("Client stop incomplete: threads did not exit") + with self.__lifecycle_lock: + self.__is_started = bool(surviving_threads) + self.__is_stopping = False + self.__stopping_thread = None + stop_result["error"] = str(failure) if failure is not None else None + stop_complete.set() + if failure is not None: + _log.error("Stop incomplete: {0}".format(failure)) + else: _log.info("Stopped") + if failure is not None: + raise failure def get_stats(self) -> tuple[int, int, int]: - return _dataMsgCount, _txtMsgCount, self.__data.qsize() + with self.__stats_lock: + data_count = self.__stats["data"] + text_count = self.__stats["text"] + return data_count, text_count, self.__data.qsize() diff --git a/tests/test_equities_client.py b/tests/test_equities_client.py index 3122da9..6b454ef 100644 --- a/tests/test_equities_client.py +++ b/tests/test_equities_client.py @@ -1,8 +1,13 @@ import logging +import queue +import threading import time import unittest +from types import SimpleNamespace from unittest.mock import Mock, patch +import websocket + from intriniorealtime.equities_client import ( CONNECT_TIMEOUT_SECONDS, EquitiesQuoteHandler, @@ -53,6 +58,16 @@ def test_on_connect_does_not_skip_join_when_joined_channels_are_stale(self): client.on_connect() self.assertEqual(client.ws.send.call_count, 2) + def test_on_connect_does_not_restore_ready_after_disconnect(self): + client = _make_client() + client.ws = Mock() + client._stop_event.set() + + client.on_connect() + + self.assertFalse(client.ready) + client.ws.send.assert_not_called() + def test_watchdog_aborts_app_when_handshake_times_out(self): client = _make_client() client.token = "tok" @@ -64,7 +79,7 @@ def fake_run_forever(app, *args, **kwargs): receiver.enabled = False client._stop_event.set() - with patch("intriniorealtime.equities_client.abort_websocket_app", side_effect=lambda app: aborted.append(app)), \ + with patch("intriniorealtime.equities_client.close_websocket_app", side_effect=lambda app: aborted.append(app)), \ patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", 0.05), \ patch("websocket.setdefaulttimeout"), \ patch("websocket.WebSocketApp.run_forever", fake_run_forever): @@ -72,14 +87,12 @@ def fake_run_forever(app, *args, **kwargs): self.assertTrue(aborted) - def test_run_sets_connect_timeout_and_clears_it_on_open(self): + def test_run_does_not_change_websocket_client_global_timeout(self): client = _make_client() client.token = "tok" receiver = EquitiesQuoteReceiver(client) - sock = Mock() def fake_run_forever(app, *args, **kwargs): - app.sock = sock app.on_open(app) receiver.enabled = False client._stop_event.set() @@ -88,9 +101,45 @@ def fake_run_forever(app, *args, **kwargs): patch("websocket.WebSocketApp.run_forever", fake_run_forever): receiver.run() - self.assertIn(((CONNECT_TIMEOUT_SECONDS,), {}), set_timeout.call_args_list) - set_timeout.assert_called_with(None) - sock.settimeout.assert_called_with(None) + set_timeout.assert_not_called() + + def test_late_open_after_disconnect_is_aborted_before_becoming_ready(self): + client = _make_client() + client.token = "tok" + receiver = EquitiesQuoteReceiver(client) + aborted = [] + + def fake_run_forever(app, *args, **kwargs): + receiver.enabled = False + client._stop_event.set() + app.on_open(app) + + with patch("websocket.WebSocketApp.run_forever", fake_run_forever), \ + patch("intriniorealtime.equities_client.close_websocket_app", side_effect=aborted.append): + receiver.run() + + self.assertFalse(client.ready) + self.assertTrue(aborted) + + def test_late_open_after_timeout_is_aborted_before_becoming_ready(self): + client = _make_client() + client.token = "tok" + receiver = EquitiesQuoteReceiver(client) + aborted = [] + + def fake_run_forever(app, *args, **kwargs): + time.sleep(0.08) + app.on_open(app) + receiver.enabled = False + + with patch("websocket.WebSocketApp.run_forever", fake_run_forever), \ + patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", 0.03), \ + patch("intriniorealtime.equities_client.should_abort_handshake", return_value=False), \ + patch("intriniorealtime.equities_client.close_websocket_app", side_effect=aborted.append): + receiver.run() + + self.assertFalse(client.ready) + self.assertTrue(aborted) def test_refresh_token_failure_backs_off_and_skips_socket_until_auth_succeeds(self): client = _make_client() @@ -116,7 +165,7 @@ def fake_run_forever(app, *args, **kwargs): with patch("websocket.setdefaulttimeout"), \ patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", 0.05), \ patch("websocket.WebSocketApp.run_forever", fake_run_forever), \ - patch("intriniorealtime.equities_client.abort_websocket_app"): + patch("intriniorealtime.equities_client.close_websocket_app"): receiver.run() self.assertEqual(len(run_forever_calls), 2) @@ -137,7 +186,7 @@ def join_and_exit(timeout=None): client.quote_receiver = old client.ws = Mock() - with patch("intriniorealtime.equities_client.abort_websocket_app"): + with patch("intriniorealtime.equities_client.close_websocket_app"): client.disconnect() self.assertFalse(old.enabled) @@ -154,6 +203,131 @@ def join_and_exit(timeout=None): self.assertIs(client.quote_receiver, new_receiver) self.assertFalse(client._stop_event.is_set()) + def test_quote_handler_stops_before_dispatching_rest_of_batch(self): + messages = queue.Queue() + callbacks = [] + fake_client = SimpleNamespace( + logger=Mock(), + quotes=messages, + on_quote=None, + on_trade=None, + ) + handler = EquitiesQuoteHandler(fake_client, bypass_parsing=True) + + def on_trade(item, backlog): + callbacks.append(item) + handler.stop() + + fake_client.on_trade = on_trade + datum = bytearray(49) + datum[0] = 2 + datum[1] = 0 + datum[2] = 24 + datum[25] = 0 + datum[26] = 24 + messages.put(bytes(datum)) + + handler.start() + handler.join(timeout=1) + + self.assertFalse(handler.is_alive()) + self.assertEqual(len(callbacks), 1) + + def test_bypass_parsing_returns_complete_raw_record(self): + for message_type, callback_name in ((0, "on_trade"), (1, "on_quote")): + with self.subTest(message_type=message_type): + callback = Mock() + fake_client = SimpleNamespace( + logger=Mock(), + on_quote=None, + on_trade=None, + ) + setattr(fake_client, callback_name, callback) + handler = EquitiesQuoteHandler(fake_client, bypass_parsing=True) + record = bytearray(range(24)) + record[0] = message_type + record[1] = len(record) + + next_index = handler.parse_message(bytes(record), 0, 7) + + self.assertEqual(next_index, len(record)) + callback.assert_called_once_with(bytes(record), 7) + + def test_quote_handler_ignores_text_and_malformed_binary_messages(self): + client = _make_client() + client.on_trade = Mock() + handler = EquitiesQuoteHandler(client, client.bypass_parsing) + client.quote_handler = handler + handler.start() + client.quotes.put("server text message") + malformed = bytearray(48) + malformed[0] = 2 + malformed[2] = 255 + client.quotes.put(bytes(malformed)) + undersized = bytearray(25) + undersized[0] = 1 + undersized[2] = 1 + client.quotes.put(bytes(undersized)) + + time.sleep(0.05) + + self.assertTrue(handler.is_alive()) + client.on_trade.assert_not_called() + client.disconnect() + + def test_connect_replaces_stopped_quote_handler_and_queue(self): + client = _make_client() + old_handler = client.quote_handler + old_queue = client.quotes + old_handler.stop() + client.refresh_token = Mock() + receiver = Mock() + + with patch.object(EquitiesQuoteHandler, "start") as start, \ + patch("intriniorealtime.equities_client.EquitiesQuoteReceiver", return_value=receiver): + client.connect() + + self.assertIsNot(client.quote_handler, old_handler) + self.assertIsNot(client.quotes, old_queue) + start.assert_called_once() + receiver.start.assert_called_once() + + def test_disconnect_and_reconnect_from_callback_replaces_current_handler(self): + callback_finished = threading.Event() + client_holder = {} + + def on_trade(item, backlog): + client = client_holder["client"] + client.disconnect() + client.connect() + callback_finished.set() + + client = IntrinioRealtimeEquitiesClient( + { + "api_key": "test-key", + "provider": "IEX", + "bypass_parsing": True, + }, + on_trade, + lambda *args: None, + ) + client_holder["client"] = client + original_handler = client.quote_handler + client.refresh_token = Mock() + client.refresh_websocket = Mock() + datum = bytearray(25) + datum[0] = 1 + datum[1] = 0 + datum[2] = 24 + + client.quotes.put(bytes(datum)) + + self.assertTrue(callback_finished.wait(timeout=1)) + original_handler.join(timeout=1) + self.assertFalse(original_handler.is_alive()) + self.assertIsNot(client.quote_handler, original_handler) + client.disconnect() + def test_refresh_websocket_replaces_disabled_live_receiver(self): client = _make_client() old = Mock() @@ -210,6 +384,278 @@ def test_connect_while_receiver_is_starting_does_not_close_socket(self): client.refresh_token.assert_not_called() self.assertIs(client.quote_receiver, receiver) + def test_concurrent_connects_start_only_one_receiver(self): + client = _make_client() + client.refresh_quote_handler = Mock() + client.refresh_token = Mock() + start_entered = threading.Event() + release_start = threading.Event() + receivers = [] + + class Receiver: + enabled = True + + def __init__(self, owner): + receivers.append(self) + + def start(self): + start_entered.set() + release_start.wait(timeout=1) + + def is_alive(self): + return True + + with patch("intriniorealtime.equities_client.EquitiesQuoteReceiver", Receiver): + first = threading.Thread(target=client.connect) + second = threading.Thread(target=client.connect) + first.start() + self.assertTrue(start_entered.wait(timeout=1)) + second.start() + time.sleep(0.05) + self.assertEqual(len(receivers), 1) + release_start.set() + first.join(timeout=1) + second.join(timeout=1) + + self.assertFalse(first.is_alive()) + self.assertFalse(second.is_alive()) + self.assertEqual(len(receivers), 1) + + def test_connect_raises_while_disconnect_is_in_progress(self): + client = _make_client() + client.refresh_quote_handler = Mock() + client.refresh_token = Mock() + old_join_entered = threading.Event() + release_disconnect = threading.Event() + + class OldReceiver: + enabled = True + alive = True + join_calls = 0 + + def is_alive(self): + return self.alive + + def join(self, timeout=None): + self.join_calls += 1 + if self.join_calls == 1: + old_join_entered.set() + release_disconnect.wait(timeout=2) + else: + self.alive = False + + old_receiver = OldReceiver() + new_receiver = Mock() + new_receiver.enabled = True + new_receiver.is_alive.return_value = True + client.quote_receiver = old_receiver + client.ws = Mock() + + with patch("intriniorealtime.equities_client.EquitiesQuoteReceiver", return_value=new_receiver), \ + patch("intriniorealtime.equities_client.close_websocket_app"): + disconnect_thread = threading.Thread(target=client.disconnect) + disconnect_thread.start() + self.assertTrue(old_join_entered.wait(timeout=1)) + + with self.assertRaisesRegex(RuntimeError, "Cannot connect while client is disconnecting"): + client.connect() + release_disconnect.set() + disconnect_thread.join(timeout=1) + + client.connect() + + self.assertFalse(disconnect_thread.is_alive()) + self.assertIs(client.quote_receiver, new_receiver) + + def test_disconnect_clears_websocket_published_late_by_current_receiver(self): + client = _make_client() + late_websocket = Mock() + + class Receiver: + enabled = True + + def join(self, timeout=None): + client.ws = late_websocket + + def is_alive(self): + return False + + receiver = Receiver() + client.quote_receiver = receiver + client.ws = None + + with patch("intriniorealtime.equities_client.close_websocket_app") as abort: + client.disconnect() + + self.assertIsNone(client.quote_receiver) + self.assertIsNone(client.ws) + abort.assert_called_with(late_websocket) + + def test_concurrent_joins_serialize_channel_refreshes(self): + client = _make_client() + client.ready = True + client.ws = Mock() + first_send_entered = threading.Event() + release_first_send = threading.Event() + + def send(_message, _opcode): + if client.ws.send.call_count == 1: + first_send_entered.set() + release_first_send.wait(timeout=1) + + client.ws.send.side_effect = send + first = threading.Thread(target=client.join, args=("AAPL",)) + second = threading.Thread(target=client.join, args=("MSFT",)) + + first.start() + self.assertTrue(first_send_entered.wait(timeout=1)) + second.start() + time.sleep(0.05) + self.assertEqual(client.ws.send.call_count, 1) + release_first_send.set() + first.join(timeout=1) + second.join(timeout=1) + + self.assertFalse(first.is_alive()) + self.assertFalse(second.is_alive()) + self.assertEqual(client.ws.send.call_count, 2) + self.assertEqual(client.joined_channels, {"AAPL", "MSFT"}) + + def test_disconnect_aborts_socket_before_waiting_for_blocked_join(self): + client = _make_client() + client.ready = True + websocket_app = Mock() + client.ws = websocket_app + send_entered = threading.Event() + release_send = threading.Event() + + def send(_message, _opcode): + send_entered.set() + release_send.wait(timeout=1) + + websocket_app.send.side_effect = send + joiner = threading.Thread(target=client.join, args=("AAPL",)) + disconnecter = threading.Thread(target=client.disconnect) + + with patch( + "intriniorealtime.equities_client.close_websocket_app", + side_effect=lambda app: release_send.set(), + ) as abort: + joiner.start() + self.assertTrue(send_entered.wait(timeout=1)) + disconnecter.start() + disconnecter.join(timeout=1) + joiner.join(timeout=1) + + self.assertFalse(disconnecter.is_alive()) + self.assertFalse(joiner.is_alive()) + abort.assert_any_call(websocket_app) + + def test_disconnect_aborts_socket_during_blocked_initial_channel_replay(self): + client = _make_client() + client.channels = {"AAPL"} + client.ws = Mock() + send_entered = threading.Event() + release_send = threading.Event() + + def send(*args, **kwargs): + send_entered.set() + release_send.wait(timeout=1) + + client.ws.send.side_effect = send + opener = threading.Thread(target=client.on_connect) + stopper = threading.Thread(target=client.disconnect) + + with patch( + "intriniorealtime.equities_client.close_websocket_app", + side_effect=lambda app: release_send.set(), + ) as abort: + opener.start() + self.assertTrue(send_entered.wait(timeout=1)) + stopper.start() + stopper.join(timeout=1) + opener.join(timeout=1) + + self.assertFalse(stopper.is_alive()) + self.assertFalse(opener.is_alive()) + abort.assert_called() + self.assertFalse(client.ready) + + def test_partial_channel_refresh_records_successful_sends(self): + client = _make_client() + client.ready = True + client.channels = {"AAPL", "MSFT"} + client.join_binary_message = Mock(side_effect=lambda channel: channel) + client.ws = Mock() + client.ws.send.side_effect = [None, RuntimeError("send failed")] + + with self.assertRaises(RuntimeError): + client.refresh_channels() + + successful_channel = client.ws.send.call_args_list[0].args[0] + failed_channel = client.ws.send.call_args_list[1].args[0] + self.assertEqual(client.joined_channels, {successful_channel}) + + client.ws.send.reset_mock(side_effect=True) + client.refresh_channels() + + client.ws.send.assert_called_once_with( + failed_channel, + websocket.ABNF.OPCODE_BINARY, + ) + self.assertEqual(client.joined_channels, {"AAPL", "MSFT"}) + + client.channels = set() + client.leave_binary_message = Mock(side_effect=lambda channel: channel) + client.ws.send.reset_mock(side_effect=True) + client.ws.send.side_effect = [None, RuntimeError("send failed")] + + with self.assertRaises(RuntimeError): + client.refresh_channels() + + successful_leave = client.ws.send.call_args_list[0].args[0] + failed_leave = client.ws.send.call_args_list[1].args[0] + self.assertNotIn(successful_leave, client.joined_channels) + self.assertEqual(client.joined_channels, {failed_leave}) + + client.ws.send.reset_mock(side_effect=True) + client.refresh_channels() + + client.ws.send.assert_called_once_with( + failed_leave, + websocket.ABNF.OPCODE_BINARY, + ) + self.assertEqual(client.joined_channels, set()) + + def test_connect_queued_before_disconnect_does_not_restart(self): + client = _make_client() + auth_entered = threading.Event() + auth_calls = [] + + def refresh_token(): + auth_calls.append(1) + auth_entered.set() + client._stop_event.wait(timeout=1) + + client.refresh_quote_handler = Mock() + client.refresh_token = Mock(side_effect=refresh_token) + client.refresh_websocket = Mock() + first = threading.Thread(target=client.connect) + second = threading.Thread(target=client.connect) + + first.start() + self.assertTrue(auth_entered.wait(timeout=1)) + second.start() + time.sleep(0.05) + client.disconnect() + first.join(timeout=1) + second.join(timeout=1) + + self.assertFalse(first.is_alive()) + self.assertFalse(second.is_alive()) + self.assertEqual(auth_calls, [1]) + client.refresh_websocket.assert_not_called() + def test_connect_resets_backoff_on_new_attempt(self): client = _make_client() client.last_self_heal_backoff = 4 @@ -230,44 +676,24 @@ def test_connect_healthy_noop_does_not_reset_backoff(self): client.connect() self.assertEqual(client.last_self_heal_backoff, 4) - def test_receiver_clears_continuation_queue_on_each_connect_attempt(self): + def test_receiver_relies_on_websocket_client_fragment_aggregation(self): client = _make_client() client.token = "tok" - client.do_backoff = Mock() - client.refresh_token = Mock() receiver = EquitiesQuoteReceiver(client) - receiver.continuation_queue.put(b"partial-from-prior-session") - seen_empty = [] - run_count = {"n": 0} + continuation_callbacks = [] def fake_run_forever(app, *args, **kwargs): - run_count["n"] += 1 - seen_empty.append(receiver.continuation_queue.empty()) - if run_count["n"] == 1: - receiver.continuation_queue.put(b"leftover-mid-fragment") - return + continuation_callbacks.append(app.on_cont_message) receiver.enabled = False client._stop_event.set() with patch("websocket.setdefaulttimeout"), \ patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", 0.05), \ patch("websocket.WebSocketApp.run_forever", fake_run_forever), \ - patch("intriniorealtime.equities_client.abort_websocket_app"): + patch("intriniorealtime.equities_client.close_websocket_app"): receiver.run() - self.assertEqual(run_count["n"], 2) - self.assertEqual(seen_empty, [True, True]) - - def test_fragmented_message_is_reassembled_and_delivered(self): - client = _make_client() - receiver = EquitiesQuoteReceiver(client) - ws = Mock() - receiver.on_message = Mock() - - receiver.on_cont_message(ws, b"abc", 0) - receiver.on_cont_message(ws, b"def", 1) - - receiver.on_message.assert_called_once_with(ws, b"abcdef") + self.assertEqual(continuation_callbacks, [None]) def test_refresh_token_passes_timeout(self): client = _make_client() @@ -298,7 +724,7 @@ def test_disconnect_keeps_receiver_if_join_times_out(self): client.ws = Mock() ws = client.ws - with patch("intriniorealtime.equities_client.abort_websocket_app") as abort: + with patch("intriniorealtime.equities_client.close_websocket_app") as abort: client.disconnect() self.assertFalse(old.enabled) @@ -316,7 +742,7 @@ def test_refresh_websocket_does_not_replace_receiver_still_alive_after_join(self client.quote_receiver = old client.ws = Mock() - with patch("intriniorealtime.equities_client.abort_websocket_app") as abort, \ + with patch("intriniorealtime.equities_client.close_websocket_app") as abort, \ patch("intriniorealtime.equities_client.EquitiesQuoteReceiver") as ctor: with self.assertRaises(RuntimeError): client.refresh_websocket() @@ -345,7 +771,7 @@ def fake_run_forever(app, *args, **kwargs): patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", 0.05), \ patch("intriniorealtime.equities_client.STABLE_SESSION_SECONDS", 10), \ patch("websocket.WebSocketApp.run_forever", fake_run_forever), \ - patch("intriniorealtime.equities_client.abort_websocket_app"): + patch("intriniorealtime.equities_client.close_websocket_app"): receiver.run() self.assertEqual(run_count["n"], 2) @@ -372,7 +798,7 @@ def fake_run_forever(app, *args, **kwargs): patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", 0.05), \ patch("intriniorealtime.equities_client.STABLE_SESSION_SECONDS", 0.05), \ patch("websocket.WebSocketApp.run_forever", fake_run_forever), \ - patch("intriniorealtime.equities_client.abort_websocket_app"): + patch("intriniorealtime.equities_client.close_websocket_app"): receiver.run() self.assertEqual(run_count["n"], 2) diff --git a/tests/test_options_client.py b/tests/test_options_client.py index a373454..d77fdd2 100644 --- a/tests/test_options_client.py +++ b/tests/test_options_client.py @@ -5,19 +5,27 @@ import unittest from unittest.mock import Mock, patch +import websocket + from intriniorealtime.options_client import ( Config, IntrinioRealtimeOptionsClient, Providers, _CONNECT_TIMEOUT_SECONDS, _WebSocket, + _thread_fn, ) -def _make_ws(stop_flag, get_token=None, get_url=None, channels=None): +def _make_ws(stop_flag, get_token=None, get_url=None, channels=None, data_queue=None, + channels_lock=None, stats=None, stats_lock=None, sent_channels=None): return _WebSocket( "ws://example/old", threading.Lock(), + channels_lock or threading.RLock(), + sent_channels if sent_channels is not None else set(), + stats if stats is not None else {"data": 0, "text": 0}, + stats_lock or threading.Lock(), [], lambda: channels if channels is not None else set(), get_token or (lambda: "tok"), @@ -26,7 +34,7 @@ def _make_ws(stop_flag, get_token=None, get_url=None, channels=None): False, False, False, - queue.Queue(), + data_queue if data_queue is not None else queue.Queue(), stop_flag, ) @@ -48,6 +56,43 @@ def setUp(self): def tearDown(self): logging.disable(logging.NOTSET) + def test_config_rejects_non_positive_or_non_integer_worker_counts(self): + for num_threads in (0, -1, True, 1.5, "2"): + with self.subTest(num_threads=num_threads): + with self.assertRaises(ValueError): + Config( + api_key="test-key", + provider=Providers.MANUAL, + num_threads=num_threads, + manual_ip_address="127.0.0.1", + ) + + config = Config( + api_key="test-key", + provider=Providers.MANUAL, + num_threads=1, + manual_ip_address="127.0.0.1", + ) + config.num_threads = 0 + with self.assertRaises(ValueError): + IntrinioRealtimeOptionsClient(config, on_trade=lambda trade: None) + + def test_config_symbol_set_initializes_channels(self): + config = Config( + api_key="test-key", + provider=Providers.MANUAL, + num_threads=1, + manual_ip_address="127.0.0.1", + symbols={"AAPL__220101C00140000", "MSFT"}, + ) + + client = IntrinioRealtimeOptionsClient(config, on_trade=lambda trade: None) + + self.assertEqual( + client._IntrinioRealtimeOptionsClient__channels, + {"AAPL_220101C140.00", "MSFT"}, + ) + def test_stale_watchdog_does_not_abort_later_attempt(self): stop = threading.Event() second_opened = threading.Event() @@ -71,7 +116,7 @@ def fake_run_forever(app, *args, **kwargs): with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0]), \ patch("websocket.setdefaulttimeout"), \ - patch("intriniorealtime.options_client.abort_websocket_app", side_effect=track_abort), \ + patch("intriniorealtime.options_client.close_websocket_app", side_effect=track_abort), \ patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): ws.start() @@ -97,7 +142,7 @@ def fake_run_forever(app, *args, **kwargs): with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0]), \ patch("websocket.setdefaulttimeout"), \ - patch("intriniorealtime.options_client.abort_websocket_app"), \ + patch("intriniorealtime.options_client.close_websocket_app"), \ patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): ws.start() @@ -128,7 +173,7 @@ def fake_run_forever(app, *args, **kwargs): with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0]), \ patch("websocket.setdefaulttimeout"), \ - patch("intriniorealtime.options_client.abort_websocket_app"), \ + patch("intriniorealtime.options_client.close_websocket_app"), \ patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): ws.start() @@ -189,57 +234,208 @@ def fake_run_forever(app, *args, **kwargs): with patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever), \ patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ patch("websocket.setdefaulttimeout"), \ - patch("intriniorealtime.options_client.abort_websocket_app"): + patch("intriniorealtime.options_client.close_websocket_app"): ws.start() self.assertEqual(calls, [1]) - def test_clears_continuation_state_on_each_connect_attempt(self): + def test_late_open_after_stop_is_aborted_before_becoming_ready(self): + stop = threading.Event() + stop.set() + ws = _make_ws(stop) + + with patch("intriniorealtime.options_client.close_websocket_app") as abort: + ws.on_open(ws) + + self.assertFalse(ws.isReady) + abort.assert_called_once_with(ws) + + def test_late_open_after_timeout_is_aborted_before_becoming_ready(self): stop = threading.Event() ws = _make_ws(stop) - ws._WebSocket__currently_continuing = True - ws._WebSocket__continuation_queue.put(b"partial-from-prior-session") - seen = [] - run_count = {"n": 0} def fake_run_forever(app, *args, **kwargs): - run_count["n"] += 1 - seen.append(( - app._WebSocket__currently_continuing, - app._WebSocket__continuation_queue.empty(), - )) - if run_count["n"] == 1: - app._WebSocket__currently_continuing = True - app._WebSocket__continuation_queue.put(b"leftover-mid-fragment") - return + time.sleep(0.08) + app.on_open(app) stop.set() - with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ - patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0]), \ - patch("websocket.setdefaulttimeout"), \ - patch("intriniorealtime.options_client.abort_websocket_app"), \ + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.03), \ + patch("intriniorealtime.options_client.should_abort_handshake", return_value=False), \ + patch("intriniorealtime.options_client.close_websocket_app") as abort, \ patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): ws.start() - self.assertEqual(run_count["n"], 2) - self.assertEqual(seen, [(False, True), (False, True)]) + self.assertFalse(ws.isReady) + self.assertTrue(abort.called) + + def test_relies_on_websocket_client_fragment_aggregation(self): + stop = threading.Event() + data_queue = queue.Queue() + ws = _make_ws(stop, data_queue=data_queue) + + self.assertIsNone(ws.on_cont_message) + ws.on_data(ws, b"complete-message", websocket.ABNF.OPCODE_BINARY, True) + self.assertEqual(data_queue.get_nowait(), b"complete-message") + + def test_open_does_not_publish_ready_before_channel_replay_lock(self): + class GateLock: + def __init__(self): + self.entered = threading.Event() + self.release = threading.Event() + + def __enter__(self): + self.entered.set() + self.release.wait(timeout=1) + return self + + def __exit__(self, exc_type, exc_value, traceback): + return False + + stop = threading.Event() + channels_lock = GateLock() + ws = _make_ws(stop, channels_lock=channels_lock) + opener = threading.Thread(target=ws.on_open, args=(ws,)) + + opener.start() + self.assertTrue(channels_lock.entered.wait(timeout=1)) + self.assertFalse(ws.isReady) + channels_lock.release.set() + opener.join(timeout=1) + + self.assertFalse(opener.is_alive()) + self.assertTrue(ws.isReady) + + def test_open_rechecks_stop_before_publishing_ready(self): + class GateLock: + def __init__(self): + self.entered = threading.Event() + self.release = threading.Event() + + def __enter__(self): + self.entered.set() + self.release.wait(timeout=1) + return self + + def __exit__(self, exc_type, exc_value, traceback): + return False + + stop = threading.Event() + channels_lock = GateLock() + ws = _make_ws(stop, channels_lock=channels_lock) + opener = threading.Thread(target=ws.on_open, args=(ws,)) + + opener.start() + self.assertTrue(channels_lock.entered.wait(timeout=1)) + stop.set() + channels_lock.release.set() + opener.join(timeout=1) + + self.assertFalse(opener.is_alive()) + self.assertFalse(ws.isReady) + + def test_open_aborts_connection_when_channel_replay_fails(self): + stop = threading.Event() + sent_channels = set() + ws = _make_ws( + stop, + channels={"AAPL"}, + sent_channels=sent_channels, + ) + ws.send_binary = Mock(side_effect=RuntimeError("send failed")) + + with patch("intriniorealtime.options_client.close_websocket_app") as abort: + with self.assertRaises(RuntimeError): + ws.on_open(ws) + + self.assertFalse(ws.isReady) + self.assertEqual(sent_channels, set()) + abort.assert_called_once_with(ws) + + def test_get_channels_returns_snapshot(self): + client = _make_client() + client._IntrinioRealtimeOptionsClient__channels = {"AAPL"} + + snapshot = client._IntrinioRealtimeOptionsClient__get_channels() + snapshot.add("MSFT") + + self.assertEqual(client._IntrinioRealtimeOptionsClient__channels, {"AAPL"}) + + def test_stats_are_isolated_between_clients(self): + first = _make_client() + second = _make_client() + first_data = first._IntrinioRealtimeOptionsClient__data + first_stats = first._IntrinioRealtimeOptionsClient__stats + first_stats_lock = first._IntrinioRealtimeOptionsClient__stats_lock + ws = _make_ws( + threading.Event(), + data_queue=first_data, + stats=first_stats, + stats_lock=first_stats_lock, + ) + + ws.on_data(ws, b"message", websocket.ABNF.OPCODE_BINARY, True) + + self.assertEqual(first.get_stats(), (1, 0, 1)) + self.assertEqual(second.get_stats(), (0, 0, 0)) + + def test_worker_stops_before_dispatching_rest_of_batch(self): + data = queue.Queue() + stop = threading.Event() + callbacks = [] + + def on_quote(quote): + callbacks.append(quote) + stop.set() + + datum = bytearray(105) + datum[0] = 2 + datum[1 + 22] = 1 + datum[1 + 52 + 22] = 1 + data.put(bytes(datum)) + + worker = threading.Thread( + target=_thread_fn, + args=(0, data, None, on_quote, None, None, stop), + ) + worker.start() + worker.join(timeout=1) + + self.assertFalse(worker.is_alive()) + self.assertEqual(len(callbacks), 1) + + def test_truncated_message_does_not_terminate_worker(self): + data = queue.Queue() + stop = threading.Event() + worker = threading.Thread( + target=_thread_fn, + args=(0, data, None, None, None, None, stop), + ) + worker.start() + + data.put(b"\x01") + time.sleep(0.05) + + self.assertTrue(worker.is_alive()) + stop.set() + worker.join(timeout=2) + self.assertFalse(worker.is_alive()) def test_stop_preserves_channels_and_clears_started(self): client = _make_client() client._IntrinioRealtimeOptionsClient__channels = {"AAPL"} + client._IntrinioRealtimeOptionsClient__sent_channels = {"AAPL"} client._IntrinioRealtimeOptionsClient__is_started = True ws = Mock() ws.isReady = True client._IntrinioRealtimeOptionsClient__webSocket = ws client._IntrinioRealtimeOptionsClient__socket_thread = None - with patch("intriniorealtime.options_client.time.sleep"): - client.stop() + client.stop() self.assertEqual(client._IntrinioRealtimeOptionsClient__channels, {"AAPL"}) + self.assertEqual(client._IntrinioRealtimeOptionsClient__sent_channels, set()) self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) - ws.send_binary.assert_called() - leave = ws.send_binary.call_args[0][0] - self.assertEqual(leave[0], 76) + ws.send_binary.assert_not_called() + ws.stop.assert_called() def test_join_after_stop_does_not_wait_for_ready(self): client = _make_client() @@ -253,6 +449,31 @@ def test_join_after_stop_does_not_wait_for_ready(self): ready.assert_not_called() self.assertIn("MSFT", client._IntrinioRealtimeOptionsClient__channels) + def test_join_does_not_wait_while_connection_is_unavailable(self): + client = _make_client() + client._IntrinioRealtimeOptionsClient__is_started = True + client._IntrinioRealtimeOptionsClient__webSocket = None + + with patch.object(client, "_IntrinioRealtimeOptionsClient__all_ready") as ready: + client.join("MSFT") + client.join_firehose() + + ready.assert_not_called() + self.assertIn("MSFT", client._IntrinioRealtimeOptionsClient__channels) + self.assertIn("$FIREHOSE", client._IntrinioRealtimeOptionsClient__channels) + + def test_reconnect_token_refresh_invalidates_cached_token(self): + client = _make_client() + client._IntrinioRealtimeOptionsClient__token = ("cached", time.time()) + + with patch("intriniorealtime.options_client.requests.get") as get: + get.return_value.status_code = 200 + get.return_value.text = "fresh" + token = client._IntrinioRealtimeOptionsClient__refresh_token() + + self.assertEqual(token, "fresh") + get.assert_called_once() + def test_leave_drops_channel_if_send_raises(self): client = _make_client() client._IntrinioRealtimeOptionsClient__channels = {"AAPL"} @@ -266,15 +487,47 @@ def test_leave_drops_channel_if_send_raises(self): self.assertNotIn("AAPL", client._IntrinioRealtimeOptionsClient__channels) + def test_join_retries_send_after_previous_send_raises(self): + client = _make_client() + ws = Mock() + ws.isReady = True + ws.send_binary.side_effect = RuntimeError("connection closed") + client._IntrinioRealtimeOptionsClient__webSocket = ws + + with self.assertRaises(RuntimeError): + client.join("AAPL") + + ws.send_binary.side_effect = None + client.join("AAPL") + + self.assertIn("AAPL", client._IntrinioRealtimeOptionsClient__channels) + self.assertEqual(ws.send_binary.call_count, 2) + def test_stop_joins_socket_thread_with_timeout(self): client = _make_client() thread = Mock() thread.ident = 1 + thread.is_alive.return_value = False client._IntrinioRealtimeOptionsClient__socket_thread = thread client._IntrinioRealtimeOptionsClient__webSocket = None with patch("intriniorealtime.options_client.time.sleep"): client.stop() - thread.join.assert_called_with(timeout=_CONNECT_TIMEOUT_SECONDS + 1) + timeout = thread.join.call_args.kwargs["timeout"] + self.assertGreater(timeout, 0) + self.assertLessEqual(timeout, _CONNECT_TIMEOUT_SECONDS + 1) + + def test_stop_from_worker_callback_does_not_join_current_thread(self): + client = _make_client() + current = threading.current_thread() + client._IntrinioRealtimeOptionsClient__worker_threads = [current] + client._IntrinioRealtimeOptionsClient__socket_thread = None + client._IntrinioRealtimeOptionsClient__webSocket = None + + with patch("intriniorealtime.options_client.time.sleep"): + client.stop() + + self.assertTrue(client._stop_event.is_set()) + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) def test_open_then_drop_twice_backs_off_on_second_reconnect(self): stop = threading.Event() @@ -298,7 +551,7 @@ def fake_run_forever(app, *args, **kwargs): patch("intriniorealtime.options_client._STABLE_SESSION_SECONDS", 10), \ patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0.01]), \ patch("websocket.setdefaulttimeout"), \ - patch("intriniorealtime.options_client.abort_websocket_app"), \ + patch("intriniorealtime.options_client.close_websocket_app"), \ patch("threading.Event.wait", tracking_wait), \ patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): ws.start() @@ -330,7 +583,7 @@ def fake_run_forever(app, *args, **kwargs): patch("intriniorealtime.options_client._STABLE_SESSION_SECONDS", 0.05), \ patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [10]), \ patch("websocket.setdefaulttimeout"), \ - patch("intriniorealtime.options_client.abort_websocket_app"), \ + patch("intriniorealtime.options_client.close_websocket_app"), \ patch("threading.Event.wait", tracking_wait), \ patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): ws.start() @@ -359,7 +612,7 @@ def fake_run_forever(app, *args, **kwargs): with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.05), \ patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0.01]), \ patch("websocket.setdefaulttimeout"), \ - patch("intriniorealtime.options_client.abort_websocket_app"), \ + patch("intriniorealtime.options_client.close_websocket_app"), \ patch("threading.Event.wait", tracking_wait), \ patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever): ws.start() @@ -367,18 +620,48 @@ def fake_run_forever(app, *args, **kwargs): self.assertEqual(run_count["n"], 2) self.assertIn(0.01, waits) - def test_stop_sets_flag_before_sleep(self): + def test_stop_sets_flag_before_socket_abort(self): client = _make_client() saw_stop = {} + ws = Mock() + ws.isReady = True + client._IntrinioRealtimeOptionsClient__webSocket = ws - def fake_sleep(seconds): + def stop(): saw_stop["set"] = client._stop_event.is_set() - with patch("intriniorealtime.options_client.time.sleep", side_effect=fake_sleep): - client.stop() + ws.stop.side_effect = stop + client.stop() self.assertTrue(saw_stop["set"]) - def test_start_does_not_spawn_thread_when_stopped_during_auth(self): + def test_stop_aborts_socket_before_waiting_for_blocked_join(self): + client = _make_client() + ws = Mock() + ws.isReady = True + send_entered = threading.Event() + release_send = threading.Event() + + def send_binary(_message): + send_entered.set() + release_send.wait(timeout=1) + + ws.send_binary.side_effect = send_binary + ws.stop.side_effect = release_send.set + client._IntrinioRealtimeOptionsClient__webSocket = ws + joiner = threading.Thread(target=client.join, args=("AAPL",)) + stopper = threading.Thread(target=client.stop) + + joiner.start() + self.assertTrue(send_entered.wait(timeout=1)) + stopper.start() + stopper.join(timeout=1) + joiner.join(timeout=1) + + self.assertFalse(stopper.is_alive()) + self.assertFalse(joiner.is_alive()) + ws.stop.assert_called() + + def test_start_raises_if_stopped_during_auth(self): client = _make_client() def get_token(): @@ -387,7 +670,8 @@ def get_token(): with patch.object(client, "_IntrinioRealtimeOptionsClient__get_token", side_effect=get_token), \ patch.object(threading.Thread, "start") as thread_start: - client.start() + with self.assertRaisesRegex(RuntimeError, "Cannot start while client is stopping"): + client.start() thread_start.assert_not_called() self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) @@ -415,7 +699,7 @@ def fake_ws_start(self): self.assertEqual(start_calls, [1]) self.assertTrue(client._IntrinioRealtimeOptionsClient__is_started) - def test_stop_during_socket_thread_start_does_not_publish_started_state(self): + def test_stop_during_socket_thread_start_raises_lifecycle_error(self): client = _make_client() def stop_before_thread_runs(): @@ -424,7 +708,8 @@ def stop_before_thread_runs(): with patch.object(client, "_IntrinioRealtimeOptionsClient__get_token", return_value="tok"), \ patch.object(threading.Thread, "start", side_effect=stop_before_thread_runs), \ patch("intriniorealtime.options_client.time.sleep"): - client.start() + with self.assertRaisesRegex(RuntimeError, "Cannot start while client is stopping"): + client.start() self.assertTrue(client._stop_event.is_set()) self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) @@ -460,7 +745,7 @@ def test_start_waits_for_previous_socket_thread_then_starts(self): client = _make_client() leftover = Mock() leftover.ident = 1 - leftover.is_alive.side_effect = [True, True, False] + leftover.is_alive.side_effect = [True, False] leftover.join = Mock() client._IntrinioRealtimeOptionsClient__socket_thread = leftover client._IntrinioRealtimeOptionsClient__is_started = False @@ -479,23 +764,80 @@ def fake_ws_start(self): leftover.join.assert_called() self.assertEqual(start_calls, [1]) - def test_stop_completes_teardown_if_leave_raises(self): + def test_start_fails_when_previous_socket_thread_does_not_exit(self): + client = _make_client() + leftover = Mock() + leftover.ident = 1 + leftover.is_alive.return_value = True + client._IntrinioRealtimeOptionsClient__socket_thread = leftover + client._IntrinioRealtimeOptionsClient__is_started = False + client._stop_event.set() + + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0), \ + patch("intriniorealtime.options_client.close_websocket_app"): + with self.assertRaisesRegex(RuntimeError, "Previous socket thread"): + client.start() + + leftover.join.assert_called_once_with(timeout=1) + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_starting) + + def test_start_raises_clear_error_while_stop_is_in_progress(self): + client = _make_client() + client._IntrinioRealtimeOptionsClient__is_stopping = True + + with self.assertRaisesRegex(RuntimeError, "Cannot start while client is stopping"): + client.start() + + def test_concurrent_start_waits_and_reports_active_start_failure(self): + client = _make_client() + token_entered = threading.Event() + release_token = threading.Event() + errors = [] + + def fail_token(): + token_entered.set() + release_token.wait(timeout=1) + raise RuntimeError("auth failed") + + def start_client(): + try: + client.start() + except Exception as error: + errors.append(str(error)) + + with patch.object(client, "_IntrinioRealtimeOptionsClient__get_token", side_effect=fail_token): + first = threading.Thread(target=start_client) + second = threading.Thread(target=start_client) + first.start() + self.assertTrue(token_entered.wait(timeout=1)) + second.start() + second.join(timeout=0.05) + self.assertTrue(second.is_alive()) + + release_token.set() + first.join(timeout=1) + second.join(timeout=1) + + self.assertFalse(first.is_alive()) + self.assertFalse(second.is_alive()) + self.assertEqual(len(errors), 2) + self.assertTrue(any(error == "auth failed" for error in errors)) + self.assertTrue(any(error.startswith("Concurrent start failed: auth failed") for error in errors)) + + def test_stop_completes_teardown_with_active_socket(self): client = _make_client() client._IntrinioRealtimeOptionsClient__channels = {"AAPL"} client._IntrinioRealtimeOptionsClient__is_started = True ws = Mock() ws.isReady = True - ws.sock = Mock() - ws.send_binary.side_effect = RuntimeError("connection closed") client._IntrinioRealtimeOptionsClient__webSocket = ws thread = Mock() thread.ident = 1 + thread.is_alive.return_value = False client._IntrinioRealtimeOptionsClient__socket_thread = thread - with patch("intriniorealtime.options_client.time.sleep"): - client.stop() + client.stop() - ws.sock.settimeout.assert_called_with(1) ws.stop.assert_called() thread.join.assert_called() self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) @@ -506,6 +848,7 @@ def test_start_uses_new_stop_event_so_old_workers_stay_stopped(self): old_event.set() old_worker = Mock() old_worker.ident = 1 + old_worker.is_alive.return_value = False client._IntrinioRealtimeOptionsClient__worker_threads = [old_worker] client._IntrinioRealtimeOptionsClient__is_started = False start_calls = [] @@ -525,11 +868,109 @@ def fake_ws_start(self): self.assertNotIn(old_worker, client._IntrinioRealtimeOptionsClient__worker_threads) self.assertEqual(start_calls, [1]) + def test_stop_reports_workers_that_do_not_exit(self): + client = _make_client() + worker = Mock() + worker.ident = 1 + worker.is_alive.return_value = True + client._IntrinioRealtimeOptionsClient__worker_threads = [worker] + client._IntrinioRealtimeOptionsClient__socket_thread = None + + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0): + with self.assertRaisesRegex(RuntimeError, "stop incomplete"): + client.stop() + + self.assertTrue(client._IntrinioRealtimeOptionsClient__is_started) + with self.assertRaisesRegex(RuntimeError, "Previous worker threads"): + client.start() + + def test_stop_rechecks_survivors_after_all_joins(self): + client = _make_client() + worker = Mock() + worker.ident = 1 + worker.is_alive.side_effect = [True, False] + client._IntrinioRealtimeOptionsClient__worker_threads = [worker] + client._IntrinioRealtimeOptionsClient__socket_thread = None + + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0): + client.stop() + + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) + + def test_concurrent_stop_waits_for_active_stop(self): + client = _make_client() + join_entered = threading.Event() + release_join = threading.Event() + worker_alive = {"value": True} + worker = Mock() + worker.ident = 1 + worker.is_alive.side_effect = lambda: worker_alive["value"] + + def join(timeout=None): + join_entered.set() + release_join.wait(timeout=1) + worker_alive["value"] = False + + worker.join.side_effect = join + client._IntrinioRealtimeOptionsClient__worker_threads = [worker] + client._IntrinioRealtimeOptionsClient__socket_thread = None + failures = [] + + def stop_client(): + try: + client.stop() + except Exception as error: + failures.append(error) + + first = threading.Thread(target=stop_client) + second = threading.Thread(target=stop_client) + first.start() + self.assertTrue(join_entered.wait(timeout=1)) + second.start() + second.join(timeout=0.05) + self.assertTrue(second.is_alive()) + + release_join.set() + first.join(timeout=1) + second.join(timeout=1) + + self.assertFalse(first.is_alive()) + self.assertFalse(second.is_alive()) + self.assertEqual(failures, []) + + def test_stop_uses_one_deadline_for_all_thread_joins(self): + client = _make_client() + clock = {"now": 100.0} + workers = [] + + for _ in range(2): + worker = Mock() + worker.ident = 1 + worker.is_alive.return_value = False + + def join(timeout=None, clock=clock): + clock["now"] += 0.6 + + worker.join.side_effect = join + workers.append(worker) + + client._IntrinioRealtimeOptionsClient__worker_threads = workers + client._IntrinioRealtimeOptionsClient__socket_thread = None + + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0), \ + patch("intriniorealtime.options_client.time.monotonic", side_effect=lambda: clock["now"]): + client.stop() + + first_timeout = workers[0].join.call_args.kwargs["timeout"] + second_timeout = workers[1].join.call_args.kwargs["timeout"] + self.assertEqual(first_timeout, 1.0) + self.assertAlmostEqual(second_timeout, 0.4) + def test_start_aborts_leftover_websocket_while_joining(self): client = _make_client() leftover = Mock() leftover.ident = 1 - leftover.is_alive.side_effect = [True, True, False] + leftover.is_alive.side_effect = [True, False] leftover.join = Mock() ws = Mock() client._IntrinioRealtimeOptionsClient__socket_thread = leftover @@ -541,7 +982,7 @@ def test_start_aborts_leftover_websocket_while_joining(self): def fake_ws_start(self): start_calls.append(1) - with patch("intriniorealtime.options_client.abort_websocket_app") as abort, \ + with patch("intriniorealtime.options_client.close_websocket_app") as abort, \ patch("intriniorealtime.options_client.requests.get") as get, \ patch.object(_WebSocket, "start", fake_ws_start): get.return_value.status_code = 200 @@ -552,7 +993,7 @@ def fake_ws_start(self): leftover.join.assert_called() self.assertEqual(start_calls, [1]) - def test_start_returns_if_stop_during_leftover_join(self): + def test_start_raises_if_stop_occurs_during_leftover_join(self): client = _make_client() leftover = Mock() leftover.ident = 1 @@ -567,10 +1008,11 @@ def join(timeout=None): client._IntrinioRealtimeOptionsClient__is_started = False client._stop_event.set() - with patch("intriniorealtime.options_client.abort_websocket_app"), \ + with patch("intriniorealtime.options_client.close_websocket_app"), \ patch.object(client, "_IntrinioRealtimeOptionsClient__get_token") as get_token, \ patch.object(threading.Thread, "start") as thread_start: - client.start() + with self.assertRaisesRegex(RuntimeError, "Cannot start while client is stopping"): + client.start() get_token.assert_not_called() thread_start.assert_not_called() diff --git a/tests/test_websocket_abort.py b/tests/test_websocket_abort.py index f38bc16..be232d5 100644 --- a/tests/test_websocket_abort.py +++ b/tests/test_websocket_abort.py @@ -1,122 +1,60 @@ -import socket import threading import time import unittest -from unittest.mock import Mock, call +from unittest.mock import Mock -import websocket +from intriniorealtime._websocket import close_websocket_app, should_abort_handshake -from intriniorealtime._websocket import ( - abort_websocket_app, - run_forever_with_connect_timeout, - should_abort_handshake, -) +class CloseWebsocketAppTests(unittest.TestCase): + def test_uses_only_public_close_method(self): + app = Mock(spec=["close"]) -class FakeWebSocket: - def __init__(self, raw, connected=False): - self.sock = raw - self.connected = connected + close_websocket_app(app) - def close(self, **kwargs): - if not self.connected: - return - self.shutdown() - - def shutdown(self): - if self.sock: - self.sock.close() - self.sock = None - self.connected = False - - -class FakeApp: - def __init__(self, ws): - self.sock = ws - self.keep_running = True - - def close(self, **kwargs): - self.keep_running = False - if self.sock: - self.sock.close() - self.sock = None - - -class AbortWebsocketAppTests(unittest.TestCase): - def test_shutdowns_raw_socket_when_websocket_is_not_connected(self): - raw = Mock() - ws = FakeWebSocket(raw, connected=False) - app = FakeApp(ws) - - abort_websocket_app(app) - - self.assertFalse(app.keep_running) - raw.shutdown.assert_called_with(socket.SHUT_RDWR) - raw.close.assert_called() - self.assertEqual(raw.mock_calls[0], call.shutdown(socket.SHUT_RDWR)) - self.assertEqual(raw.mock_calls[1], call.close()) - self.assertIsNone(app.sock) - self.assertIsNone(ws.sock) - self.assertFalse(ws.connected) + app.close.assert_called_once_with() def test_noop_on_none(self): - abort_websocket_app(None) - - def test_sets_keep_running_false_without_sock(self): - app = FakeApp(None) - abort_websocket_app(app) - self.assertFalse(app.keep_running) - - -class ConnectTimeoutTests(unittest.TestCase): - def test_connect_timeout_is_restored_as_soon_as_socket_opens(self): - original_timeout = websocket.getdefaulttimeout() - websocket.setdefaulttimeout(17) - observed = [] - sock = Mock() - app = Mock() - app.sock = sock - - def on_open(ws): - observed.append(websocket.getdefaulttimeout()) - - def run_forever(**kwargs): - self.assertEqual(websocket.getdefaulttimeout(), 3) - app.on_open(app) + close_websocket_app(None) - app.on_open = on_open - app.run_forever.side_effect = run_forever - try: - run_forever_with_connect_timeout(app, 3) - self.assertEqual(observed, [17]) - self.assertEqual(websocket.getdefaulttimeout(), 17) - sock.settimeout.assert_called_once_with(None) - finally: - websocket.setdefaulttimeout(original_timeout) + def test_swallows_close_error(self): + app = Mock(spec=["close"]) + app.close.side_effect = RuntimeError("close failed") + close_websocket_app(app) class ShouldAbortHandshakeTests(unittest.TestCase): def test_true_on_timeout_for_current_generation(self): handshake = threading.Event() stop = threading.Event() - self.assertTrue(should_abort_handshake(handshake, stop, 1, 1, 0.01)) + self.assertTrue(should_abort_handshake( + handshake, stop, 1, 1, time.monotonic() + 0.01 + )) def test_false_when_handshake_completes(self): handshake = threading.Event() handshake.set() stop = threading.Event() - self.assertFalse(should_abort_handshake(handshake, stop, 1, 1, 0.01)) + self.assertFalse(should_abort_handshake( + handshake, stop, 1, 1, time.monotonic() + 0.01 + )) def test_false_for_stale_generation(self): handshake = threading.Event() stop = threading.Event() - self.assertFalse(should_abort_handshake(handshake, stop, 1, 2, 0.01)) + self.assertFalse(should_abort_handshake( + handshake, stop, 1, 2, time.monotonic() + 0.01 + )) - def test_true_when_stopped_and_handshake_timed_out(self): + def test_false_immediately_when_stopped(self): handshake = threading.Event() stop = threading.Event() stop.set() - self.assertTrue(should_abort_handshake(handshake, stop, 1, 1, 0.01)) + start = time.monotonic() + self.assertFalse(should_abort_handshake( + handshake, stop, 1, 1, time.monotonic() + 1.0 + )) + self.assertLess(time.monotonic() - start, 0.1) def test_callable_current_generation_is_read_after_wait(self): handshake = threading.Event() @@ -128,7 +66,9 @@ def bump(): state["g"] = 2 threading.Thread(target=bump, daemon=True).start() - self.assertFalse(should_abort_handshake(handshake, stop, 1, lambda: state["g"], 0.05)) + self.assertFalse(should_abort_handshake( + handshake, stop, 1, lambda: state["g"], time.monotonic() + 0.05 + )) if __name__ == "__main__": From 32232678d46e9f89adf0ffc545e6dcb5f0916100 Mon Sep 17 00:00:00 2001 From: Shawn Snyder Date: Mon, 28 Sep 2026 12:52:28 -0500 Subject: [PATCH 6/7] More cases --- intriniorealtime/equities_client.py | 61 ++++++++++-- intriniorealtime/options_client.py | 122 ++++++++++++++++++----- tests/test_equities_client.py | 123 ++++++++++++++++++++++-- tests/test_options_client.py | 144 ++++++++++++++++++++++++++-- 4 files changed, 403 insertions(+), 47 deletions(-) diff --git a/intriniorealtime/equities_client.py b/intriniorealtime/equities_client.py index 530186a..53a4d88 100644 --- a/intriniorealtime/equities_client.py +++ b/intriniorealtime/equities_client.py @@ -14,7 +14,7 @@ SELF_HEAL_BACKOFFS = [10, 30, 60, 300, 600] CONNECT_TIMEOUT_SECONDS = 30 -STABLE_SESSION_SECONDS = 1.0 +STABLE_SESSION_SECONDS = 60.0 _EMPTY_STRING = "" _NAN = float("NAN") REALTIME = "REALTIME" @@ -275,9 +275,26 @@ def connect(self): self.logger.error(f"Cannot connect: {repr(e)}") self.do_backoff() + def _finish_incomplete_disconnect(self, receiver, handler, survivors): + for thread in survivors: + thread.join() + with self._lifecycle_lock: + if (receiver is not None + and self.quote_receiver is receiver + and not receiver.is_alive()): + self.quote_receiver = None + self.ws = None + if (handler is not None + and self.quote_handler is handler + and not handler.is_alive()): + self.quote_handler = None + self._disconnecting = False + def disconnect(self): with self._disconnect_lock: with self._lifecycle_lock: + if self._disconnecting: + raise RuntimeError("Client disconnect is already in progress") self._disconnecting = True self._lifecycle_generation += 1 self._stop_event.set() @@ -286,6 +303,8 @@ def disconnect(self): websocket_app = self.ws if receiver is not None: receiver.enabled = False + shutdown_deadline = time.monotonic() + CONNECT_TIMEOUT_SECONDS + 1 + deferred_completion = False try: if handler is not None: handler.stop() @@ -297,7 +316,10 @@ def disconnect(self): self.joined_channels = set() if receiver is not None: - receiver.join(timeout=CONNECT_TIMEOUT_SECONDS + 1) + if (getattr(receiver, "ident", True) is not None + and receiver is not threading.current_thread()): + receiver.join(timeout=max( + 0.0, shutdown_deadline - time.monotonic())) if receiver.is_alive(): close_websocket_app(websocket_app) else: @@ -315,17 +337,40 @@ def disconnect(self): if ( handler is not None - and handler.ident is not None + and getattr(handler, "ident", True) is not None and handler is not threading.current_thread() ): - handler.join(timeout=CONNECT_TIMEOUT_SECONDS + 1) + handler.join(timeout=max( + 0.0, shutdown_deadline - time.monotonic())) if handler is not None and not handler.is_alive(): with self._lifecycle_lock: if self.quote_handler is handler: self.quote_handler = None + + current_thread = threading.current_thread() + non_current_survivors = [ + thread for thread in (receiver, handler) + if (thread is not None + and thread is not current_thread + and thread.is_alive()) + ] + if non_current_survivors: + deferred_survivors = list(non_current_survivors) + if (current_thread in (receiver, handler) + and current_thread.is_alive()): + deferred_survivors.append(current_thread) + deferred_completion = True + threading.Thread( + target=self._finish_incomplete_disconnect, + args=(receiver, handler, deferred_survivors), + daemon=True, + ).start() + raise RuntimeError( + "Client disconnect incomplete: threads did not exit") finally: - with self._lifecycle_lock: - self._disconnecting = False + if not deferred_completion: + with self._lifecycle_lock: + self._disconnecting = False def refresh_quote_handler(self): handler = self.quote_handler @@ -577,6 +622,10 @@ def watchdog( ) # skip_utf8_validation for more performance except Exception as e: self.client.logger.error(f"Websocket ERROR: {repr(e)}") + finally: + # The attempt is over even if the websocket never opened. Wake + # its watchdog so it cannot fire later during reconnect backoff. + handshake_event.set() with self.client._channels_lock: self.client.ready = False diff --git a/intriniorealtime/options_client.py b/intriniorealtime/options_client.py index a93c2ba..c451bf5 100644 --- a/intriniorealtime/options_client.py +++ b/intriniorealtime/options_client.py @@ -13,7 +13,7 @@ _SELF_HEAL_BACKOFFS = [10, 30, 60, 300, 600] _CONNECT_TIMEOUT_SECONDS = 30 -_STABLE_SESSION_SECONDS = 1.0 +_STABLE_SESSION_SECONDS = 60.0 _EMPTY_STRING = "" _OPTIONS_TRADE_MESSAGE_SIZE = 72 # 61 used + 11 pad _OPTIONS_QUOTE_MESSAGE_SIZE = 52 # 48 used + 4 pad @@ -295,7 +295,8 @@ def __init__(self, use_on_refresh: bool, use_on_ua: bool, data_queue: queue.Queue, - stop_flag: threading.Event): + stop_flag: threading.Event, + worker_factory: Callable[[int], threading.Thread] = None): self.url: str = ws_url self.__wsLock: threading.Lock = ws_lock self.__channels_lock: threading.RLock = channels_lock @@ -325,6 +326,7 @@ def __init__(self, self.__session_opened: bool = False self.__session_opened_at: float = 0.0 self.__stop_flag: threading.Event = stop_flag + self.__worker_factory = worker_factory def __close_leftover_sock(self): with self.__app_lock: @@ -356,11 +358,22 @@ def __on_open(self, ws): self.__sent_channels.clear() self.__wsLock.acquire() try: + for index, worker in enumerate(self.__worker_threads): + if worker.is_alive(): + continue + if worker.ident is not None: + if self.__worker_factory is None: + raise RuntimeError("Stopped worker thread cannot be restarted") + worker = self.__worker_factory(index) + self.__worker_threads[index] = worker + worker.start() self.isReady = True self.__is_reconnecting = False - for worker in self.__worker_threads: - if not worker.is_alive(): - worker.start() + except Exception: + self.isReady = False + self.__sent_channels.clear() + close_websocket_app(ws) + raise finally: self.__wsLock.release() if self.__get_channels and callable(self.__get_channels): @@ -476,6 +489,9 @@ def watchdog( except Exception as e: _log.warning("Websocket - Attempt failed: {0}".format(e)) finally: + # The attempt is over even if the websocket never opened. Wake + # its watchdog so it cannot fire later during reconnect backoff. + handshake_event.set() with self.__app_lock: if self.__app is app: self.__app = None @@ -867,13 +883,19 @@ def __init__(self, config: Config, on_trade: Callable[[OptionsTrade], None], on_ _log.setLevel(config.log_level) def __make_worker_threads(self) -> list[threading.Thread]: - return [threading.Thread( + return [self.__make_worker_thread(i, self.__data, self._stop_event) + for i in range(self.__num_threads)] + + def __make_worker_thread(self, index: int, data_queue: queue.Queue, + stop_event: threading.Event) -> threading.Thread: + return threading.Thread( group=None, target=_thread_fn, - args=(i, self.__data, self.__on_trade, self.__on_quote, self.__on_refresh, self.__on_unusual_activity, self._stop_event), + args=(index, data_queue, self.__on_trade, self.__on_quote, + self.__on_refresh, self.__on_unusual_activity, stop_event), kwargs={}, daemon=True - ) for i in range(self.__num_threads)] + ) def __all_ready(self) -> bool: self.__ws_lock.acquire() @@ -1038,7 +1060,9 @@ def __socket_start_fn(self, token: str, stop_event: threading.Event, self.__use_on_refresh, self.__use_on_unusual_activity, data_queue, - stop_event) + stop_event, + lambda index: self.__make_worker_thread( + index, data_queue, stop_event)) self.__webSocket = web_socket web_socket.start() @@ -1137,11 +1161,51 @@ def start(self): start_result["error"] = failure start_complete.set() + def __publish_stop_result(self, stop_complete: threading.Event, + stop_result: dict, surviving_threads: list, + failure: BaseException = None): + with self.__lifecycle_lock: + self.__is_started = bool(surviving_threads) + self.__is_stopping = False + self.__stopping_thread = None + stop_result["error"] = str(failure) if failure is not None else None + stop_complete.set() + if failure is not None: + _log.error("Stop incomplete: {0}".format(failure)) + else: + _log.info("Stopped") + + def __finish_stop_after_owned_thread_exits( + self, owned_thread: threading.Thread, + candidate_threads: list[threading.Thread], + stop_complete: threading.Event, stop_result: dict, + failure: BaseException = None): + try: + owned_thread.join() + except BaseException as error: + if failure is None: + failure = error + surviving_threads = [ + thread for thread in candidate_threads if thread.is_alive() + ] + if failure is None and surviving_threads: + failure = RuntimeError("Client stop incomplete: threads did not exit") + self.__publish_stop_result( + stop_complete, stop_result, surviving_threads, failure) + def stop(self): current_thread = threading.current_thread() with self.__lifecycle_lock: if self.__is_stopping: - if self.__stopping_thread is current_thread: + owned_threads = self.__worker_threads + ( + [self.__socket_thread] + if self.__socket_thread is not None else [] + ) + if (self.__stopping_thread is current_thread + or current_thread in owned_threads): + # An owned thread must return to its run loop before the + # active shutdown can join it. Waiting here would make the + # shutdown wait on its own completion event. return stop_complete = self.__stop_complete stop_result = self.__stop_result @@ -1168,15 +1232,13 @@ def stop(self): return candidate_threads = [ - worker for worker in worker_threads - if worker.ident is not None and worker is not current_thread + worker for worker in worker_threads if worker.ident is not None ] - if (socket_thread is not None - and socket_thread.ident is not None - and socket_thread is not current_thread): + if socket_thread is not None and socket_thread.ident is not None: candidate_threads.append(socket_thread) surviving_threads = [] failure = None + deferred_completion = False try: try: if ws is not None: @@ -1213,20 +1275,30 @@ def stop(self): # Recheck after all joins and cleanup so threads that exited while a # later thread was being joined are not retained as false survivors. surviving_threads = [thread for thread in candidate_threads if thread.is_alive()] - if failure is None and surviving_threads: + owned_current_thread = ( + current_thread if current_thread in surviving_threads else None + ) + other_survivors = [ + thread for thread in surviving_threads + if thread is not owned_current_thread + ] + if failure is None and other_survivors: failure = RuntimeError("Client stop incomplete: threads did not exit") - with self.__lifecycle_lock: - self.__is_started = bool(surviving_threads) - self.__is_stopping = False - self.__stopping_thread = None - stop_result["error"] = str(failure) if failure is not None else None - stop_complete.set() - if failure is not None: - _log.error("Stop incomplete: {0}".format(failure)) + if owned_current_thread is not None: + deferred_completion = True + threading.Thread( + target=self.__finish_stop_after_owned_thread_exits, + args=(owned_current_thread, candidate_threads, + stop_complete, stop_result, failure), + daemon=True, + ).start() else: - _log.info("Stopped") + self.__publish_stop_result( + stop_complete, stop_result, surviving_threads, failure) if failure is not None: raise failure + if deferred_completion: + _log.debug("Stop completion deferred until callback thread exits") def get_stats(self) -> tuple[int, int, int]: with self.__stats_lock: diff --git a/tests/test_equities_client.py b/tests/test_equities_client.py index 6b454ef..189315e 100644 --- a/tests/test_equities_client.py +++ b/tests/test_equities_client.py @@ -10,6 +10,7 @@ from intriniorealtime.equities_client import ( CONNECT_TIMEOUT_SECONDS, + STABLE_SESSION_SECONDS, EquitiesQuoteHandler, EquitiesQuoteReceiver, IntrinioRealtimeEquitiesClient, @@ -32,6 +33,9 @@ def setUp(self): def tearDown(self): logging.disable(logging.NOTSET) + def test_default_stable_session_requires_sustained_health(self): + self.assertGreaterEqual(STABLE_SESSION_SECONDS, 60) + def test_on_connect_rejoins_channels_after_reconnect(self): client = _make_client() client.channels = {"AAPL"} @@ -441,8 +445,7 @@ def join(self, timeout=None): if self.join_calls == 1: old_join_entered.set() release_disconnect.wait(timeout=2) - else: - self.alive = False + self.alive = False old_receiver = OldReceiver() new_receiver = Mock() @@ -714,24 +717,97 @@ def refresh(): client.connect() ctor.assert_not_called() - def test_disconnect_keeps_receiver_if_join_times_out(self): + def test_incomplete_disconnect_blocks_connect_until_receiver_exits(self): client = _make_client() - old = Mock() - old.is_alive.return_value = True + release_receiver = threading.Event() + old = threading.Thread(target=release_receiver.wait, daemon=True) old.enabled = True - old.join = Mock() + old.start() client.quote_receiver = old client.ws = Mock() ws = client.ws - with patch("intriniorealtime.equities_client.close_websocket_app") as abort: - client.disconnect() + with patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", -1), \ + patch("intriniorealtime.equities_client.close_websocket_app") as abort: + with self.assertRaisesRegex(RuntimeError, "disconnect incomplete"): + client.disconnect() self.assertFalse(old.enabled) - old.join.assert_called() self.assertIs(client.quote_receiver, old) self.assertIs(client.ws, ws) + self.assertTrue(client._disconnecting) self.assertGreaterEqual(abort.call_count, 2) + with self.assertRaisesRegex(RuntimeError, "disconnecting"): + client.connect() + + release_receiver.set() + old.join(timeout=1) + deadline = time.time() + 1 + while client._disconnecting and time.time() < deadline: + time.sleep(0.01) + + self.assertFalse(client._disconnecting) + self.assertIsNone(client.quote_receiver) + self.assertIsNone(client.ws) + + def test_deferred_disconnect_waits_for_calling_handler_too(self): + client = _make_client() + release_receiver = threading.Event() + release_handler = threading.Event() + disconnect_returned = threading.Event() + disconnect_errors = [] + + class Receiver(threading.Thread): + def __init__(self): + super().__init__(target=release_receiver.wait, daemon=True) + self.enabled = True + + class Handler(threading.Thread): + def __init__(self): + super().__init__(target=self.run_callback, daemon=True) + self.stopped = False + + def stop(self): + self.stopped = True + + def run_callback(self): + try: + client.disconnect() + except Exception as error: + disconnect_errors.append(error) + disconnect_returned.set() + release_handler.wait() + + receiver = Receiver() + handler = Handler() + receiver.start() + client.quote_receiver = receiver + client.quote_handler = handler + client.ws = Mock() + + with patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", -1): + handler.start() + self.assertTrue(disconnect_returned.wait(timeout=1)) + + self.assertEqual(len(disconnect_errors), 1) + self.assertRegex(str(disconnect_errors[0]), "disconnect incomplete") + self.assertTrue(client._disconnecting) + + release_receiver.set() + receiver.join(timeout=1) + time.sleep(0.05) + self.assertTrue(client._disconnecting) + self.assertTrue(handler.is_alive()) + + release_handler.set() + handler.join(timeout=1) + deadline = time.time() + 1 + while client._disconnecting and time.time() < deadline: + time.sleep(0.01) + + self.assertFalse(client._disconnecting) + self.assertIsNone(client.quote_receiver) + self.assertIsNone(client.quote_handler) def test_refresh_websocket_does_not_replace_receiver_still_alive_after_join(self): client = _make_client() @@ -805,6 +881,35 @@ def fake_run_forever(app, *args, **kwargs): self.assertEqual(client.refresh_token.call_count, 1) client.do_backoff.assert_not_called() + def test_completed_handshake_attempt_cancels_its_watchdog(self): + client = _make_client() + client.token = "tok" + receiver = EquitiesQuoteReceiver(client) + run_count = {"n": 0} + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + if run_count["n"] >= 2: + receiver.enabled = False + client._stop_event.set() + + def backoff(): + client._stop_event.wait(timeout=0.06) + + client.do_backoff = backoff + client.refresh_token = Mock() + with patch("intriniorealtime.equities_client.CONNECT_TIMEOUT_SECONDS", 0.02), \ + patch("intriniorealtime.equities_client.websocket.WebSocketApp.run_forever", fake_run_forever), \ + patch.object(client.logger, "warning") as warning: + receiver.run() + + timeout_warnings = [ + call for call in warning.call_args_list + if call.args and "connect timed out" in call.args[0] + ] + self.assertEqual(run_count["n"], 2) + self.assertEqual(timeout_warnings, []) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_options_client.py b/tests/test_options_client.py index d77fdd2..f422ca8 100644 --- a/tests/test_options_client.py +++ b/tests/test_options_client.py @@ -12,13 +12,15 @@ IntrinioRealtimeOptionsClient, Providers, _CONNECT_TIMEOUT_SECONDS, + _STABLE_SESSION_SECONDS, _WebSocket, _thread_fn, ) def _make_ws(stop_flag, get_token=None, get_url=None, channels=None, data_queue=None, - channels_lock=None, stats=None, stats_lock=None, sent_channels=None): + channels_lock=None, stats=None, stats_lock=None, sent_channels=None, + worker_threads=None, worker_factory=None): return _WebSocket( "ws://example/old", threading.Lock(), @@ -26,7 +28,7 @@ def _make_ws(stop_flag, get_token=None, get_url=None, channels=None, data_queue= sent_channels if sent_channels is not None else set(), stats if stats is not None else {"data": 0, "text": 0}, stats_lock or threading.Lock(), - [], + worker_threads if worker_threads is not None else [], lambda: channels if channels is not None else set(), get_token or (lambda: "tok"), get_url or (lambda token: "ws://example/" + token), @@ -36,6 +38,7 @@ def _make_ws(stop_flag, get_token=None, get_url=None, channels=None, data_queue= False, data_queue if data_queue is not None else queue.Queue(), stop_flag, + worker_factory, ) @@ -56,6 +59,9 @@ def setUp(self): def tearDown(self): logging.disable(logging.NOTSET) + def test_default_stable_session_requires_sustained_health(self): + self.assertGreaterEqual(_STABLE_SESSION_SECONDS, 60) + def test_config_rejects_non_positive_or_non_integer_worker_counts(self): for num_threads in (0, -1, True, 1.5, "2"): with self.subTest(num_threads=num_threads): @@ -332,6 +338,47 @@ def __exit__(self, exc_type, exc_value, traceback): self.assertFalse(opener.is_alive()) self.assertFalse(ws.isReady) + def test_open_replaces_a_terminated_worker_before_publishing_ready(self): + stop = threading.Event() + terminated = threading.Thread(target=lambda: None) + terminated.start() + terminated.join(timeout=1) + replacements = [] + + def make_worker(_index): + worker = threading.Thread(target=stop.wait) + replacements.append(worker) + return worker + + workers = [terminated] + ws = _make_ws( + stop, + worker_threads=workers, + worker_factory=make_worker, + ) + + ws.on_open(ws) + + self.assertTrue(ws.isReady) + self.assertIs(workers[0], replacements[0]) + self.assertTrue(replacements[0].is_alive()) + stop.set() + replacements[0].join(timeout=1) + + def test_open_rolls_back_ready_when_worker_start_fails(self): + stop = threading.Event() + terminated = threading.Thread(target=lambda: None) + terminated.start() + terminated.join(timeout=1) + ws = _make_ws(stop, worker_threads=[terminated]) + + with patch("intriniorealtime.options_client.close_websocket_app") as close: + with self.assertRaisesRegex(RuntimeError, "cannot be restarted"): + ws.on_open(ws) + + self.assertFalse(ws.isReady) + close.assert_called_once_with(ws) + def test_open_aborts_connection_when_channel_replay_fails(self): stop = threading.Event() sent_channels = set() @@ -516,17 +563,77 @@ def test_stop_joins_socket_thread_with_timeout(self): self.assertGreater(timeout, 0) self.assertLessEqual(timeout, _CONNECT_TIMEOUT_SECONDS + 1) - def test_stop_from_worker_callback_does_not_join_current_thread(self): + def test_stop_from_worker_callback_finishes_after_worker_exits(self): client = _make_client() - current = threading.current_thread() - client._IntrinioRealtimeOptionsClient__worker_threads = [current] + callback_entered = threading.Event() + stop_returned = threading.Event() + release_callback = threading.Event() + + def callback_worker(): + callback_entered.set() + client.stop() + stop_returned.set() + release_callback.wait(timeout=1) + + worker = threading.Thread(target=callback_worker) + client._IntrinioRealtimeOptionsClient__worker_threads = [worker] client._IntrinioRealtimeOptionsClient__socket_thread = None client._IntrinioRealtimeOptionsClient__webSocket = None + client._IntrinioRealtimeOptionsClient__is_started = True + worker.start() - with patch("intriniorealtime.options_client.time.sleep"): - client.stop() + self.assertTrue(callback_entered.wait(timeout=1)) + self.assertTrue(stop_returned.wait(timeout=1)) + self.assertTrue(client._IntrinioRealtimeOptionsClient__is_stopping) + + waiting_stop = threading.Thread(target=client.stop) + waiting_stop.start() + waiting_stop.join(timeout=0.05) + self.assertTrue(waiting_stop.is_alive()) + + with self.assertRaisesRegex(RuntimeError, "client is stopping"): + client.start() + + release_callback.set() + worker.join(timeout=1) + waiting_stop.join(timeout=1) self.assertTrue(client._stop_event.is_set()) + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_stopping) + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) + self.assertFalse(waiting_stop.is_alive()) + + def test_concurrent_worker_stops_do_not_wait_on_each_other(self): + client = _make_client() + barrier = threading.Barrier(3) + results = [] + + def callback_worker(): + barrier.wait() + try: + client.stop() + results.append(None) + except Exception as error: + results.append(error) + + workers = [threading.Thread(target=callback_worker) for _ in range(2)] + client._IntrinioRealtimeOptionsClient__worker_threads = workers + client._IntrinioRealtimeOptionsClient__socket_thread = None + client._IntrinioRealtimeOptionsClient__is_started = True + for worker in workers: + worker.start() + + barrier.wait() + for worker in workers: + worker.join(timeout=2) + + deadline = time.time() + 1 + while client._IntrinioRealtimeOptionsClient__is_stopping and time.time() < deadline: + time.sleep(0.01) + + self.assertEqual(results, [None, None]) + self.assertFalse(any(worker.is_alive() for worker in workers)) + self.assertFalse(client._IntrinioRealtimeOptionsClient__is_stopping) self.assertFalse(client._IntrinioRealtimeOptionsClient__is_started) def test_open_then_drop_twice_backs_off_on_second_reconnect(self): @@ -620,6 +727,29 @@ def fake_run_forever(app, *args, **kwargs): self.assertEqual(run_count["n"], 2) self.assertIn(0.01, waits) + def test_completed_handshake_attempt_cancels_its_watchdog(self): + stop = threading.Event() + run_count = {"n": 0} + + def fake_run_forever(app, *args, **kwargs): + run_count["n"] += 1 + if run_count["n"] >= 2: + stop.set() + + ws = _make_ws(stop) + with patch("intriniorealtime.options_client._CONNECT_TIMEOUT_SECONDS", 0.02), \ + patch("intriniorealtime.options_client._SELF_HEAL_BACKOFFS", [0.06]), \ + patch("intriniorealtime.options_client.websocket.WebSocketApp.run_forever", fake_run_forever), \ + patch("intriniorealtime.options_client._log.warning") as warning: + ws.start() + + timeout_warnings = [ + call for call in warning.call_args_list + if call.args and "Connect timed out" in call.args[0] + ] + self.assertEqual(run_count["n"], 2) + self.assertEqual(timeout_warnings, []) + def test_stop_sets_flag_before_socket_abort(self): client = _make_client() saw_stop = {} From 15381ffd818fc1b80157eda40f90fb502e7a203b Mon Sep 17 00:00:00 2001 From: Shawn Snyder Date: Mon, 28 Sep 2026 13:23:24 -0500 Subject: [PATCH 7/7] redact api key in certain error logging messages. --- intriniorealtime/_websocket.py | 5 +++++ intriniorealtime/equities_client.py | 6 +++--- intriniorealtime/options_client.py | 4 ++-- 3 files changed, 10 insertions(+), 5 deletions(-) diff --git a/intriniorealtime/_websocket.py b/intriniorealtime/_websocket.py index a5f71ab..d7a7d40 100644 --- a/intriniorealtime/_websocket.py +++ b/intriniorealtime/_websocket.py @@ -32,3 +32,8 @@ def should_abort_handshake( if remaining <= 0: return True handshake_event.wait(min(remaining, 0.1)) + + +def _redact_api_key(message, api_key): + text = str(message) + return text.replace(api_key, "") if api_key else text diff --git a/intriniorealtime/equities_client.py b/intriniorealtime/equities_client.py index 53a4d88..8049f7e 100644 --- a/intriniorealtime/equities_client.py +++ b/intriniorealtime/equities_client.py @@ -10,7 +10,7 @@ from enum import IntEnum, unique from typing import Optional, Dict, Any -from ._websocket import close_websocket_app, should_abort_handshake +from ._websocket import close_websocket_app, should_abort_handshake, _redact_api_key SELF_HEAL_BACKOFFS = [10, 30, 60, 300, 600] CONNECT_TIMEOUT_SECONDS = 30 @@ -272,7 +272,7 @@ def connect(self): self.refresh_websocket(requested_generation) return except Exception as e: - self.logger.error(f"Cannot connect: {repr(e)}") + self.logger.error("Cannot connect: %s", _redact_api_key(repr(e), self.api_key)) self.do_backoff() def _finish_incomplete_disconnect(self, receiver, handler, survivors): @@ -650,7 +650,7 @@ def watchdog( self.client.refresh_token() break except Exception as e: - self.client.logger.error(f"Cannot connect: {repr(e)}") + self.client.logger.error("Cannot connect: %s", _redact_api_key(repr(e), self.client.api_key)) self.client.logger.debug("QuoteReceiver exiting") diff --git a/intriniorealtime/options_client.py b/intriniorealtime/options_client.py index c451bf5..073ac6b 100644 --- a/intriniorealtime/options_client.py +++ b/intriniorealtime/options_client.py @@ -9,7 +9,7 @@ from collections.abc import Callable from enum import IntEnum, unique -from ._websocket import close_websocket_app, should_abort_handshake +from ._websocket import close_websocket_app, should_abort_handshake, _redact_api_key _SELF_HEAL_BACKOFFS = [10, 30, 60, 300, 600] _CONNECT_TIMEOUT_SECONDS = 30 @@ -947,7 +947,7 @@ def __try_set_token(self) -> bool: _log.error("Authorization Failure: The request timed out.") return False except requests.exceptions.ConnectionError as err: - _log.error("Authorization Failure: {0}".format(err)) + _log.error("Authorization Failure: %s", _redact_api_key(err, self.__apiKey)) return False def __get_token(self) -> str: