KeyError not raised in _on_subscribe callback

[kuriakosejohn]

In our code we are subscribing to topics and publishing concurrently, from two different asyncio tasks. In rare situations when the subscription response is delayed, and context switch happens to publishing task we are getting a KeyError exception. But in the _on_subcribe callback the exception is only logged not raised, causing the subscribe call to wait infinitely.

We are using version 2.0.1 of aiomqtt and I am not sure why the TimeoutError was not raised after the default 10sec timeout interval. According to MQTT Broker logs the connection got disconnected due to no keep alive message but aiomqtt client was stuck until we restarted the whole application.

def _on_subscribe(  # noqa: PLR0913
        self,
        client: mqtt.Client,
        userdata: Any,
        mid: int,
        reason_codes: list[ReasonCode],
        properties: Properties | None = None,
    ) -> None:
        """Called when we receive a SUBACK message from the broker."""
        try:
            fut = self._pending_subscribes.pop(mid)
            if not fut.done():
                fut.set_result(reason_codes)
        except KeyError:
            self._logger.exception(
                'Unexpected message ID "%d" in on_subscribe callback', mid
            )

Here is the message from our client log:

2025-03-06 05:38:24.418 - client - ERROR - Unexpected message ID "5" in on_subscribe callback
Traceback (most recent call last):
  File "/home/compulab/.venvs/lib/python3.11/site-packages/aiomqtt/client.py", line 576, in _on_subscribe
    fut = self._pending_subscribes.pop(mid)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
KeyError: 5

[PhiSchmidt]

I got the same or a similar problem with unsubscribing in 2.4.0
I'm running aiomqtt with fastapi and within a streaming endpoint a subscription is dynamically created while the client is listening for events and removed afterwards.

After calling client.unsubscribe() I get the KeyError in _on_unsubscribe. Further debugging revealed, that the finally block in _pending_call is removing the call from pending_dict right before on_unsubscribe is called.

@contextlib.contextmanager
    def _pending_call(
        self,
        mid: int,
        value: T,
        pending_dict: dict[int, T],
    ) -> Iterator[None]:
        if mid in self._pending_calls:
            msg = f'There already exists a pending call for message ID "{mid}"'
            raise RuntimeError(msg)
        pending_dict[mid] = value  # [1]
        try:
            # Log a warning if there is a concerning number of pending calls
            pending = len(list(self._pending_calls))
            if pending > self.pending_calls_threshold:
                self._logger.warning("There are %d pending publish calls.", pending)
            # Back to the caller (run whatever is inside the with statement)
            yield
        finally:
            # The normal procedure is:
            #  * We add the item at [1]
            #  * A callback will remove the item
            #
            # However, if the callback doesn't get called (e.g., due to a
            # network error) we still need to remove the item from the dict.
            with contextlib.suppress(KeyError):
                del pending_dict[mid]

It looks like the await

        with self._pending_call(mid, confirmation, self._pending_unsubscribes):
            # Wait for confirmation
            await self._wait_for(confirmation.wait(), timeout=timeout)

from ´unsubscribe()` is returning prematurely. But I can't find the reason for it... the timeout is not triggered.

[PhiSchmidt]

I was able to solve my problem. Maybe it helps the OP:
The problem was that fastapi executes requests in their own asyncio tasks, which are ended directly after the response is sent. This led to the behaviour mentioned above, where the await in the subscribe method finished before the subscription was actually removed. Decoupling the message generator and the subscribe / unsubscribe logic into an independet task did the trick.