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
 
@@ -16,6 +17,7 @@
 )
 from roborock.devices.device import DeviceReadyCallback, RoborockDevice
 from roborock.exceptions import RoborockException
+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
@@ -58,6 +60,7 @@ def __init__(
         device_creator: DeviceCreator,
         mqtt_session: MqttSession,
         cache: Cache,
+        diagnostics: Diagnostics,
     ) -> None:
         """Initialize the DeviceManager with user data and optional cache storage.
 
@@ -68,12 +71,15 @@ 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, prefer_cache: bool = True) -> 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 or not prefer_cache:
             _LOGGER.debug("Fetching home data (prefer_cache=%s)", prefer_cache)
+            self._diagnostics.increment("fetch_home_data")
             try:
                 cache_data.home_data = await self._web_api.get_home_data()
             except RoborockException as ex:
@@ -116,6 +122,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:
@@ -182,7 +192,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:
@@ -226,6 +239,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 .health_manager import HealthManager
 from .session import MqttParams, MqttSession, MqttSessionException, MqttSessionUnauthorized
@@ -76,6 +77,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
         self._health_manager = HealthManager(self.restart)
 
     @property
@@ -96,24 +98,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
@@ -136,6 +144,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:
@@ -144,6 +153,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))
@@ -164,6 +174,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)
 
@@ -175,17 +186,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)
@@ -251,6 +264,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)
 
@@ -262,13 +276,15 @@ 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()
                         self._client_subscribed_topics.discard(topic)
                         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() -> None:
@@ -301,10 +317,10 @@ 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):
-                _LOGGER.debug("Unsubscribing topic %s, starting idle timer", topic)
                 schedule_unsubscribe()
             else:
                 _LOGGER.debug("Unsubscribing topic %s, still have active callbacks", topic)
@@ -320,7 +336,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
 
@@ -333,11 +350,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:
@@ -353,6 +371,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
 
@@ -403,4 +422,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
 from roborock.mqtt.health_manager import HealthManager
 
@@ -32,6 +33,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

@@ -348,3 +348,17 @@ async def test_start_connect_unexpected_error(home_data: HomeData, channel_failu
     """Test that some unexpected errors from start_connect are propagated."""
     with pytest.raises(Exception, match="Unexpected error"):
         await create_device_manager(USER_PARAMS)
+
+
+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()