feat: Add diagnostics library for tracking stats/counters

Add a generic library that can be used to track counters, or elapsed time (e.g. how long it takes on average to connect to mqtt).  This is a copy of https://github.com/allenporter/python-google-nest-sdm/blob/main/google_nest_sdm/diagnostics.py

This is an initial pass to add a few initial example metrics for MQTT, but we can add more as we need fine grained details in diagnostics.

Author: allenporter

roborock/devices/device_manager.py

@@ -3,8 +3,9 @@
 import asyncio
 import enum
 import logging
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from dataclasses import dataclass
+from typing import Any
 
 import aiohttp
 
@@ -15,6 +16,7 @@
     UserData,
 )
 from roborock.devices.device import DeviceReadyCallback, RoborockDevice
+from roborock.diagnostics import Diagnostics
 from roborock.map.map_parser import MapParserConfig
 from roborock.mqtt.roborock_session import create_lazy_mqtt_session
 from roborock.mqtt.session import MqttSession
@@ -57,6 +59,7 @@ def __init__(
         device_creator: DeviceCreator,
         mqtt_session: MqttSession,
         cache: Cache,
+        diagnostics: Diagnostics,
     ) -> None:
         """Initialize the DeviceManager with user data and optional cache storage.
 
@@ -67,11 +70,14 @@ def __init__(
         self._device_creator = device_creator
         self._devices: dict[str, RoborockDevice] = {}
         self._mqtt_session = mqtt_session
+        self._diagnostics = diagnostics
 
     async def discover_devices(self) -> list[RoborockDevice]:
         """Discover all devices for the logged-in user."""
+        self._diagnostics.increment("discover_devices")
         cache_data = await self._cache.get()
         if not cache_data.home_data:
+            self._diagnostics.increment("fetch_home_data")
             _LOGGER.debug("No cached home data found, fetching from API")
             cache_data.home_data = await self._web_api.get_home_data()
             await self._cache.set(cache_data)
@@ -109,6 +115,10 @@ async def close(self) -> None:
         tasks.append(self._mqtt_session.close())
         await asyncio.gather(*tasks)
 
+    def diagnostic_data(self) -> Mapping[str, Any]:
+        """Return diagnostics information about the device manager."""
+        return self._diagnostics.as_dict()
+
 
 @dataclass
 class UserParams:
@@ -175,7 +185,10 @@ async def create_device_manager(
     web_api = create_web_api_wrapper(user_params, session=session, cache=cache)
     user_data = user_params.user_data
 
+    diagnostics = Diagnostics()
+
     mqtt_params = create_mqtt_params(user_data.rriot)
+    mqtt_params.diagnostics = diagnostics.subkey("mqtt_session")
     mqtt_session = await create_lazy_mqtt_session(mqtt_params)
 
     def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDataProduct) -> RoborockDevice:
@@ -219,6 +232,6 @@ def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDat
             dev.add_ready_callback(ready_callback)
         return dev
 
-    manager = DeviceManager(web_api, device_creator, mqtt_session=mqtt_session, cache=cache)
+    manager = DeviceManager(web_api, device_creator, mqtt_session=mqtt_session, cache=cache, diagnostics=diagnostics)
     await manager.discover_devices()
     return manager

roborock/diagnostics.py

@@ -0,0 +1,84 @@
+"""Diagnostics for debugging.
+
+A Diagnostics object can be used to track counts and latencies of various
+operations within a module. This can be useful for debugging performance issues
+or understanding usage patterns.
+
+This is an internal facing module and is not intended for public use. Diagnostics
+data is collected and exposed to clients via higher level APIs like the
+DeviceManager.
+"""
+
+from __future__ import annotations
+
+import time
+from collections import Counter
+from collections.abc import Generator, Mapping
+from contextlib import contextmanager
+from typing import Any
+
+
+class Diagnostics:
+    """A class that holdes diagnostics information for a module.
+
+    You can use this class to hold counter or for recording timing information
+    that can be exported as a dictionary for debugging purposes.
+    """
+
+    def __init__(self) -> None:
+        """Initialize Diagnostics."""
+        self._counter: Counter = Counter()
+        self._subkeys: dict[str, Diagnostics] = {}
+
+    def increment(self, key: str, count: int = 1) -> None:
+        """Increment a counter for the specified key/event."""
+        self._counter.update(Counter({key: count}))
+
+    def elapsed(self, key_prefix: str, elapsed_ms: int = 1) -> None:
+        """Track a latency event for the specified key/event prefix."""
+        self.increment(f"{key_prefix}_count", 1)
+        self.increment(f"{key_prefix}_sum", elapsed_ms)
+
+    def as_dict(self) -> Mapping[str, Any]:
+        """Return diagnostics as a debug dictionary."""
+        data: dict[str, Any] = {k: self._counter[k] for k in self._counter}
+        for k, d in self._subkeys.items():
+            v = d.as_dict()
+            if not v:
+                continue
+            data[k] = v
+        return data
+
+    def subkey(self, key: str) -> Diagnostics:
+        """Return sub-Diagnositics object with the specified subkey.
+
+        This will create a new Diagnostics object if one does not already exist
+        for the specified subkey. Stats from the sub-Diagnostics will be included
+        in the parent Diagnostics when exported as a dictionary.
+
+        Args:
+            key: The subkey for the diagnostics.
+
+        Returns:
+            The Diagnostics object for the specified subkey.
+        """
+        if key not in self._subkeys:
+            self._subkeys[key] = Diagnostics()
+        return self._subkeys[key]
+
+    @contextmanager
+    def timer(self, key_prefix: str) -> Generator[None, None, None]:
+        """A context manager that records the timing of operations as a diagnostic."""
+        start = time.perf_counter()
+        try:
+            yield
+        finally:
+            end = time.perf_counter()
+            ms = int((end - start) * 1000)
+            self.elapsed(key_prefix, ms)
+
+    def reset(self) -> None:
+        """Clear all diagnostics, for testing."""
+        self._counter = Counter()
+        for d in self._subkeys.values():
+            d.reset()

roborock/mqtt/roborock_session.py

@@ -18,6 +18,7 @@
 from aiomqtt import MqttCodeError, MqttError, TLSParameters
 
 from roborock.callbacks import CallbackMap
+from roborock.diagnostics import Diagnostics
 
 from .session import MqttParams, MqttSession, MqttSessionException, MqttSessionUnauthorized
 
@@ -74,6 +75,7 @@ def __init__(
         self._connection_task: asyncio.Task[None] | None = None
         self._topic_idle_timeout = topic_idle_timeout
         self._idle_timers: dict[str, asyncio.Task[None]] = {}
+        self._diagnostics = params.diagnostics
 
     @property
     def connected(self) -> bool:
@@ -88,24 +90,30 @@ async def start(self) -> None:
         handle the failure and retry if desired itself. Once connected,
         the session will retry connecting in the background.
         """
+        self._diagnostics.increment("start_attempt")
         start_future: asyncio.Future[None] = asyncio.Future()
         loop = asyncio.get_event_loop()
         self._reconnect_task = loop.create_task(self._run_reconnect_loop(start_future))
         try:
             await start_future
         except MqttCodeError as err:
+            self._diagnostics.increment(f"start_failure:{err.rc}")
             if err.rc == MqttReasonCode.RC_ERROR_UNAUTHORIZED:
                 raise MqttSessionUnauthorized(f"Authorization error starting MQTT session: {err}") from err
             raise MqttSessionException(f"Error starting MQTT session: {err}") from err
         except MqttError as err:
+            self._diagnostics.increment("start_failure:unknown")
             raise MqttSessionException(f"Error starting MQTT session: {err}") from err
         except Exception as err:
+            self._diagnostics.increment("start_failure:uncaught")
             raise MqttSessionException(f"Unexpected error starting session: {err}") from err
         else:
+            self._diagnostics.increment("start_success")
             _LOGGER.debug("MQTT session started successfully")
 
     async def close(self) -> None:
         """Cancels the MQTT loop and shutdown the client library."""
+        self._diagnostics.increment("close")
         self._stop = True
         tasks = [task for task in [self._connection_task, self._reconnect_task, *self._idle_timers.values()] if task]
         self._connection_task = None
@@ -128,6 +136,7 @@ async def restart(self) -> None:
         the reconnect loop. This is a no-op if there is no active connection.
         """
         _LOGGER.info("Forcing MQTT session restart")
+        self._diagnostics.increment("restart")
         if self._connection_task:
             self._connection_task.cancel()
         else:
@@ -136,6 +145,7 @@ async def restart(self) -> None:
     async def _run_reconnect_loop(self, start_future: asyncio.Future[None] | None) -> None:
         """Run the MQTT loop."""
         _LOGGER.info("Starting MQTT session")
+        self._diagnostics.increment("start_loop")
         while True:
             try:
                 self._connection_task = asyncio.create_task(self._run_connection(start_future))
@@ -156,6 +166,7 @@ async def _run_reconnect_loop(self, start_future: asyncio.Future[None] | None) -
                 _LOGGER.debug("MQTT session closed, stopping retry loop")
                 return
             _LOGGER.info("MQTT session disconnected, retrying in %s seconds", self._backoff.total_seconds())
+            self._diagnostics.increment("reconnect_wait")
             await asyncio.sleep(self._backoff.total_seconds())
             self._backoff = min(self._backoff * BACKOFF_MULTIPLIER, MAX_BACKOFF_INTERVAL)
 
@@ -167,17 +178,19 @@ async def _run_connection(self, start_future: asyncio.Future[None] | None) -> No
         is lost, this method will exit.
         """
         try:
-            async with self._mqtt_client(self._params) as client:
-                self._backoff = MIN_BACKOFF_INTERVAL
-                self._healthy = True
-                _LOGGER.info("MQTT Session connected.")
-                if start_future and not start_future.done():
-                    start_future.set_result(None)
-
-                _LOGGER.debug("Processing MQTT messages")
-                async for message in client.messages:
-                    _LOGGER.debug("Received message: %s", message)
-                    self._listeners(message.topic.value, message.payload)
+            with self._diagnostics.timer("connection"):
+                async with self._mqtt_client(self._params) as client:
+                    self._backoff = MIN_BACKOFF_INTERVAL
+                    self._healthy = True
+                    _LOGGER.info("MQTT Session connected.")
+                    if start_future and not start_future.done():
+                        start_future.set_result(None)
+
+                    _LOGGER.debug("Processing MQTT messages")
+                    async for message in client.messages:
+                        _LOGGER.debug("Received message: %s", message)
+                        with self._diagnostics.timer("dispatch_message"):
+                            self._listeners(message.topic.value, message.payload)
         except MqttError as err:
             if start_future and not start_future.done():
                 _LOGGER.info("MQTT error starting session: %s", err)
@@ -219,6 +232,7 @@ async def _mqtt_client(self, params: MqttParams) -> aiomqtt.Client:
                 async with self._client_lock:
                     self._client = client
                     for topic in self._listeners.keys():
+                        self._diagnostics.increment("resubscribe")
                         _LOGGER.debug("Re-establishing subscription to topic %s", topic)
                         # TODO: If this fails it will break the whole connection. Make
                         # this retry again in the background with backoff.
@@ -243,6 +257,7 @@ async def subscribe(self, topic: str, callback: Callable[[bytes], None]) -> Call
 
         # If there is an idle timer for this topic, cancel it (reuse subscription)
         if idle_timer := self._idle_timers.pop(topic, None):
+            self._diagnostics.increment("unsubscribe_idle_cancel")
             idle_timer.cancel()
             _LOGGER.debug("Cancelled idle timer for topic %s (reused subscription)", topic)
 
@@ -252,12 +267,14 @@ async def subscribe(self, topic: str, callback: Callable[[bytes], None]) -> Call
             if self._client:
                 _LOGGER.debug("Establishing subscription to topic %s", topic)
                 try:
-                    await self._client.subscribe(topic)
+                    with self._diagnostics.timer("subscribe"):
+                        await self._client.subscribe(topic)
                 except MqttError as err:
                     # Clean up the callback if subscription fails
                     unsub()
                     raise MqttSessionException(f"Error subscribing to topic: {err}") from err
             else:
+                self._diagnostics.increment("subscribe_pending")
                 _LOGGER.debug("Client not connected, will establish subscription later")
 
         def schedule_unsubscribe():
@@ -283,9 +300,11 @@ async def idle_unsubscribe():
             self._idle_timers[topic] = task
 
         def delayed_unsub():
+            self._diagnostics.increment("unsubscribe")
             unsub()  # Remove the callback from CallbackMap
             # If no more callbacks for this topic, start idle timer
             if not self._listeners.get_callbacks(topic):
+                self._diagnostics.increment("unsubscribe_idle_start")
                 schedule_unsubscribe()
 
         return delayed_unsub
@@ -299,7 +318,8 @@ async def publish(self, topic: str, message: bytes) -> None:
                 raise MqttSessionException("Could not publish message, MQTT client not connected")
             client = self._client
         try:
-            await client.publish(topic, message)
+            with self._diagnostics.timer("publish"):
+                await client.publish(topic, message)
         except MqttError as err:
             raise MqttSessionException(f"Error publishing message: {err}") from err
 
@@ -312,11 +332,12 @@ class LazyMqttSession(MqttSession):
     is made.
     """
 
-    def __init__(self, session: RoborockMqttSession) -> None:
+    def __init__(self, session: RoborockMqttSession, diagnostics: Diagnostics) -> None:
         """Initialize the lazy session with an existing session."""
         self._lock = asyncio.Lock()
         self._started = False
         self._session = session
+        self._diagnostics = diagnostics
 
     @property
     def connected(self) -> bool:
@@ -327,6 +348,7 @@ async def _maybe_start(self) -> None:
         """Start the MQTT session if not already started."""
         async with self._lock:
             if not self._started:
+                self._diagnostics.increment("start")
                 await self._session.start()
                 self._started = True
 
@@ -377,4 +399,4 @@ async def create_lazy_mqtt_session(params: MqttParams) -> MqttSession:
     This function is a factory for creating an MQTT session that will
     only connect when the first attempt to subscribe or publish is made.
     """
-    return LazyMqttSession(RoborockMqttSession(params))
+    return LazyMqttSession(RoborockMqttSession(params), diagnostics=params.diagnostics.subkey("lazy_mqtt"))

roborock/mqtt/session.py

@@ -4,6 +4,7 @@
 from collections.abc import Callable
 from dataclasses import dataclass
 
+from roborock.diagnostics import Diagnostics
 from roborock.exceptions import RoborockException
 
 DEFAULT_TIMEOUT = 30.0
@@ -31,6 +32,14 @@ class MqttParams:
     timeout: float = DEFAULT_TIMEOUT
     """Timeout for communications with the broker in seconds."""
 
+    diagnostics: Diagnostics = Diagnostics()
+    """Diagnostics object for tracking MQTT session stats.
+
+    This defaults to a new Diagnostics object, but the common case is the
+    caller will provide their own (e.g., from a DeviceManager) so that the
+    shared MQTT session diagnostics are included in the overall diagnostics.
+    """
+
 
 class MqttSession(ABC):
     """An MQTT session for sending and receiving messages."""

tests/devices/test_device_manager.py

@@ -231,3 +231,17 @@ async def test_start_connect_failure(home_data: HomeData, channel_failure: Mock,
 
     await device_manager.close()
     assert mock_unsub.call_count == 1
+
+
+async def test_diagnostics_collection(home_data: HomeData) -> None:
+    """Test that diagnostics are collected correctly in the DeviceManager."""
+    device_manager = await create_device_manager(USER_PARAMS)
+    devices = await device_manager.get_devices()
+    assert len(devices) == 1
+
+    diagnostics = device_manager.diagnostic_data()
+    assert diagnostics is not None
+    assert diagnostics.get("discover_devices") == 1
+    assert diagnostics.get("fetch_home_data") == 1
+
+    await device_manager.close()