Reconnection in fastapi

[tomliptrot]

Hi,

I'm using aiomqtt in a FastAPI application, and it has been working fine until got an unexpected disconnection. I'm not sure how to best deal with this. Part of the problem is that I find it hard to recreate a disconnection for testing purposes.

I was thinking of something like this:

import asyncio
import aiomqtt
import contextlib
import fastapi


async def listen(client):
    async with client.messages() as messages:
        await client.subscribe("humidity/#")
        async for message in messages:
            print(message.payload)


client = None


@contextlib.asynccontextmanager
async def lifespan(app):
    global client
    async with aiomqtt.Client("test.mosquitto.org") as c:
        # Make client globally available
        client = c
        # Listen for MQTT messages in (unawaited) asyncio task
        loop = asyncio.get_event_loop()
        task = loop.create_task(listen(client))
        yield
        # Cancel the task
        task.cancel()
        # Wait for the task to be cancelled
        try:
            await task
        except asyncio.CancelledError:
            pass


app = fastapi.FastAPI(lifespan=lifespan)


@router.get("/")
async def publish():
    interval = 5
    while True:
        try:
            await client.publish("humidity/outside", 0.38)
            break
        except aiomqtt.MqttCodeError:
            print(f"Connection lost; Reconnecting in {interval} seconds ...")
            await asyncio.sleep(interval)

My questions are:

1. what is the best way to test/simulate an MQTT client disconnecting
2. does the approach I put there look like it will work? Would I need to put the same logic around each client.publish() or is there a nicer way to do it once (I guess I could make a new publish_with_reconnect() function.)

[empicano]

Hi Tom,

This is a good discussion to have, aiomqtt needs improvement in this sense 🙂 To address your questions:

1. One option is to spin up a local MQTT broker, connect to it, and then shut the broker down. This is not very lightweight, but works. In the tests we use the internal _disconnected future to simulate a disconnection.
2. I don't expect it to (I didn't test it). In the example in the docs we re-enter the client context (async with client:) in addition to retrying the publish/subscribe call. Re-entering the client context rebuilds the connection.

The tricky thing is that with FastAPI we can't go back into the lifespan function (or want to for that matter; there's potentially other expensive stuff that's initialized there, like a database pool).

I think the best way would be to implement reconnection directly in aiomqtt such that it's automatic (e.g. like a database pool). For context, PR #216 made the context manager reusable and issue #6 has some background on reconnection.

@frederikaalund, what are your thoughts on this?

[empicano]

If you use __aenter__ and __aexit__ directly anyway (which I think you do), you should be able to call __aenter__ again to reconnect. Note that you will need to also redo your subscriptions (or use a persistent session). I'm not sure how calling __aenter__ from somewhere else affects the listener task that's waiting inside async for message in messages:, that could get a bit messy.

You could also try to use two clients, one to publish messages, and one for the listener that you initialize inside the listener task (not inside lifespan) like in the reconnection example.

[empicano]

@frederikaalund, do you have a rough idea of how we could implement this? 🙂 Like Tom, I'd be interested in implementing this, but I'm not quite sure how. I guess we could draw some inspiration from other MQTT clients or even database clients here.

[frederikaalund]

@frederikaalund, do you have a rough idea of how we could implement this? 🙂

You mean the reconnect part, or? :) My take is to split it into two layers: low-level and high-level. We already have the "low-level" client (that is just aiomqtt.Client. The high-level client is then something like:

from __future__ import annotations
import anyio

class HighLevelClient:
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        self._client_args = args
        self._client_kwargs = kwargs
        self._tg = anyio.create_task_group()
        self._send_stream, self._receive_stream = anyio.create_memory_object_stream()
        self._message_broadcast: BroadcastValue[Message] = BroadcastValue()  # See cyto.cytio

    def subscribe(self, topic) -> None:
        self._send_stream.send_nowait({"type": "subscribe", "topic": topic})

    def unsubscribe(self, topic) -> None:
        self._send_stream.send_nowait({"type": "unsubscribe", "topic": topic})

    def publish(self, *args: Any, **kwargs: Any) -> None:
        self._send_stream.send_nowait({"type": "publish", "args": args, "kwargs": kwargs})

    async def messages(self) -> Iterable[Message]:
        async with self._message_broadcast.subscribe(seed="nothing") as stream:
            async with stream as messages:
                async for message in messages:
                    yield message

    async def _maintain_connection(self) -> None:
        """Continously attempt to reconnect in the background."""
        client = aiomqtt.Client(*self._client_args, self._client_kwargs)
        while True:
            try:
                async with client, anyio.create_task_group() as tg:
                    # Listen for events
                    tg.start_soon(self._receive_events, client)
                    # Forward all received messages
                    tg.start_soon(self._broadcast_messages, client)
            except Exception:
                await anyio.sleep(...)

    async def _receive_events(self, client: aiomqtt.Client) -> None:
        async for event in self._receive_stream:
            match event:
                case {"type": "subscribe", "topic": topic}:
                    await client.subscribe(topic)
                case {"type": "unsubscribe", "topic": topic}:
                    await client.unsubscribe(topic)
                case {"type": "publish", "args": args, "kwargs": kwargs}:
                    await client.publish(*args, **kwargs)

    async def _broadcast_messages(self, client: aiomqtt.Client) -> None:
        async for message in client.messages():
            self._message_broadcast.publish(message)

    async def __aenter__(self) -> HighLevelClient:
        await self._tg.__aenter__()
        self._tg.start_soon(self._maintain_connection)
        return self

    async def __aexit__(self, ...) -> bool | None:
        return await self._tg.__aexit__(...)

Does it make sense? Basically, the high-level client queues up events (e.g., subscribe, unsubscribe, publish) from the user. Internally, the high-level client uses a background task to maintain I/O between the user and the low-level client. This way, the user does not have to bother with the connect/disconnect/reconnect at all. It happens transparently in the background.

On the flip side: The user does not know if/when their events actually go through. For that, they still have to use the low-level client.

Disclaimer: I have not run the code above. There are also some edge cases that we don't deal with yet. :)

[empicano]

Thanks for the great example, that makes things clear! 😊 It's very good to have an idea of how we could implement this in my head, even though I would like to do some of the interface changes we already discussed before starting with this.

---

@tomliptrot, in case you haven't seen it, I've added the ./script/develop script to run a local mosquitto broker via Docker that you can connect to via aiomqtt.Client("localhost", port=1883). That hopefully makes testing a disconnection more straightforward 😊

[tomliptrot]

Thanks both. I've been away, but when I get a chance i will try out your approach @frederikaalund