From a59076d140e316af686a9ed705281e0d6255ac9c Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Wed, 22 Nov 2023 23:27:17 +0100 Subject: [PATCH] Speed up ESPHome connection setup (#104304) --- .../components/esphome/bluetooth/__init__.py | 49 ++++++++---- homeassistant/components/esphome/manager.py | 77 +++++++++++++------ 2 files changed, 86 insertions(+), 40 deletions(-) diff --git a/homeassistant/components/esphome/bluetooth/__init__.py b/homeassistant/components/esphome/bluetooth/__init__.py index 9ef298145d3..6936afac714 100644 --- a/homeassistant/components/esphome/bluetooth/__init__.py +++ b/homeassistant/components/esphome/bluetooth/__init__.py @@ -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(), + ], + ) diff --git a/homeassistant/components/esphome/manager.py b/homeassistant/components/esphome/manager.py index 8282940a71d..85c311ecc81 100644 --- a/homeassistant/components/esphome/manager.py +++ b/homeassistant/components/esphome/manager.py @@ -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."""