Wait for disconnect when we are out of connection ble slots in esphome (#79246)

homeassistant/components/esphome/bluetooth/client.py

@@ -8,6 +8,7 @@ from typing import Any, TypeVar, cast
 import uuid
 
 from aioesphomeapi.connection import APIConnectionError, TimeoutAPIError
+import async_timeout
 from bleak.backends.characteristic import BleakGATTCharacteristic
 from bleak.backends.client import BaseBleakClient, NotifyCallback
 from bleak.backends.device import BLEDevice
@@ -24,6 +25,10 @@ from .service import BleakGATTServiceESPHome
 
 DEFAULT_MTU = 23
 GATT_HEADER_SIZE = 3
+DISCONNECT_TIMEOUT = 5.0
+CONNECT_FREE_SLOT_TIMEOUT = 2.0
+GATT_READ_TIMEOUT = 30.0
+
 DEFAULT_MAX_WRITE_WITHOUT_RESPONSE = DEFAULT_MTU - GATT_HEADER_SIZE
 _LOGGER = logging.getLogger(__name__)
 
@@ -37,6 +42,19 @@ def mac_to_int(address: str) -> int:
     return int(address.replace(":", ""), 16)
 
 
+def verify_connected(func: _WrapFuncType) -> _WrapFuncType:
+    """Define a wrapper throw BleakError if not connected."""
+
+    async def _async_wrap_bluetooth_connected_operation(
+        self: "ESPHomeClient", *args: Any, **kwargs: Any
+    ) -> Any:
+        if not self._is_connected:  # pylint: disable=protected-access
+            raise BleakError("Not connected")
+        return await func(self, *args, **kwargs)
+
+    return cast(_WrapFuncType, _async_wrap_bluetooth_connected_operation)
+
+
 def api_error_as_bleak_error(func: _WrapFuncType) -> _WrapFuncType:
     """Define a wrapper throw esphome api errors as BleakErrors."""
 
@@ -128,6 +146,7 @@ class ESPHomeClient(BaseBleakClient):
         Returns:
             Boolean representing connection status.
         """
+        await self._wait_for_free_connection_slot(CONNECT_FREE_SLOT_TIMEOUT)
 
         connected_future: asyncio.Future[bool] = asyncio.Future()
 
@@ -179,8 +198,20 @@ class ESPHomeClient(BaseBleakClient):
         """Disconnect from the peripheral device."""
         self._unsubscribe_connection_state()
         await self._client.bluetooth_device_disconnect(self._address_as_int)
+        await self._wait_for_free_connection_slot(DISCONNECT_TIMEOUT)
         return True
 
+    async def _wait_for_free_connection_slot(self, timeout: float) -> None:
+        """Wait for a free connection slot."""
+        entry_data = self._async_get_entry_data()
+        if entry_data.ble_connections_free:
+            return
+        _LOGGER.debug(
+            "%s: Out of connection slots, waiting for a free one", self._source
+        )
+        async with async_timeout.timeout(timeout):
+            await entry_data.wait_for_ble_connections_free()
+
     @property
     def is_connected(self) -> bool:
         """Is Connected."""
@@ -191,11 +222,13 @@ class ESPHomeClient(BaseBleakClient):
         """Get ATT MTU size for active connection."""
         return self._mtu or DEFAULT_MTU
 
+    @verify_connected
     @api_error_as_bleak_error
     async def pair(self, *args: Any, **kwargs: Any) -> bool:
         """Attempt to pair."""
         raise NotImplementedError("Pairing is not available in ESPHome.")
 
+    @verify_connected
     @api_error_as_bleak_error
     async def unpair(self) -> bool:
         """Attempt to unpair."""
@@ -272,6 +305,7 @@ class ESPHomeClient(BaseBleakClient):
             raise BleakError(f"Characteristic {char_specifier} was not found!")
         return characteristic
 
+    @verify_connected
     @api_error_as_bleak_error
     async def read_gatt_char(
         self,
@@ -289,9 +323,10 @@ class ESPHomeClient(BaseBleakClient):
         """
         characteristic = self._resolve_characteristic(char_specifier)
         return await self._client.bluetooth_gatt_read(
-            self._address_as_int, characteristic.handle
+            self._address_as_int, characteristic.handle, GATT_READ_TIMEOUT
         )
 
+    @verify_connected
     @api_error_as_bleak_error
     async def read_gatt_descriptor(self, handle: int, **kwargs: Any) -> bytearray:
         """Perform read operation on the specified GATT descriptor.
@@ -302,9 +337,10 @@ class ESPHomeClient(BaseBleakClient):
             (bytearray) The read data.
         """
         return await self._client.bluetooth_gatt_read_descriptor(
-            self._address_as_int, handle
+            self._address_as_int, handle, GATT_READ_TIMEOUT
         )
 
+    @verify_connected
     @api_error_as_bleak_error
     async def write_gatt_char(
         self,
@@ -326,6 +362,7 @@ class ESPHomeClient(BaseBleakClient):
             self._address_as_int, characteristic.handle, bytes(data), response
         )
 
+    @verify_connected
     @api_error_as_bleak_error
     async def write_gatt_descriptor(
         self, handle: int, data: bytes | bytearray | memoryview
@@ -340,6 +377,7 @@ class ESPHomeClient(BaseBleakClient):
             self._address_as_int, handle, bytes(data)
         )
 
+    @verify_connected
     @api_error_as_bleak_error
     async def start_notify(
         self,

homeassistant/components/esphome/entry_data.py

@@ -89,6 +89,9 @@ class RuntimeEntryData:
     _storage_contents: dict[str, Any] | None = None
     ble_connections_free: int = 0
     ble_connections_limit: int = 0
+    _ble_connection_free_futures: list[asyncio.Future[int]] = field(
+        default_factory=list
+    )
 
     @callback
     def async_update_ble_connection_limits(self, free: int, limit: int) -> None:
@@ -97,6 +100,18 @@ class RuntimeEntryData:
         _LOGGER.debug("%s: BLE connection limits: %s/%s", name, free, limit)
         self.ble_connections_free = free
         self.ble_connections_limit = limit
+        if free:
+            for fut in self._ble_connection_free_futures:
+                fut.set_result(free)
+            self._ble_connection_free_futures.clear()
+
+    async def wait_for_ble_connections_free(self) -> int:
+        """Wait until there are free BLE connections."""
+        if self.ble_connections_free > 0:
+            return self.ble_connections_free
+        fut: asyncio.Future[int] = asyncio.Future()
+        self._ble_connection_free_futures.append(fut)
+        return await fut
 
     @callback
     def async_remove_entity(