Merge pull request #133 from opentensor/feat/thewhaleking/add-websocket-logging

Adds ability to log raw websockets for debugging.

Author: thewhaleking

async_substrate_interface/async_substrate.py

@@ -75,6 +75,7 @@
 ResultHandler = Callable[[dict, Any], Awaitable[tuple[dict, bool]]]
 
 logger = logging.getLogger("async_substrate_interface")
+raw_websocket_logger = logging.getLogger("raw_websocket")
 
 
 class AsyncExtrinsicReceipt:
@@ -505,6 +506,7 @@ def __init__(
         max_connections=100,
         shutdown_timer=5,
         options: Optional[dict] = None,
+        _log_raw_websockets: bool = False,
     ):
         """
         Websocket manager object. Allows for the use of a single websocket connection by multiple
@@ -532,6 +534,8 @@ def __init__(
         self._exit_task = None
         self._open_subscriptions = 0
         self._options = options if options else {}
+        self._log_raw_websockets = _log_raw_websockets
+
         try:
             now = asyncio.get_running_loop().time()
         except RuntimeError:
@@ -615,7 +619,10 @@ async def shutdown(self):
     async def _recv(self) -> None:
         try:
             # TODO consider wrapping this in asyncio.wait_for and use that for the timeout logic
-            response = json.loads(await self.ws.recv(decode=False))
+            recd = await self.ws.recv(decode=False)
+            if self._log_raw_websockets:
+                raw_websocket_logger.debug(f"WEBSOCKET_RECEIVE> {recd.decode()}")
+            response = json.loads(recd)
             self.last_received = await self.loop_time()
             async with self._lock:
                 # note that these 'subscriptions' are all waiting sent messages which have not received
@@ -660,7 +667,10 @@ async def send(self, payload: dict) -> int:
         # self._open_subscriptions += 1
         await self.max_subscriptions.acquire()
         try:
-            await self.ws.send(json.dumps({**payload, **{"id": original_id}}))
+            to_send = {**payload, **{"id": original_id}}
+            if self._log_raw_websockets:
+                raw_websocket_logger.debug(f"WEBSOCKET_SEND> {to_send}")
+            await self.ws.send(json.dumps(to_send))
             self.last_sent = await self.loop_time()
             return original_id
         except (ConnectionClosed, ssl.SSLError, EOFError):
@@ -699,6 +709,7 @@ def __init__(
         max_retries: int = 5,
         retry_timeout: float = 60.0,
         _mock: bool = False,
+        _log_raw_websockets: bool = False,
     ):
         """
         The asyncio-compatible version of the subtensor interface commands we use in bittensor. It is important to
@@ -716,16 +727,19 @@ def __init__(
             max_retries: number of times to retry RPC requests before giving up
             retry_timeout: how to long wait since the last ping to retry the RPC request
             _mock: whether to use mock version of the subtensor interface
+            _log_raw_websockets: whether to log raw websocket requests during RPC requests
 
         """
         self.max_retries = max_retries
         self.retry_timeout = retry_timeout
         self.chain_endpoint = url
         self.url = url
         self._chain = chain_name
+        self._log_raw_websockets = _log_raw_websockets
         if not _mock:
             self.ws = Websocket(
                 url,
+                _log_raw_websockets=_log_raw_websockets,
                 options={
                     "max_size": self.ws_max_size,
                     "write_limit": 2**16,

async_substrate_interface/substrate_addons.py

@@ -117,6 +117,7 @@ def __init__(
         max_retries: int = 5,
         retry_timeout: float = 60.0,
         _mock: bool = False,
+        _log_raw_websockets: bool = False,
         archive_nodes: Optional[list[str]] = None,
     ):
         fallback_chains = fallback_chains or []
@@ -151,6 +152,7 @@ def __init__(
                     _mock=_mock,
                     retry_timeout=retry_timeout,
                     max_retries=max_retries,
+                    _log_raw_websockets=_log_raw_websockets,
                 )
                 initialized = True
                 logger.info(f"Connected to {chain_url}")
@@ -260,6 +262,7 @@ def __init__(
         max_retries: int = 5,
         retry_timeout: float = 60.0,
         _mock: bool = False,
+        _log_raw_websockets: bool = False,
         archive_nodes: Optional[list[str]] = None,
     ):
         fallback_chains = fallback_chains or []
@@ -287,6 +290,7 @@ def __init__(
             _mock=_mock,
             retry_timeout=retry_timeout,
             max_retries=max_retries,
+            _log_raw_websockets=_log_raw_websockets,
         )
         self._original_methods = {
             method: getattr(self, method) for method in RETRY_METHODS

async_substrate_interface/sync_substrate.py

@@ -53,6 +53,7 @@
 ResultHandler = Callable[[dict, Any], tuple[dict, bool]]
 
 logger = logging.getLogger("async_substrate_interface")
+raw_websocket_logger = logging.getLogger("raw_websocket")
 
 
 class ExtrinsicReceipt:
@@ -485,6 +486,7 @@ def __init__(
         max_retries: int = 5,
         retry_timeout: float = 60.0,
         _mock: bool = False,
+        _log_raw_websockets: bool = False,
     ):
         """
         The sync compatible version of the subtensor interface commands we use in bittensor. Use this instance only
@@ -501,6 +503,7 @@ def __init__(
             max_retries: number of times to retry RPC requests before giving up
             retry_timeout: how to long wait since the last ping to retry the RPC request
             _mock: whether to use mock version of the subtensor interface
+            _log_raw_websockets: whether to log raw websocket requests during RPC requests
 
         """
         self.max_retries = max_retries
@@ -527,6 +530,7 @@ def __init__(
         self.registry_type_map = {}
         self.type_id_to_name = {}
         self._mock = _mock
+        self.log_raw_websockets = _log_raw_websockets
         if not _mock:
             self.ws = self.connect(init=True)
             self.initialize()
@@ -1831,12 +1835,18 @@ def _make_rpc_request(
         ws = self.connect(init=False if attempt == 1 else True)
         for payload in payloads:
             item_id = get_next_id()
-            ws.send(json.dumps({**payload["payload"], **{"id": item_id}}))
+            to_send = {**payload["payload"], **{"id": item_id}}
+            if self.log_raw_websockets:
+                raw_websocket_logger.debug(f"WEBSOCKET_SEND> {to_send}")
+            ws.send(json.dumps(to_send))
             request_manager.add_request(item_id, payload["id"])
 
         while True:
             try:
-                response = json.loads(ws.recv(timeout=self.retry_timeout, decode=False))
+                recd = ws.recv(timeout=self.retry_timeout, decode=False)
+                if self.log_raw_websockets:
+                    raw_websocket_logger.debug(f"WEBSOCKET_RECEIVE> {recd.decode()}")
+                response = json.loads(recd)
             except (TimeoutError, ConnectionClosed):
                 if attempt >= self.max_retries:
                     logger.warning(