Speed up ESPHome connection setup (#104304)

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(
+        coros.append(
-            bluetooth_device.async_update_ble_connection_limits
+            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(
+        coros.append(
-            scanner.async_on_raw_advertisements
+            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
+    await asyncio.gather(*coros)
-    def _async_unload() -> None:
+    return partial(
-        for callback in unload_callbacks:
+        _async_unload,
-            callback()
+        [
-
+            async_register_scanner(hass, scanner, connectable),
-    return _async_unload
+            scanner.async_setup(),
+        ],
+    )

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(
+            setup_coros_with_disconnect_callbacks.append(
-                await async_connect_scanner(
+                async_connect_scanner(
                     hass, entry, cli, entry_data, self.domain_data.bluetooth_cache
                 )
             )
 
-        self.device_id = _async_setup_device_registry(hass, entry, entry_data)
+        if device_info.voice_assistant_version:
-        entry_data.async_update_device_state(hass)
+            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()
+            setup_results = await asyncio.gather(
-            await entry_data.async_update_static_infos(
+                *setup_coros_with_disconnect_callbacks,
-                hass, entry, entity_infos, device_info.mac_address
-            )
-            await _setup_services(hass, entry_data, services)
-            await asyncio.gather(
                 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:
+            return
-            _async_check_firmware_version(hass, device_info, entry_data.api_version)
+
-            _async_check_using_api_password(hass, device_info, bool(self.password))
+        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."""