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

This commit is contained in:
J. Nick Koston 2022-09-29 02:42:55 -10:00 committed by GitHub
parent 616b85df31
commit 0b5289f748
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 55 additions and 2 deletions

View File

@ -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,

View File

@ -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(