From 26b702eb38880de4d88a4eaba875a83f198a1287 Mon Sep 17 00:00:00 2001 From: stan Date: Mon, 14 Apr 2025 03:04:34 +0800 Subject: [PATCH 1/9] Feat: replace native websocket module with picows --- polygon/websocket/__init__.py | 229 ++++++++++++++++++++++------------ 1 file changed, 150 insertions(+), 79 deletions(-) diff --git a/polygon/websocket/__init__.py b/polygon/websocket/__init__.py index 77865d3f..20f86724 100644 --- a/polygon/websocket/__init__.py +++ b/polygon/websocket/__init__.py @@ -7,8 +7,7 @@ import ssl import certifi from .models import * -from websockets.client import connect, WebSocketClientProtocol -from websockets.exceptions import ConnectionClosedOK, ConnectionClosedError +from picows import WSFrame, WSListener, WSMsgType, WSCloseCode, WSTransport, ws_connect from ..logging import get_logger import logging from ..exceptions import AuthError @@ -17,6 +16,78 @@ logger = get_logger("WebSocketClient") +class PolygonWSListener(WSListener): + def __init__(self, client): + self.client = client + self.transport = None + self.processor = None + self.reconnects = 0 + + def on_ws_connected(self, transport: WSTransport): + """Called when WebSocket connection is established""" + self.transport = transport + logger.debug("connected") + # Send auth message + transport.send(WSMsgType.TEXT, self.client.json.dumps({"action": "auth", "params": self.client.api_key})) + + def on_ws_frame(self, transport: WSTransport, frame: WSFrame): + """Called when a WebSocket frame is received""" + if frame.msg_type == WSMsgType.CLOSE: + close_code = frame.get_close_code() + close_reason = frame.get_close_reason() + logger.debug(f"connection closed: code={close_code}, reason={close_reason}") + transport.send_close(close_code) + return + + if frame.msg_type != WSMsgType.TEXT: + logger.debug(f"Received unexpected frame type: {frame.msg_type}") + return + + message = frame.get_payload_as_ascii_text() + + # Process the message + try: + msgJson = self.client.json.loads(message) + + # Handle auth response + if len(msgJson) > 0 and "status" in msgJson[0]: + if msgJson[0]["status"] == "auth_failed": + logger.error(f"Authentication failed: {msgJson[0]['message']}") + transport.send_close(WSCloseCode.PROTOCOL_ERROR) + return + elif msgJson[0]["status"] == "connected": + logger.debug(f"authed: {message}") + # Handle subscriptions after successful auth + if self.client.schedule_resub: + self.client._handle_subscriptions() + return + + # Handle regular messages + if not self.client.raw: + for m in msgJson: + if "ev" in m and m["ev"] == "status": + logger.debug(f"status: {m.get('message', '')}") + continue + + cmsg = parse(msgJson, logger) + else: + cmsg = message + + if len(cmsg) > 0 and self.client.processor: + asyncio.create_task(self.client.processor(cmsg)) + + except Exception as e: + logger.error(f"Error processing message: {e}") + + def on_ws_disconnected(self, transport): + """Called when WebSocket connection is closed""" + logger.debug("WebSocket connection closed") + self.reconnects += 1 + self.client.scheduled_subs = set(self.client.subs) + self.client.subs = set() + self.client.schedule_resub = True + + class WebSocketClient: def __init__( self, @@ -62,7 +133,9 @@ def __init__( self.subscribed = False self.subs: Set[str] = set() self.max_reconnects = max_reconnects - self.websocket: Optional[WebSocketClientProtocol] = None + self.transport = None + self.listener = PolygonWSListener(self) + self.processor = None if subscriptions is None: subscriptions = [] self.scheduled_subs: Set[str] = set(subscriptions) @@ -72,7 +145,6 @@ def __init__( else: self.json = json - # https://websockets.readthedocs.io/en/stable/reference/client.html#opening-a-connection async def connect( self, processor: Union[ @@ -89,73 +161,72 @@ async def connect( :param close_timeout: How long to wait for handshake when calling .close. :raises AuthError: If invalid API key is supplied. """ - reconnects = 0 - logger.debug("connect: %s", self.url) - # darwin needs some extra <3 - ssl_context = None - if self.url.startswith("wss://"): - ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) - ssl_context.load_verify_locations(certifi.where()) - - async for s in connect( - self.url, close_timeout=close_timeout, ssl=ssl_context, **kwargs - ): - self.websocket = s + logger.debug(f"connect: {self.url}") + self.processor = processor + + # For picows, we don't need to explicitly pass SSL context + # The library handles secure connections based on the URL scheme + + while True: try: - msg = await s.recv() - logger.debug("connected: %s", msg) - logger.debug("authing...") - await s.send( - self.json.dumps({"action": "auth", "params": self.api_key}) + # Connect using picows + _, client = await ws_connect( + lambda: self.listener, + self.url, + **kwargs ) - auth_msg = await s.recv() - auth_msg_parsed = self.json.loads(auth_msg) - logger.debug("authed: %s", auth_msg) - if auth_msg_parsed[0]["status"] == "auth_failed": - raise AuthError(auth_msg_parsed[0]["message"]) - while True: - if self.schedule_resub: - logger.debug( - "reconciling: %s %s", self.subs, self.scheduled_subs - ) - new_subs = self.scheduled_subs.difference(self.subs) - await self._subscribe(new_subs) - old_subs = self.subs.difference(self.scheduled_subs) - await self._unsubscribe(old_subs) - self.subs = self.scheduled_subs - self.subs = set(self.scheduled_subs) - self.schedule_resub = False - - try: - cmsg: Union[List[WebSocketMessage], Union[str, bytes]] = ( - await asyncio.wait_for(s.recv(), timeout=1) - ) - except asyncio.TimeoutError: - continue - - if not self.raw: - # we know cmsg is Data - msgJson = self.json.loads(cmsg) # type: ignore - for m in msgJson: - if m["ev"] == "status": - logger.debug("status: %s", m["message"]) - continue - cmsg = parse(msgJson, logger) - - if len(cmsg) > 0: - await processor(cmsg) # type: ignore - except ConnectionClosedOK as e: - logger.debug("connection closed (OK): %s", e) - return - except ConnectionClosedError as e: - logger.debug("connection closed (ERR): %s", e) - reconnects += 1 - self.scheduled_subs = set(self.subs) - self.subs = set() - self.schedule_resub = True - if self.max_reconnects is not None and reconnects > self.max_reconnects: - return - continue + self.transport = client.transport + + # Wait for disconnection + await self.transport.wait_disconnected() + + # Check if we should reconnect + if (self.max_reconnects is not None and + self.listener.reconnects > self.max_reconnects): + logger.debug(f"Max reconnects ({self.max_reconnects}) reached") + break + + # Wait before reconnecting + await asyncio.sleep(1) + + except Exception as e: + logger.error(f"Connection error: {e}") + await asyncio.sleep(1) + + # Check if we should reconnect + self.listener.reconnects += 1 + if (self.max_reconnects is not None and + self.listener.reconnects > self.max_reconnects): + logger.debug(f"Max reconnects ({self.max_reconnects}) reached") + break + + def _handle_subscriptions(self): + """Handle subscription reconciliation""" + if not self.transport: + return + + logger.debug(f"reconciling: {self.subs} {self.scheduled_subs}") + + # Handle new subscriptions + new_subs = self.scheduled_subs.difference(self.subs) + if new_subs: + subs = ",".join(new_subs) + logger.debug(f"subbing: {subs}") + self.transport.send(WSMsgType.TEXT, + self.json.dumps({"action": "subscribe", "params": subs}) + ) + + # Handle unsubscriptions + old_subs = self.subs.difference(self.scheduled_subs) + if old_subs: + subs = ",".join(old_subs) + logger.debug(f"unsubbing: {subs}") + self.transport.send(WSMsgType.TEXT, + self.json.dumps({"action": "unsubscribe", "params": subs}) + ) + + self.subs = set(self.scheduled_subs) + self.schedule_resub = False def run( self, @@ -180,20 +251,20 @@ async def handle_msg_wrapper(msgs): asyncio.run(self.connect(handle_msg_wrapper, close_timeout, **kwargs)) async def _subscribe(self, topics: Union[List[str], Set[str]]): - if self.websocket is None or len(topics) == 0: + if self.transport is None or len(topics) == 0: return subs = ",".join(topics) - logger.debug("subbing: %s", subs) - await self.websocket.send( + logger.debug(f"subbing: {subs}") + self.transport.send(WSMsgType.TEXT, self.json.dumps({"action": "subscribe", "params": subs}) ) async def _unsubscribe(self, topics: Union[List[str], Set[str]]): - if self.websocket is None or len(topics) == 0: + if self.transport is None or len(topics) == 0: return subs = ",".join(topics) - logger.debug("unsubbing: %s", subs) - await self.websocket.send( + logger.debug(f"unsubbing: {subs}") + self.transport.send(WSMsgType.TEXT, self.json.dumps({"action": "unsubscribe", "params": subs}) ) @@ -261,8 +332,8 @@ async def close(self): """ logger.debug("closing") - if self.websocket: - await self.websocket.close() - self.websocket = None + if self.transport: + self.transport.send_close(WSCloseCode.GOING_AWAY) + self.transport = None else: - logger.warning("no websocket open to close") + logger.warning("no websocket connection open to close") From e07f98a40c1978ac2716bbce4ae75bcc6415dfd8 Mon Sep 17 00:00:00 2001 From: stan Date: Mon, 14 Apr 2025 17:27:51 +0800 Subject: [PATCH 2/9] fix to use get_close_message() --- polygon/websocket/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/polygon/websocket/__init__.py b/polygon/websocket/__init__.py index 20f86724..bcdf0772 100644 --- a/polygon/websocket/__init__.py +++ b/polygon/websocket/__init__.py @@ -34,8 +34,8 @@ def on_ws_frame(self, transport: WSTransport, frame: WSFrame): """Called when a WebSocket frame is received""" if frame.msg_type == WSMsgType.CLOSE: close_code = frame.get_close_code() - close_reason = frame.get_close_reason() - logger.debug(f"connection closed: code={close_code}, reason={close_reason}") + close_message = frame.get_close_message() + logger.debug(f"connection closed: code={close_code}, reason={close_message}") transport.send_close(close_code) return From e5011c3dec234fa2f5adb244e522c27d71b12732 Mon Sep 17 00:00:00 2001 From: stan Date: Mon, 14 Apr 2025 21:04:19 +0800 Subject: [PATCH 3/9] handle json.loads error --- ...GU1NjQ1YzQyM2U3NzJhOSZzb3J0PXRpY2tlcg.json | 34 ------------------- 1 file changed, 34 deletions(-) delete mode 100644 test_rest/mocks/v3/reference/tickers&cursor=YWN0aXZlPXRydWUmZGF0ZT0yMDIyLTA0LTI3JmxpbWl0PTImb3JkZXI9YXNjJnBhZ2VfbWFya2VyPUFBJTdDZjEyMmJjYmY4YWQwNzRmZmJlMTZmNjkxOWQ0ZDc3NjZlMzA3MWNmNmU1Nzg3OGE0OGU1NjQ1YzQyM2U3NzJhOSZzb3J0PXRpY2tlcg.json diff --git a/test_rest/mocks/v3/reference/tickers&cursor=YWN0aXZlPXRydWUmZGF0ZT0yMDIyLTA0LTI3JmxpbWl0PTImb3JkZXI9YXNjJnBhZ2VfbWFya2VyPUFBJTdDZjEyMmJjYmY4YWQwNzRmZmJlMTZmNjkxOWQ0ZDc3NjZlMzA3MWNmNmU1Nzg3OGE0OGU1NjQ1YzQyM2U3NzJhOSZzb3J0PXRpY2tlcg.json b/test_rest/mocks/v3/reference/tickers&cursor=YWN0aXZlPXRydWUmZGF0ZT0yMDIyLTA0LTI3JmxpbWl0PTImb3JkZXI9YXNjJnBhZ2VfbWFya2VyPUFBJTdDZjEyMmJjYmY4YWQwNzRmZmJlMTZmNjkxOWQ0ZDc3NjZlMzA3MWNmNmU1Nzg3OGE0OGU1NjQ1YzQyM2U3NzJhOSZzb3J0PXRpY2tlcg.json deleted file mode 100644 index 0cad5983..00000000 --- a/test_rest/mocks/v3/reference/tickers&cursor=YWN0aXZlPXRydWUmZGF0ZT0yMDIyLTA0LTI3JmxpbWl0PTImb3JkZXI9YXNjJnBhZ2VfbWFya2VyPUFBJTdDZjEyMmJjYmY4YWQwNzRmZmJlMTZmNjkxOWQ0ZDc3NjZlMzA3MWNmNmU1Nzg3OGE0OGU1NjQ1YzQyM2U3NzJhOSZzb3J0PXRpY2tlcg.json +++ /dev/null @@ -1,34 +0,0 @@ -{ - "results": [ - { - "ticker": "AAA", - "name": "AAF First Priority CLO Bond ETF", - "market": "stocks", - "locale": "us", - "primary_exchange": "ARCX", - "type": "ETF", - "active": true, - "currency_name": "usd", - "composite_figi": "BBG00X5FSP48", - "share_class_figi": "BBG00X5FSPZ4", - "last_updated_utc": "2022-04-27T00:00:00Z" - }, - { - "ticker": "AAAU", - "name": "Goldman Sachs Physical Gold ETF Shares", - "market": "stocks", - "locale": "us", - "primary_exchange": "BATS", - "type": "ETF", - "active": true, - "currency_name": "usd", - "cik": "0001708646", - "composite_figi": "BBG00LPXX872", - "share_class_figi": "BBG00LPXX8Z1", - "last_updated_utc": "2022-04-27T00:00:00Z" - } - ], - "status": "OK", - "request_id": "40d60d83fa0628503b4d13387b7bde2a", - "count": 2 -} \ No newline at end of file From d96ca96f551f25de1be6cc0fddbb5fa697e45dd7 Mon Sep 17 00:00:00 2001 From: stan Date: Mon, 14 Apr 2025 21:32:37 +0800 Subject: [PATCH 4/9] log.debug -> log.info --- polygon/websocket/__init__.py | 60 ++++++++++++++++++++++++++--------- 1 file changed, 45 insertions(+), 15 deletions(-) diff --git a/polygon/websocket/__init__.py b/polygon/websocket/__init__.py index bcdf0772..126752fe 100644 --- a/polygon/websocket/__init__.py +++ b/polygon/websocket/__init__.py @@ -26,7 +26,7 @@ def __init__(self, client): def on_ws_connected(self, transport: WSTransport): """Called when WebSocket connection is established""" self.transport = transport - logger.debug("connected") + logger.info("WebSocket connected") # Send auth message transport.send(WSMsgType.TEXT, self.client.json.dumps({"action": "auth", "params": self.client.api_key})) @@ -35,19 +35,49 @@ def on_ws_frame(self, transport: WSTransport, frame: WSFrame): if frame.msg_type == WSMsgType.CLOSE: close_code = frame.get_close_code() close_message = frame.get_close_message() - logger.debug(f"connection closed: code={close_code}, reason={close_message}") + logger.info(f"WebSocket connection closed: code={close_code}, reason={close_message}") transport.send_close(close_code) return if frame.msg_type != WSMsgType.TEXT: - logger.debug(f"Received unexpected frame type: {frame.msg_type}") + logger.info(f"Received unexpected frame type: {frame.msg_type}") return message = frame.get_payload_as_ascii_text() # Process the message try: - msgJson = self.client.json.loads(message) + # Handle potential JSON parsing errors more gracefully + try: + msgJson = self.client.json.loads(message) + except json.JSONDecodeError as json_err: + # Log detailed information about the JSON parsing error + error_pos = json_err.pos + # Get a snippet of the message around the error position + start_pos = max(0, error_pos - 50) + end_pos = min(len(message), error_pos + 50) + context = message[start_pos:end_pos] + + logger.error(f"JSON decode error at position {error_pos}: {json_err}") + # Removed verbose debug logging of message context + + # Try to recover by trimming the message if it appears to be truncated + if "unexpected end of data" in str(json_err): + # Find the last complete JSON object by looking for the last '}]' sequence + last_complete = message.rfind('}]') + if last_complete > 0: + try: + # Try parsing up to the last complete object + fixed_msg = message[:last_complete+2] + msgJson = self.client.json.loads(fixed_msg) + logger.info(f"Recovered from truncated JSON by trimming to length {len(fixed_msg)}") + except json.JSONDecodeError: + # If recovery fails, re-raise the original error + raise json_err + else: + raise json_err + else: + raise json_err # Handle auth response if len(msgJson) > 0 and "status" in msgJson[0]: @@ -56,7 +86,7 @@ def on_ws_frame(self, transport: WSTransport, frame: WSFrame): transport.send_close(WSCloseCode.PROTOCOL_ERROR) return elif msgJson[0]["status"] == "connected": - logger.debug(f"authed: {message}") + logger.info("Authentication successful") # Handle subscriptions after successful auth if self.client.schedule_resub: self.client._handle_subscriptions() @@ -66,7 +96,7 @@ def on_ws_frame(self, transport: WSTransport, frame: WSFrame): if not self.client.raw: for m in msgJson: if "ev" in m and m["ev"] == "status": - logger.debug(f"status: {m.get('message', '')}") + logger.info(f"Status message: {m.get('message', '')}") continue cmsg = parse(msgJson, logger) @@ -81,7 +111,7 @@ def on_ws_frame(self, transport: WSTransport, frame: WSFrame): def on_ws_disconnected(self, transport): """Called when WebSocket connection is closed""" - logger.debug("WebSocket connection closed") + logger.info("WebSocket connection closed") self.reconnects += 1 self.client.scheduled_subs = set(self.client.subs) self.client.subs = set() @@ -161,7 +191,7 @@ async def connect( :param close_timeout: How long to wait for handshake when calling .close. :raises AuthError: If invalid API key is supplied. """ - logger.debug(f"connect: {self.url}") + logger.info(f"Connecting to: {self.url}") self.processor = processor # For picows, we don't need to explicitly pass SSL context @@ -183,7 +213,7 @@ async def connect( # Check if we should reconnect if (self.max_reconnects is not None and self.listener.reconnects > self.max_reconnects): - logger.debug(f"Max reconnects ({self.max_reconnects}) reached") + logger.info(f"Max reconnects ({self.max_reconnects}) reached") break # Wait before reconnecting @@ -197,7 +227,7 @@ async def connect( self.listener.reconnects += 1 if (self.max_reconnects is not None and self.listener.reconnects > self.max_reconnects): - logger.debug(f"Max reconnects ({self.max_reconnects}) reached") + logger.info(f"Max reconnects ({self.max_reconnects}) reached") break def _handle_subscriptions(self): @@ -205,13 +235,13 @@ def _handle_subscriptions(self): if not self.transport: return - logger.debug(f"reconciling: {self.subs} {self.scheduled_subs}") + logger.info("Reconciling subscriptions") # Handle new subscriptions new_subs = self.scheduled_subs.difference(self.subs) if new_subs: subs = ",".join(new_subs) - logger.debug(f"subbing: {subs}") + logger.info(f"Subscribing to: {subs}") self.transport.send(WSMsgType.TEXT, self.json.dumps({"action": "subscribe", "params": subs}) ) @@ -220,7 +250,7 @@ def _handle_subscriptions(self): old_subs = self.subs.difference(self.scheduled_subs) if old_subs: subs = ",".join(old_subs) - logger.debug(f"unsubbing: {subs}") + logger.info(f"Unsubscribing from: {subs}") self.transport.send(WSMsgType.TEXT, self.json.dumps({"action": "unsubscribe", "params": subs}) ) @@ -288,7 +318,7 @@ def subscribe(self, *subscriptions: str): topic, sym = self._parse_subscription(s) if topic == None: continue - logger.debug("sub desired: %s", s) + logger.info("Adding subscription: %s", s) self.scheduled_subs.add(s) # If user subs to X.*, remove other X.\w+ if sym == "*": @@ -308,7 +338,7 @@ def unsubscribe(self, *subscriptions: str): topic, sym = self._parse_subscription(s) if topic == None: continue - logger.debug("sub undesired: %s", s) + logger.info("Removing subscription: %s", s) self.scheduled_subs.discard(s) # If user unsubs to X.*, remove other X.\w+ From f0af7824d50cd04fee9dc4f78fe370ccfdfbae04 Mon Sep 17 00:00:00 2001 From: stan Date: Mon, 14 Apr 2025 21:59:10 +0800 Subject: [PATCH 5/9] print trackback --- polygon/websocket/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/polygon/websocket/__init__.py b/polygon/websocket/__init__.py index 126752fe..cb81fe37 100644 --- a/polygon/websocket/__init__.py +++ b/polygon/websocket/__init__.py @@ -11,6 +11,7 @@ from ..logging import get_logger import logging from ..exceptions import AuthError +import traceback env_key = "POLYGON_API_KEY" logger = get_logger("WebSocketClient") @@ -108,6 +109,7 @@ def on_ws_frame(self, transport: WSTransport, frame: WSFrame): except Exception as e: logger.error(f"Error processing message: {e}") + traceback.print_exc() def on_ws_disconnected(self, transport): """Called when WebSocket connection is closed""" From c08b647a9764b33be54320c673b9c6f480c601c6 Mon Sep 17 00:00:00 2001 From: stan Date: Mon, 14 Apr 2025 22:48:57 +0800 Subject: [PATCH 6/9] handle subscription --- polygon/websocket/__init__.py | 59 ++++++++++++++++++++++++++++++----- 1 file changed, 51 insertions(+), 8 deletions(-) diff --git a/polygon/websocket/__init__.py b/polygon/websocket/__init__.py index cb81fe37..0e6e710e 100644 --- a/polygon/websocket/__init__.py +++ b/polygon/websocket/__init__.py @@ -23,6 +23,8 @@ def __init__(self, client): self.transport = None self.processor = None self.reconnects = 0 + self.subscription_confirmed = False + self.last_subscription_time = None def on_ws_connected(self, transport: WSTransport): """Called when WebSocket connection is established""" @@ -96,8 +98,21 @@ def on_ws_frame(self, transport: WSTransport, frame: WSFrame): # Handle regular messages if not self.client.raw: for m in msgJson: + # Check for subscription status messages if "ev" in m and m["ev"] == "status": - logger.info(f"Status message: {m.get('message', '')}") + status_msg = m.get('message', '') + logger.info(f"Status message: {status_msg}") + + # Check for subscription confirmation + if 'successfully subscribed to' in status_msg.lower(): + logger.info("Subscription confirmed by Polygon") + self.subscription_confirmed = True + # Check for subscription failure + elif 'failed' in status_msg.lower() and 'subscribe' in status_msg.lower(): + logger.error(f"Subscription failed: {status_msg}") + # Trigger resubscription + self.subscription_confirmed = False + self.client.schedule_resub = True continue cmsg = parse(msgJson, logger) @@ -177,6 +192,17 @@ def __init__( else: self.json = json + async def _verify_subscription(self): + """Verify that subscriptions were successful and retry if needed""" + # Wait a reasonable time for subscription confirmation + await asyncio.sleep(5) + + # If we haven't received confirmation, try to resubscribe + if not self.listener.subscription_confirmed and self.scheduled_subs: + logger.warning("No subscription confirmation received, resubscribing...") + self.schedule_resub = True + self._handle_subscriptions() + async def connect( self, processor: Union[ @@ -198,6 +224,13 @@ async def connect( # For picows, we don't need to explicitly pass SSL context # The library handles secure connections based on the URL scheme + + # Reset subscription state on new connection + self.listener.subscription_confirmed = False + self.listener.last_subscription_time = None + + # Always resubscribe on reconnect + self.schedule_resub = True while True: try: @@ -223,6 +256,7 @@ async def connect( except Exception as e: logger.error(f"Connection error: {e}") + traceback.print_exc() await asyncio.sleep(1) # Check if we should reconnect @@ -232,6 +266,8 @@ async def connect( logger.info(f"Max reconnects ({self.max_reconnects}) reached") break + + def _handle_subscriptions(self): """Handle subscription reconciliation""" if not self.transport: @@ -239,14 +275,17 @@ def _handle_subscriptions(self): logger.info("Reconciling subscriptions") - # Handle new subscriptions - new_subs = self.scheduled_subs.difference(self.subs) - if new_subs: - subs = ",".join(new_subs) - logger.info(f"Subscribing to: {subs}") + # Always send all subscriptions to ensure they're properly registered + # This is more reliable than just sending the difference + if self.scheduled_subs: + subs = ",".join(self.scheduled_subs) + logger.info(f"Subscribing to all: {subs}") self.transport.send(WSMsgType.TEXT, self.json.dumps({"action": "subscribe", "params": subs}) ) + # Record subscription time for verification + self.listener.last_subscription_time = asyncio.get_event_loop().time() + self.listener.subscription_confirmed = False # Handle unsubscriptions old_subs = self.subs.difference(self.scheduled_subs) @@ -259,6 +298,10 @@ def _handle_subscriptions(self): self.subs = set(self.scheduled_subs) self.schedule_resub = False + + # Set up a task to verify subscription was successful + if self.scheduled_subs: + asyncio.create_task(self._verify_subscription()) def run( self, @@ -286,7 +329,7 @@ async def _subscribe(self, topics: Union[List[str], Set[str]]): if self.transport is None or len(topics) == 0: return subs = ",".join(topics) - logger.debug(f"subbing: {subs}") + logger.info(f"Subscribing to: {subs}") self.transport.send(WSMsgType.TEXT, self.json.dumps({"action": "subscribe", "params": subs}) ) @@ -295,7 +338,7 @@ async def _unsubscribe(self, topics: Union[List[str], Set[str]]): if self.transport is None or len(topics) == 0: return subs = ",".join(topics) - logger.debug(f"unsubbing: {subs}") + logger.info(f"Unsubscribing from: {subs}") self.transport.send(WSMsgType.TEXT, self.json.dumps({"action": "unsubscribe", "params": subs}) ) From 271ed6c050d1da481ac44574ad2d9f915968a189 Mon Sep 17 00:00:00 2001 From: stan Date: Tue, 15 Apr 2025 00:44:44 +0800 Subject: [PATCH 7/9] 1.delay 3sec then subscription 2.handle CONTINUATION frame --- polygon/websocket/__init__.py | 113 ++++++++++++++++++++++------------ 1 file changed, 72 insertions(+), 41 deletions(-) diff --git a/polygon/websocket/__init__.py b/polygon/websocket/__init__.py index 0e6e710e..e873a140 100644 --- a/polygon/websocket/__init__.py +++ b/polygon/websocket/__init__.py @@ -23,8 +23,14 @@ def __init__(self, client): self.transport = None self.processor = None self.reconnects = 0 - self.subscription_confirmed = False - self.last_subscription_time = None + self._full_msg = bytearray() + self._full_msg_type = None + + async def _delay_subscription(self): + # Wait for 3 seconds after authentication before subscribing + await asyncio.sleep(3) + if self.client.schedule_resub: + self.client._handle_subscriptions() def on_ws_connected(self, transport: WSTransport): """Called when WebSocket connection is established""" @@ -35,20 +41,53 @@ def on_ws_connected(self, transport: WSTransport): def on_ws_frame(self, transport: WSTransport, frame: WSFrame): """Called when a WebSocket frame is received""" + # Handle control frames if frame.msg_type == WSMsgType.CLOSE: close_code = frame.get_close_code() close_message = frame.get_close_message() logger.info(f"WebSocket connection closed: code={close_code}, reason={close_message}") transport.send_close(close_code) return - - if frame.msg_type != WSMsgType.TEXT: - logger.info(f"Received unexpected frame type: {frame.msg_type}") + elif frame.msg_type == WSMsgType.PING: + transport.send(WSMsgType.PONG, frame.get_payload_as_bytes()) + return + elif frame.msg_type == WSMsgType.PONG: + # Just ignore pong frames return - message = frame.get_payload_as_ascii_text() - - # Process the message + # Handle data frames (with fragmentation support) + if frame.fin: + if self._full_msg: + # Last fragment of a fragmented message + if frame.msg_type == 0: # Continuation frame + self._full_msg.extend(frame.get_payload_as_memoryview()) + message = self._full_msg + msg_type = self._full_msg_type + self._full_msg = bytearray() + self._full_msg_type = None + else: + # This shouldn't happen - fin=True but not a continuation frame + logger.warning(f"Unexpected frame type {frame.msg_type} with fin=True when processing fragmented message") + self._full_msg = bytearray() + self._full_msg_type = None + return + else: + # Single-frame message + if frame.msg_type == WSMsgType.TEXT: + message = frame.get_payload_as_memoryview() + msg_type = WSMsgType.TEXT + else: + logger.info(f"Ignoring unexpected frame type: {frame.msg_type}") + return + else: + # Fragment (not final) + if not self._full_msg: + # First fragment + self._full_msg_type = frame.msg_type + self._full_msg.extend(frame.get_payload_as_memoryview()) + return # Wait for more fragments + + # Process the complete message try: # Handle potential JSON parsing errors more gracefully try: @@ -90,9 +129,8 @@ def on_ws_frame(self, transport: WSTransport, frame: WSFrame): return elif msgJson[0]["status"] == "connected": logger.info("Authentication successful") - # Handle subscriptions after successful auth - if self.client.schedule_resub: - self.client._handle_subscriptions() + # Add a delay before subscribing to ensure auth is fully processed + asyncio.create_task(self._delay_subscription()) return # Handle regular messages @@ -103,16 +141,7 @@ def on_ws_frame(self, transport: WSTransport, frame: WSFrame): status_msg = m.get('message', '') logger.info(f"Status message: {status_msg}") - # Check for subscription confirmation - if 'successfully subscribed to' in status_msg.lower(): - logger.info("Subscription confirmed by Polygon") - self.subscription_confirmed = True - # Check for subscription failure - elif 'failed' in status_msg.lower() and 'subscribe' in status_msg.lower(): - logger.error(f"Subscription failed: {status_msg}") - # Trigger resubscription - self.subscription_confirmed = False - self.client.schedule_resub = True + # Just log status messages without special handling continue cmsg = parse(msgJson, logger) @@ -123,7 +152,23 @@ def on_ws_frame(self, transport: WSTransport, frame: WSFrame): asyncio.create_task(self.client.processor(cmsg)) except Exception as e: - logger.error(f"Error processing message: {e}") + # Log the error with the message + try: + if isinstance(message, bytearray): + message_str = message.decode('utf-8') + else: + message_str = frame.get_payload_as_ascii_text() + + # Show beginning and end of long messages + if len(message_str) > 200: + message_preview = f"{message_str[:100]} ... {message_str[-100:]}" + else: + message_preview = message_str + + logger.error(f"Error processing message: {e}\nMessage (len={len(message_str)}): {message_preview}") + except Exception as log_err: + logger.error(f"Error processing message: {e} (additionally, error logging message: {log_err})") + traceback.print_exc() def on_ws_disconnected(self, transport): @@ -192,16 +237,7 @@ def __init__( else: self.json = json - async def _verify_subscription(self): - """Verify that subscriptions were successful and retry if needed""" - # Wait a reasonable time for subscription confirmation - await asyncio.sleep(5) - - # If we haven't received confirmation, try to resubscribe - if not self.listener.subscription_confirmed and self.scheduled_subs: - logger.warning("No subscription confirmation received, resubscribing...") - self.schedule_resub = True - self._handle_subscriptions() + async def connect( self, @@ -225,9 +261,7 @@ async def connect( # For picows, we don't need to explicitly pass SSL context # The library handles secure connections based on the URL scheme - # Reset subscription state on new connection - self.listener.subscription_confirmed = False - self.listener.last_subscription_time = None + # Always resubscribe on reconnect self.schedule_resub = True @@ -238,6 +272,7 @@ async def connect( _, client = await ws_connect( lambda: self.listener, self.url, +# max_frame_size=1024*1024, # from 10K to 1MB **kwargs ) self.transport = client.transport @@ -283,9 +318,7 @@ def _handle_subscriptions(self): self.transport.send(WSMsgType.TEXT, self.json.dumps({"action": "subscribe", "params": subs}) ) - # Record subscription time for verification - self.listener.last_subscription_time = asyncio.get_event_loop().time() - self.listener.subscription_confirmed = False + # Handle unsubscriptions old_subs = self.subs.difference(self.scheduled_subs) @@ -299,9 +332,7 @@ def _handle_subscriptions(self): self.subs = set(self.scheduled_subs) self.schedule_resub = False - # Set up a task to verify subscription was successful - if self.scheduled_subs: - asyncio.create_task(self._verify_subscription()) + def run( self, From 9c8ff27488fa281e4fecfb45f3b3b49ae76bd1c3 Mon Sep 17 00:00:00 2001 From: stan Date: Tue, 15 Apr 2025 01:43:08 +0800 Subject: [PATCH 8/9] optimize: reuse bytearray rather than alloc new one --- polygon/websocket/__init__.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/polygon/websocket/__init__.py b/polygon/websocket/__init__.py index e873a140..152b7e9c 100644 --- a/polygon/websocket/__init__.py +++ b/polygon/websocket/__init__.py @@ -61,14 +61,15 @@ def on_ws_frame(self, transport: WSTransport, frame: WSFrame): # Last fragment of a fragmented message if frame.msg_type == 0: # Continuation frame self._full_msg.extend(frame.get_payload_as_memoryview()) - message = self._full_msg + # Create a copy of the message before clearing the buffer + message = bytearray(self._full_msg) msg_type = self._full_msg_type - self._full_msg = bytearray() + self._full_msg.clear() self._full_msg_type = None else: # This shouldn't happen - fin=True but not a continuation frame logger.warning(f"Unexpected frame type {frame.msg_type} with fin=True when processing fragmented message") - self._full_msg = bytearray() + self._full_msg.clear() self._full_msg_type = None return else: From 7c5c82eaa650940252ec7aa0b7f359fadbf56b8e Mon Sep 17 00:00:00 2001 From: stan Date: Wed, 16 Apr 2025 00:32:53 +0800 Subject: [PATCH 9/9] WebSocket Message Processing Improvements -Reduced Memory Allocations --- polygon/websocket/__init__.py | 59 +++++++++++++---------------------- 1 file changed, 22 insertions(+), 37 deletions(-) diff --git a/polygon/websocket/__init__.py b/polygon/websocket/__init__.py index 152b7e9c..dee944c3 100644 --- a/polygon/websocket/__init__.py +++ b/polygon/websocket/__init__.py @@ -61,11 +61,10 @@ def on_ws_frame(self, transport: WSTransport, frame: WSFrame): # Last fragment of a fragmented message if frame.msg_type == 0: # Continuation frame self._full_msg.extend(frame.get_payload_as_memoryview()) - # Create a copy of the message before clearing the buffer - message = bytearray(self._full_msg) + # Process the message directly without creating a copy + message = self._full_msg msg_type = self._full_msg_type - self._full_msg.clear() - self._full_msg_type = None + # We'll clear the buffer after processing else: # This shouldn't happen - fin=True but not a continuation frame logger.warning(f"Unexpected frame type {frame.msg_type} with fin=True when processing fragmented message") @@ -94,20 +93,13 @@ def on_ws_frame(self, transport: WSTransport, frame: WSFrame): try: msgJson = self.client.json.loads(message) except json.JSONDecodeError as json_err: - # Log detailed information about the JSON parsing error error_pos = json_err.pos - # Get a snippet of the message around the error position - start_pos = max(0, error_pos - 50) - end_pos = min(len(message), error_pos + 50) - context = message[start_pos:end_pos] - logger.error(f"JSON decode error at position {error_pos}: {json_err}") - # Removed verbose debug logging of message context # Try to recover by trimming the message if it appears to be truncated if "unexpected end of data" in str(json_err): # Find the last complete JSON object by looking for the last '}]' sequence - last_complete = message.rfind('}]') + last_complete = message.rfind(b'}]' if isinstance(message, (bytes, bytearray, memoryview)) else '}]') if last_complete > 0: try: # Try parsing up to the last complete object @@ -115,7 +107,6 @@ def on_ws_frame(self, transport: WSTransport, frame: WSFrame): msgJson = self.client.json.loads(fixed_msg) logger.info(f"Recovered from truncated JSON by trimming to length {len(fixed_msg)}") except json.JSONDecodeError: - # If recovery fails, re-raise the original error raise json_err else: raise json_err @@ -123,7 +114,7 @@ def on_ws_frame(self, transport: WSTransport, frame: WSFrame): raise json_err # Handle auth response - if len(msgJson) > 0 and "status" in msgJson[0]: + if msgJson and isinstance(msgJson, list) and len(msgJson) > 0 and "status" in msgJson[0]: if msgJson[0]["status"] == "auth_failed": logger.error(f"Authentication failed: {msgJson[0]['message']}") transport.send_close(WSCloseCode.PROTOCOL_ERROR) @@ -136,41 +127,35 @@ def on_ws_frame(self, transport: WSTransport, frame: WSFrame): # Handle regular messages if not self.client.raw: + has_status_messages = False # Initialize the variable for m in msgJson: # Check for subscription status messages if "ev" in m and m["ev"] == "status": + has_status_messages = True status_msg = m.get('message', '') logger.info(f"Status message: {status_msg}") - - # Just log status messages without special handling - continue - - cmsg = parse(msgJson, logger) + + # Only parse if we have messages to process + if msgJson and (not has_status_messages or len(msgJson) > 1): + cmsg = parse(msgJson, logger) + else: + cmsg = [] else: cmsg = message - if len(cmsg) > 0 and self.client.processor: + if cmsg and len(cmsg) > 0 and self.client.processor: + # Create a task for async processing asyncio.create_task(self.client.processor(cmsg)) except Exception as e: - # Log the error with the message - try: - if isinstance(message, bytearray): - message_str = message.decode('utf-8') - else: - message_str = frame.get_payload_as_ascii_text() - - # Show beginning and end of long messages - if len(message_str) > 200: - message_preview = f"{message_str[:100]} ... {message_str[-100:]}" - else: - message_preview = message_str - - logger.error(f"Error processing message: {e}\nMessage (len={len(message_str)}): {message_preview}") - except Exception as log_err: - logger.error(f"Error processing message: {e} (additionally, error logging message: {log_err})") - + # Log the error with minimal message details + logger.error(f"Error processing message: {e}") traceback.print_exc() + finally: + # Clear the buffer if we were using a fragmented message + if self._full_msg and msg_type == self._full_msg_type: + self._full_msg.clear() + self._full_msg_type = None def on_ws_disconnected(self, transport): """Called when WebSocket connection is closed"""