How to support reconnect and external shutdown?

[MRIIOT]

The intention of the below code is to support reconnection as well as external shutdown via asyncio.Event. The shutdown appears to work, but the reconnect does not. If I restart the broker, the code attempts to reconnect, but the message_task is always finished on reconnect. I'm not an expert in Python async, so any opinions you can share would be appreciated.

import os
import sys
import asyncio
from asyncio import Event
import aiomqtt
from aiomqtt import MqttError


BROKER_IP = "mosquitto"
BROKER_PORT = 1883
BROKER_USER = None
BROKER_PASS = None
BROKER_CID = "my-async-test"

shutdown_event = Event()


async def handle_broker_message(message):
    print(f"on message: {message.payload.decode("utf-8")}")


async def iterate_broker_messages_task(messages):
    async for message in messages:
        await asyncio.shield(handle_broker_message(message))


async def broker_task():
    client = aiomqtt.Client(
        hostname=BROKER_IP,
        port=BROKER_PORT,
        username=BROKER_USER,
        password=BROKER_PASS,
        identifier=BROKER_CID,
        clean_session=True,
        keepalive=60)

    while not shutdown_event.is_set():
        try:
            print("broker task: connecting")

            async with client:
                print("broker task: connected")

                await client.subscribe(topic="$SYS/broker/#", qos=0)

                message_task = asyncio.create_task(iterate_broker_messages_task(client.messages))
                close_task = asyncio.create_task(shutdown_event.wait())
                await asyncio.wait([message_task, close_task], return_when=asyncio.FIRST_COMPLETED)

                if message_task.done():
                    print("broker task: message task done")
                    message_task.result()
                    close_task.cancel()
                else:
                    print("broker task: close task done")
                    message_task.cancel()

        except MqttError as e:
            print("broker task: connection failed, will reconnect")
            message_task.cancel()
            close_task.cancel()
            await asyncio.sleep(5)

    print("broker task: done")


async def main():
    print("main: starting broker task")
    bt = asyncio.create_task(broker_task())
    print("main: sleeping")
    await asyncio.sleep(30)
    print("main: shutting down")
    shutdown_event.set()
    await asyncio.sleep(5)
    print("main: shut down")


if sys.platform.lower() == "win32" or os.name.lower() == "nt":
    from asyncio import set_event_loop_policy, WindowsSelectorEventLoopPolicy
    set_event_loop_policy(WindowsSelectorEventLoopPolicy())

asyncio.run(main())

[empicano]

Hi there, sorry for the late reply! You might have stumbled on the same issue as #268. The underlying problem is that when an exception is raised inside a generator, the generator will raise the StopIteration exception on each subsequent call (telling the program "I'm already finished!").

For an immediate workaround, you could try to use client._messages() instead of client.messages.

Let me know if that works 🙂

[MRIIOT]

Thanks, will give it a try.