Speed up ESPHome connection setup (#104304)

This commit is contained in:
J. Nick Koston 2023-11-22 23:27:17 +01:00 committed by GitHub
parent a3c0f36592
commit a59076d140
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 86 additions and 40 deletions

View file

@ -1,8 +1,11 @@
"""Bluetooth support for esphome."""
from __future__ import annotations
import asyncio
from collections.abc import Coroutine
from functools import partial
import logging
from typing import Any
from aioesphomeapi import APIClient, BluetoothProxyFeature
@ -43,6 +46,13 @@ def _async_can_connect(
return can_connect
@hass_callback
def _async_unload(unload_callbacks: list[CALLBACK_TYPE]) -> None:
"""Cancel all the callbacks on unload."""
for callback in unload_callbacks:
callback()
async def async_connect_scanner(
hass: HomeAssistant,
entry: ConfigEntry,
@ -92,27 +102,36 @@ async def async_connect_scanner(
hass, source, entry.title, new_info_callback, connector, connectable
)
client_data.scanner = scanner
coros: list[Coroutine[Any, Any, CALLBACK_TYPE]] = []
# These calls all return a callback that can be used to unsubscribe
# but we never unsubscribe so we don't care about the return value
if connectable:
# If its connectable be sure not to register the scanner
# until we know the connection is fully setup since otherwise
# there is a race condition where the connection can fail
await cli.subscribe_bluetooth_connections_free(
bluetooth_device.async_update_ble_connection_limits
coros.append(
cli.subscribe_bluetooth_connections_free(
bluetooth_device.async_update_ble_connection_limits
)
)
unload_callbacks = [
async_register_scanner(hass, scanner, connectable),
scanner.async_setup(),
]
if feature_flags & BluetoothProxyFeature.RAW_ADVERTISEMENTS:
await cli.subscribe_bluetooth_le_raw_advertisements(
scanner.async_on_raw_advertisements
coros.append(
cli.subscribe_bluetooth_le_raw_advertisements(
scanner.async_on_raw_advertisements
)
)
else:
await cli.subscribe_bluetooth_le_advertisements(scanner.async_on_advertisement)
coros.append(
cli.subscribe_bluetooth_le_advertisements(scanner.async_on_advertisement)
)
@hass_callback
def _async_unload() -> None:
for callback in unload_callbacks:
callback()
return _async_unload
await asyncio.gather(*coros)
return partial(
_async_unload,
[
async_register_scanner(hass, scanner, connectable),
scanner.async_setup(),
],
)

View file

@ -2,6 +2,7 @@
from __future__ import annotations
import asyncio
from collections.abc import Coroutine
import logging
from typing import TYPE_CHECKING, Any, NamedTuple
@ -10,6 +11,7 @@ from aioesphomeapi import (
APIConnectionError,
APIVersion,
DeviceInfo as EsphomeDeviceInfo,
EntityInfo,
HomeassistantServiceCall,
InvalidAuthAPIError,
InvalidEncryptionKeyAPIError,
@ -26,7 +28,14 @@ import voluptuous as vol
from homeassistant.components import tag, zeroconf
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ATTR_DEVICE_ID, CONF_MODE, EVENT_HOMEASSISTANT_STOP
from homeassistant.core import Event, HomeAssistant, ServiceCall, State, callback
from homeassistant.core import (
CALLBACK_TYPE,
Event,
HomeAssistant,
ServiceCall,
State,
callback,
)
from homeassistant.exceptions import TemplateError
from homeassistant.helpers import template
import homeassistant.helpers.config_validation as cv
@ -372,13 +381,20 @@ class ESPHomeManager:
stored_device_name = entry.data.get(CONF_DEVICE_NAME)
unique_id_is_mac_address = unique_id and ":" in unique_id
try:
device_info = await cli.device_info()
results = await asyncio.gather(
cli.device_info(),
cli.list_entities_services(),
)
except APIConnectionError as err:
_LOGGER.warning("Error getting device info for %s: %s", self.host, err)
# Re-connection logic will trigger after this
await cli.disconnect()
return
device_info: EsphomeDeviceInfo = results[0]
entity_infos_services: tuple[list[EntityInfo], list[UserService]] = results[1]
entity_infos, services = entity_infos_services
device_mac = format_mac(device_info.mac_address)
mac_address_matches = unique_id == device_mac
#
@ -439,44 +455,55 @@ class ESPHomeManager:
if device_info.name:
reconnect_logic.name = device_info.name
self.device_id = _async_setup_device_registry(hass, entry, entry_data)
entry_data.async_update_device_state(hass)
await asyncio.gather(
entry_data.async_update_static_infos(
hass, entry, entity_infos, device_info.mac_address
),
_setup_services(hass, entry_data, services),
)
setup_coros_with_disconnect_callbacks: list[
Coroutine[Any, Any, CALLBACK_TYPE]
] = []
if device_info.bluetooth_proxy_feature_flags_compat(cli.api_version):
entry_data.disconnect_callbacks.add(
await async_connect_scanner(
setup_coros_with_disconnect_callbacks.append(
async_connect_scanner(
hass, entry, cli, entry_data, self.domain_data.bluetooth_cache
)
)
self.device_id = _async_setup_device_registry(hass, entry, entry_data)
entry_data.async_update_device_state(hass)
if device_info.voice_assistant_version:
setup_coros_with_disconnect_callbacks.append(
cli.subscribe_voice_assistant(
self._handle_pipeline_start,
self._handle_pipeline_stop,
)
)
try:
entity_infos, services = await cli.list_entities_services()
await entry_data.async_update_static_infos(
hass, entry, entity_infos, device_info.mac_address
)
await _setup_services(hass, entry_data, services)
await asyncio.gather(
setup_results = await asyncio.gather(
*setup_coros_with_disconnect_callbacks,
cli.subscribe_states(entry_data.async_update_state),
cli.subscribe_service_calls(self.async_on_service_call),
cli.subscribe_home_assistant_states(self.async_on_state_subscription),
)
if device_info.voice_assistant_version:
entry_data.disconnect_callbacks.add(
await cli.subscribe_voice_assistant(
self._handle_pipeline_start,
self._handle_pipeline_stop,
)
)
hass.async_create_task(entry_data.async_save_to_store())
except APIConnectionError as err:
_LOGGER.warning("Error getting initial data for %s: %s", self.host, err)
# Re-connection logic will trigger after this
await cli.disconnect()
else:
_async_check_firmware_version(hass, device_info, entry_data.api_version)
_async_check_using_api_password(hass, device_info, bool(self.password))
return
for result_idx in range(len(setup_coros_with_disconnect_callbacks)):
cancel_callback = setup_results[result_idx]
if TYPE_CHECKING:
assert cancel_callback is not None
entry_data.disconnect_callbacks.add(cancel_callback)
hass.async_create_task(entry_data.async_save_to_store())
_async_check_firmware_version(hass, device_info, entry_data.api_version)
_async_check_using_api_password(hass, device_info, bool(self.password))
async def on_disconnect(self, expected_disconnect: bool) -> None:
"""Run disconnect callbacks on API disconnect."""