diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py index 0d5409cd..9b94d01a 100644 --- a/massive/websocket/__init__.py +++ b/massive/websocket/__init__.py @@ -1,10 +1,12 @@ import os +from collections import Counter from enum import Enum from typing import Optional, Union, List, Set, Callable, Awaitable, Any import logging import json import asyncio import ssl +import time import certifi from .models import * from websockets.asyncio.client import connect, ClientConnection @@ -16,6 +18,10 @@ env_key = "MASSIVE_API_KEY" logger = get_logger("WebSocketClient") +# Status messages that mean things are working. Any other status (e.g. "error", +# "max_connections", "slow_consumer", "force_disconnect") is logged as a warning. +OK_STATUSES = {"connected", "auth_success", "success"} + class WebSocketClient: def __init__( @@ -70,6 +76,8 @@ def __init__( subscriptions = [] self.scheduled_subs: Set[str] = set(subscriptions) self.schedule_resub = True + self._server_errors: Counter = Counter() + self._server_errors_since = 0.0 if custom_json: self.json = custom_json else: @@ -112,6 +120,7 @@ async def connect( try: msg = await s.recv() logger.debug("connected: %s", msg) + self._log_server_statuses(self.json.loads(msg)) logger.debug("authing...") await s.send( self.json.dumps({"action": "auth", "params": self.api_key}) @@ -121,6 +130,9 @@ async def connect( logger.debug("authed: %s", auth_msg) if auth_msg_parsed[0]["status"] == "auth_failed": raise AuthError(auth_msg_parsed[0]["message"]) + # e.g. max_connections, which the server sends before closing + self._log_server_statuses(auth_msg_parsed) + self._flush_server_errors(force=True) while True: if self.schedule_resub: logger.debug( @@ -139,6 +151,7 @@ async def connect( await asyncio.wait_for(s.recv(), timeout=1) ) except asyncio.TimeoutError: + self._flush_server_errors() continue if not self.raw: @@ -147,17 +160,21 @@ async def connect( for m in msgJson: if m["ev"] == "status": logger.debug("status: %s", m["message"]) - continue + self._log_server_statuses(msgJson) cmsg = parse(msgJson, logger, self.market) + self._flush_server_errors() if len(cmsg) > 0: await processor(cmsg) # type: ignore except ConnectionClosedOK as e: logger.debug("connection closed (OK): %s", e) + # the server sends a status (e.g. slow_consumer) before closing + self._flush_server_errors(force=True) return except ConnectionClosedError as e: last_exc = e logger.debug("connection closed (ERR): %s", e) + self._flush_server_errors(force=True) reconnects += 1 self.scheduled_subs = set(self.subs) self.subs = set() @@ -195,6 +212,29 @@ async def handle_msg_wrapper(msgs): asyncio.run(self.connect(handle_msg_wrapper, close_timeout, **kwargs)) + def _log_server_statuses(self, msgs: List[Any]): + # status messages are dropped by parse(), so problems (e.g. subscriptions + # the server rejected, or why it disconnected) are logged here. The server + # can send one error per channel, so identical ones are counted and logged + # once a second. + for m in msgs: + if m.get("ev") != "status" or m.get("status") in OK_STATUSES: + continue + if not self._server_errors: + self._server_errors_since = time.monotonic() + self._server_errors[(m.get("status"), m.get("message"))] += 1 + + def _flush_server_errors(self, force: bool = False): + if not self._server_errors: + return + if not force and time.monotonic() - self._server_errors_since < 1: + return + for (status, message), count in self._server_errors.items(): + logger.warning( + "server %s: %s%s", status, message, f" (x{count})" if count > 1 else "" + ) + self._server_errors = Counter() + async def _subscribe(self, topics: Union[List[str], Set[str]]): if self.websocket is None or len(topics) == 0: return @@ -276,6 +316,7 @@ async def close(self): Close the websocket connection. """ logger.debug("closing") + self._flush_server_errors(force=True) if self.websocket: await self.websocket.close() diff --git a/test_websocket/mock_server.py b/test_websocket/mock_server.py index 4cf8ab4c..fc38748b 100644 --- a/test_websocket/mock_server.py +++ b/test_websocket/mock_server.py @@ -14,12 +14,41 @@ async def mock_server(websocket): message = json.loads(message) if "action" not in message: continue + if message["action"] == "auth" and message["params"] == "MAX-CONNECTIONS": + # emulate the server refusing a connection over the account's limit + await websocket.send( + '[{"ev":"status","status":"max_connections","message":"Maximum number of websocket connections exceeded."}]' + ) + await websocket.close(code=1008) + return if message["action"] == "auth": await websocket.send( '[{"ev":"status","status":"auth_success","message":"authenticated"}]' ) elif message["action"] == "subscribe": + if message["params"] == "SLOW.X": + # emulate the server disconnecting a slow consumer + await websocket.send( + '[{"ev":"status","status":"slow_consumer","message":"Disconnected slow consumer."}]' + ) + await websocket.close(code=1000) + return for p in message["params"].split(","): + if p.startswith("ERR."): + # emulate the server rejecting a subscription + await websocket.send( + json.dumps( + [ + { + "ev": "status", + "status": "error", + "message": "Subscription limit reached for feed.", + } + ] + * 2 + ) + ) + continue subs.add(p) await websocket.send( json.dumps( diff --git a/test_websocket/test_status.py b/test_websocket/test_status.py new file mode 100644 index 00000000..d333c9c8 --- /dev/null +++ b/test_websocket/test_status.py @@ -0,0 +1,97 @@ +from massive import WebSocketClient +from massive.websocket import logger +from base_ws import BaseTest +from mock_server import port +from websockets.asyncio.client import connect +import asyncio + + +class StatusTest(BaseTest): + async def asyncSetUp(self): + await super().asyncSetUp() + # wait until the mock server is accepting connections + for _ in range(100): + try: + async with connect(f"ws://localhost:{port}"): + return + except OSError: + await asyncio.sleep(0.05) + self.fail("mock server did not start") + + async def asyncTearDown(self): + # wait for the server to shut down so the next test can bind the port + self.task.cancel() + try: + await self.task + except asyncio.CancelledError: + pass + + async def run_client(self, *subscriptions): + c = WebSocketClient(api_key="", feed=f"localhost:{port}", secure=False) + c.subscribe(*subscriptions) + task = asyncio.create_task(c.connect(lambda msgs: asyncio.sleep(0))) + await asyncio.sleep(0.5) + await c.close() + await task + + async def test_server_errors_are_logged(self): + with self.assertLogs("WebSocketClient", level="WARNING") as logs: + await self.run_client("ERR.X") + self.assertEqual( + logs.output, + [ + "WARNING:WebSocketClient:server error: " + "Subscription limit reached for feed. (x2)" + ], + ) + + async def test_server_errors_are_grouped(self): + # one subscription with many rejected channels logs a single line + with self.assertLogs("WebSocketClient", level="WARNING") as logs: + await self.run_client(*[f"ERR.X{i}" for i in range(50)]) + self.assertEqual( + logs.output, + [ + "WARNING:WebSocketClient:server error: " + "Subscription limit reached for feed. (x100)" + ], + ) + + async def test_success_statuses_are_not_logged(self): + # assertNoLogs needs Python 3.10, so log a sentinel and check it's the only line + with self.assertLogs("WebSocketClient", level="WARNING") as logs: + await self.run_client("T.AAPL") + logger.warning("sentinel") + self.assertEqual(logs.output, ["WARNING:WebSocketClient:sentinel"]) + + async def test_disconnect_reason_is_logged(self): + # the server sends a status and then closes the connection normally + with self.assertLogs("WebSocketClient", level="WARNING") as logs: + c = WebSocketClient(api_key="", feed=f"localhost:{port}", secure=False) + c.subscribe("SLOW.X") + await asyncio.wait_for(c.connect(lambda msgs: asyncio.sleep(0)), 5) + self.assertEqual( + logs.output, + [ + "WARNING:WebSocketClient:server slow_consumer: Disconnected slow consumer." + ], + ) + + async def test_auth_reply_status_is_logged(self): + # e.g. max_connections arrives as the reply to the auth message + c = WebSocketClient( + api_key="MAX-CONNECTIONS", + feed=f"localhost:{port}", + secure=False, + max_reconnects=0, + ) + with self.assertLogs("WebSocketClient", level="WARNING") as logs: + with self.assertRaises(Exception): + await asyncio.wait_for(c.connect(lambda msgs: asyncio.sleep(0)), 5) + self.assertEqual( + logs.output, + [ + "WARNING:WebSocketClient:server max_connections: " + "Maximum number of websocket connections exceeded." + ], + )