[SPT-1644] Log WebSocket server errors instead of dropping them

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()

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(

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."
+            ],
+        )