roborock_session.py

"""An MQTT session for sending and receiving messages.

See create_mqtt_session for a factory function to create an MQTT session.

This is a thin wrapper around the async MQTT client that handles dispatching messages
from a topic to a callback function, since the async MQTT client does not
support this out of the box. It also handles the authentication process and
receiving messages from the vacuum cleaner.
"""

import asyncio
import datetime
import logging
import ssl
from collections.abc import Callable
from contextlib import asynccontextmanager

import aiomqtt
from aiomqtt import MqttCodeError, MqttError, TLSParameters

from roborock.callbacks import CallbackMap
from roborock.diagnostics import Diagnostics, redact_topic_name

from .health_manager import HealthManager
from .session import MqttParams, MqttSession, MqttSessionException, MqttSessionUnauthorized

_LOGGER = logging.getLogger(__name__)
_MQTT_LOGGER = logging.getLogger(f"{__name__}.aiomqtt")

CLIENT_KEEPALIVE = datetime.timedelta(seconds=45)
TOPIC_KEEPALIVE = datetime.timedelta(seconds=60)

# Exponential backoff parameters
MIN_BACKOFF_INTERVAL = datetime.timedelta(seconds=10)
MAX_BACKOFF_INTERVAL = datetime.timedelta(hours=6)
BACKOFF_MULTIPLIER = 1.5


class MqttReasonCode:
    """MQTT Reason Codes used by Roborock devices.

    This is a subset of paho.mqtt.reasoncodes.ReasonCode where we would like
    different error handling behavior.
    """

    RC_ERROR_UNAUTHORIZED = 135


class RoborockMqttSession(MqttSession):
    """An MQTT session for sending and receiving messages.

    You can start a session invoking the start() method which will connect to
    the MQTT broker. A caller may subscribe to a topic, and the session keeps
    track of which callbacks to invoke for each topic.

    The client is run as a background task that will run until shutdown. Once
    connected, the client will wait for messages to be received in a loop. If
    the connection is lost, the client will be re-created and reconnected. There
    is backoff to avoid spamming the broker with connection attempts. The client
    will automatically re-establish any subscriptions when the connection is
    re-established.
    """

    def __init__(
        self,
        params: MqttParams,
        topic_idle_timeout: datetime.timedelta = TOPIC_KEEPALIVE,
    ):
        self._params = params
        self._reconnect_task: asyncio.Task[None] | None = None
        self._healthy = False
        self._stop = False
        self._backoff = MIN_BACKOFF_INTERVAL
        self._client: aiomqtt.Client | None = None
        self._client_subscribed_topics: set[str] = set()
        self._client_lock = asyncio.Lock()
        self._listeners: CallbackMap[str, bytes] = CallbackMap(_LOGGER)
        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)
        self._unauthorized_hook = params.unauthorized_hook

    @property
    def connected(self) -> bool:
        """True if the session is connected to the broker."""
        return self._healthy

    @property
    def health_manager(self) -> HealthManager:
        """Return the health manager for the session."""
        return self._health_manager

    async def start(self) -> None:
        """Start the MQTT session.

        This has special behavior for the first connection attempt where any
        failures are raised immediately. This is to allow the caller to
        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
        self._reconnect_task = None
        self._idle_timers.clear()

        for task in tasks:
            task.cancel()
        try:
            await asyncio.gather(*tasks, return_exceptions=True)
        except asyncio.CancelledError:
            pass

        self._healthy = False

    async def restart(self) -> None:
        """Force the session to disconnect and reconnect.

        The active connection task will be cancelled and restarted in the background, retried by
        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:
            _LOGGER.debug("No message loop task to cancel")

    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))
                await self._connection_task
            except asyncio.CancelledError:
                _LOGGER.debug("MQTT connection task cancelled")
            except Exception:
                # Exceptions are logged and handled in _run_connection.
                # There is a special case for exceptions on startup where we return
                # immediately. Otherwise, we let the reconnect loop retry with
                # backoff when the reconnect loop is active.
                if start_future and start_future.done() and start_future.exception():
                    return

            self._healthy = False
            start_future = None
            if self._stop:
                _LOGGER.debug("MQTT session closed, stopping retry loop")
                return
            if not self._client_subscribed_topics and not self._listeners.keys():
                _LOGGER.debug("MQTT session disconnected with no active subscriptions, deferring reconnect")
                self._diagnostics.increment("reconnect_deferred")
                while not self._stop and not self._client_subscribed_topics and not self._listeners.keys():
                    await asyncio.sleep(0.1)
                if self._stop:
                    _LOGGER.debug("MQTT session closed while waiting for active subscriptions")
                    return
                self._backoff = MIN_BACKOFF_INTERVAL
                continue
            _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)

    async def _run_connection(self, start_future: asyncio.Future[None] | None) -> None:
        """Connect to the MQTT broker and listen for messages.

        This is the primary connection loop for the MQTT session that is
        long running and processes incoming messages. If the connection
        is lost, this method will exit.
        """
        try:
            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 MqttCodeError as err:
            self._diagnostics.increment(f"connect_failure:{err.rc}")
            if start_future and not start_future.done():
                _LOGGER.debug("MQTT error starting session: %s", err)
                start_future.set_exception(err)
            else:
                _LOGGER.debug("MQTT error: %s", err)
            if err.rc == MqttReasonCode.RC_ERROR_UNAUTHORIZED and self._unauthorized_hook:
                _LOGGER.info("MQTT unauthorized/rate-limit error received, setting backoff to maximum")
                self._unauthorized_hook()
                self._backoff = MAX_BACKOFF_INTERVAL
            raise
        except MqttError as err:
            self._diagnostics.increment("connect_failure:unknown")
            if start_future and not start_future.done():
                _LOGGER.info("MQTT error starting session: %s", err)
                start_future.set_exception(err)
            else:
                _LOGGER.info("MQTT error: %s", err)
            raise
        except Exception as err:
            self._diagnostics.increment("connect_failure:uncaught")
            # This error is thrown when the MQTT loop is cancelled
            # and the generator is not stopped.
            if "generator didn't stop" in str(err) or "generator didn't yield" in str(err):
                _LOGGER.debug("MQTT loop was cancelled")
                return
            if start_future and not start_future.done():
                _LOGGER.error("Uncaught error starting MQTT session: %s", err)
                start_future.set_exception(err)
            else:
                _LOGGER.exception("Uncaught error during MQTT session: %s", err)
            raise

    @asynccontextmanager
    async def _mqtt_client(self, params: MqttParams) -> aiomqtt.Client:
        """Connect to the MQTT broker and listen for messages."""
        _LOGGER.debug("Connecting to %s:%s for %s", params.host, params.port, params.username)
        tls_params = None
        if params.tls:
            tls_params = TLSParameters(cert_reqs=ssl.CERT_REQUIRED if params.verify_tls else ssl.CERT_NONE)
        try:
            async with aiomqtt.Client(
                hostname=params.host,
                port=params.port,
                username=params.username,
                password=params.password,
                keepalive=int(CLIENT_KEEPALIVE.total_seconds()),
                protocol=aiomqtt.ProtocolVersion.V5,
                tls_params=tls_params,
                timeout=params.timeout,
                logger=_MQTT_LOGGER,
            ) as client:
                _LOGGER.debug("Connected to MQTT broker")
                # Re-establish any existing subscriptions
                async with self._client_lock:
                    self._client = client
                    for topic in self._client_subscribed_topics:
                        self._diagnostics.increment("resubscribe")
                        _LOGGER.debug("Re-establishing subscription to topic %s", redact_topic_name(topic))
                        # TODO: If this fails it will break the whole connection. Make
                        # this retry again in the background with backoff.
                        await client.subscribe(topic)

                yield client
        finally:
            async with self._client_lock:
                self._client = None

    async def subscribe(self, topic: str, callback: Callable[[bytes], None]) -> Callable[[], None]:
        """Subscribe to messages on the specified topic and invoke the callback for new messages.

        The callback will be called with the message payload as a bytes object. The callback
        should not block since it runs in the async loop. It should not raise any exceptions.

        The returned callable unsubscribes from the topic when called, but will delay actual
        unsubscription for the idle timeout period. If a new subscription comes in during the
        timeout, the timer is cancelled and the subscription is reused.
        """
        _LOGGER.debug("Subscribing to topic %s", redact_topic_name(topic))

        # 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)", redact_topic_name(topic))

        unsub = self._listeners.add_callback(topic, callback)

        async with self._client_lock:
            if topic not in self._client_subscribed_topics:
                self._client_subscribed_topics.add(topic)
                if self._client:
                    _LOGGER.debug("Establishing subscription to topic %s", topic)
                    try:
                        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:
            async def idle_unsubscribe():
                try:
                    await asyncio.sleep(self._topic_idle_timeout.total_seconds())
                    # Only unsubscribe if there are no callbacks left for this topic
                    if not self._listeners.get_callbacks(topic):
                        async with self._client_lock:
                            # Check again if we have listeners, in case a subscribe happened
                            # while we were waiting for the lock or after we popped the timer.
                            if self._listeners.get_callbacks(topic):
                                _LOGGER.debug("Skipping unsubscribe for %s, new listeners added", topic)
                                return

                            self._idle_timers.pop(topic, None)
                            self._client_subscribed_topics.discard(topic)

                            if self._client:
                                _LOGGER.debug("Idle timeout expired, unsubscribing from topic %s", topic)
                                try:
                                    await self._client.unsubscribe(topic)
                                except MqttError as err:
                                    _LOGGER.warning("Error unsubscribing from topic %s: %s", topic, err)
                except asyncio.CancelledError:
                    _LOGGER.debug("Idle unsubscribe for topic %s cancelled", topic)

            # Start the idle timer task
            task = asyncio.create_task(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()
            else:
                _LOGGER.debug("Unsubscribing topic %s, still have active callbacks", topic)

        return delayed_unsub

    async def publish(self, topic: str, message: bytes) -> None:
        """Publish a message on the topic."""
        _LOGGER.debug("Sending message to topic %s: %s", topic, message)
        client: aiomqtt.Client
        async with self._client_lock:
            if self._client is None:
                raise MqttSessionException("Could not publish message, MQTT client not connected")
            client = self._client
        try:
            with self._diagnostics.timer("publish"):
                await client.publish(topic, message)
        except MqttError as err:
            raise MqttSessionException(f"Error publishing message: {err}") from err


class LazyMqttSession(MqttSession):
    """An MQTT session that is started on first attempt to subscribe.

    This is a wrapper around an existing MqttSession that will only start
    the underlying session when the first attempt to subscribe or publish
    is made.
    """

    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:
        """True if the session is connected to the broker."""
        return self._session.connected

    @property
    def health_manager(self) -> HealthManager:
        """Return the health manager for the session."""
        return self._session.health_manager

    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

    async def subscribe(self, device_id: str, callback: Callable[[bytes], None]) -> Callable[[], None]:
        """Invoke the callback when messages are received on the topic.

        The returned callable unsubscribes from the topic when called.
        """
        await self._maybe_start()
        return await self._session.subscribe(device_id, callback)

    async def publish(self, topic: str, message: bytes) -> None:
        """Publish a message on the specified topic.

        This will raise an exception if the message could not be sent.
        """
        await self._maybe_start()
        return await self._session.publish(topic, message)

    async def close(self) -> None:
        """Cancels the mqtt loop.

        This will close the underlying session and will not allow it to be
        restarted again.
        """
        await self._session.close()

    async def restart(self) -> None:
        """Force the session to disconnect and reconnect."""
        await self._session.restart()


async def create_mqtt_session(params: MqttParams) -> MqttSession:
    """Create an MQTT session.

    This function is a factory for creating an MQTT session. This will
    raise an exception if initial attempt to connect fails. Once connected,
    the session will retry connecting on failure in the background.
    """
    session = RoborockMqttSession(params)
    await session.start()
    return session


async def create_lazy_mqtt_session(params: MqttParams) -> MqttSession:
    """Create a lazy MQTT session.

    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), diagnostics=params.diagnostics.subkey("lazy_mqtt"))