Fix connection subscription leak on cancelation (#316)

aioesphomeapi/client.py

@@ -518,34 +518,46 @@ class APIClient:
                 self._connection.remove_message_callback(on_msg)
 
         try:
-            async with async_timeout.timeout(timeout):
-                await event.wait()
-        except asyncio.TimeoutError as err:
-            # Disconnect before raising the exception to ensure
-            # the slot is recovered before the timeout is raised
-            # to avoid race were we run out even though we have a slot.
-            await self.bluetooth_device_disconnect(address)
-            addr = to_human_readable_address(address)
-            _LOGGER.debug("%s: Connecting timed out, waiting for disconnect", addr)
             try:
-                async with async_timeout.timeout(disconnect_timeout):
+                async with async_timeout.timeout(timeout):
                     await event.wait()
-                    disconnect_timed_out = False
-            except asyncio.TimeoutError:
-                disconnect_timed_out = True
-            _LOGGER.debug("%s: Disconnect timed out: %s", addr, disconnect_timed_out)
+            except asyncio.TimeoutError as err:
+                # Disconnect before raising the exception to ensure
+                # the slot is recovered before the timeout is raised
+                # to avoid race were we run out even though we have a slot.
+                await self.bluetooth_device_disconnect(address)
+                addr = to_human_readable_address(address)
+                _LOGGER.debug("%s: Connecting timed out, waiting for disconnect", addr)
+                try:
+                    async with async_timeout.timeout(disconnect_timeout):
+                        await event.wait()
+                        disconnect_timed_out = False
+                except asyncio.TimeoutError:
+                    disconnect_timed_out = True
+                _LOGGER.debug(
+                    "%s: Disconnect timed out: %s", addr, disconnect_timed_out
+                )
+                try:
+                    unsub()
+                except ValueError:
+                    _LOGGER.warning(
+                        "%s: Bluetooth device connection timed out but already unsubscribed",
+                        addr,
+                    )
+                raise TimeoutAPIError(
+                    f"Timeout waiting for connect response while connecting to {addr} "
+                    f"after {timeout}s, disconnect timed out: {disconnect_timed_out}, "
+                    f" after {disconnect_timeout}s"
+                ) from err
+        except asyncio.CancelledError:
             try:
                 unsub()
             except ValueError:
                 _LOGGER.warning(
-                    "%s: Bluetooth device connection timed out but already unsubscribed",
+                    "%s: Bluetooth device connection canceled but already unsubscribed",
                     addr,
                 )
-            raise TimeoutAPIError(
-                f"Timeout waiting for connect response while connecting to {addr} "
-                f"after {timeout}s, disconnect timed out: {disconnect_timed_out}, "
-                f" after {disconnect_timeout}s"
-            ) from err
+            raise
 
         return unsub
 

aioesphomeapi/connection.py

@@ -423,7 +423,11 @@ class APIConnection:
     ) -> None:
         """Send a message to the remote and register the given message handler."""
         self._message_handlers.append(on_message)
-        await self.send_message(send_msg)
+        try:
+            await self.send_message(send_msg)
+        except asyncio.CancelledError:
+            self._message_handlers.remove(on_message)
+            raise
 
     async def send_message_await_response_complex(
         self,