Simplify stage 1 in bootstrap (#137668)

* Simplify stage 1 in bootstrap

* Add timeouts to STAGE 0

* Fix test

* Clarify pre import language

* Remove timeout for frontend and recorder

* Address review

---------

Co-authored-by: J. Nick Koston <nick@koston.org>

homeassistant/bootstrap.py

@@ -134,14 +134,12 @@ DATA_REGISTRIES_LOADED: HassKey[None] = HassKey("bootstrap_registries_loaded")
 LOG_SLOW_STARTUP_INTERVAL = 60
 SLOW_STARTUP_CHECK_INTERVAL = 1
 
+STAGE_0_SUBSTAGE_TIMEOUT = 60
 STAGE_1_TIMEOUT = 120
 STAGE_2_TIMEOUT = 300
 WRAP_UP_TIMEOUT = 300
 COOLDOWN_TIME = 60
 
-
-DEBUGGER_INTEGRATIONS = {"debugpy"}
-
 # Core integrations are unconditionally loaded
 CORE_INTEGRATIONS = {"homeassistant", "persistent_notification"}
 
@@ -172,12 +170,27 @@ FRONTEND_INTEGRATIONS = {
     # add it here.
     "backup",
 }
-RECORDER_INTEGRATIONS = {
+# Stage 0 is divided into substages. Each substage has a name, a set of integrations and a timeout.
-    # Setup after frontend
+# The substage containing recorder should have no timeout, as it could cancel a database migration.
-    # To record data
+# Recorder freezes "recorder" timeout during a migration, but it does not freeze other timeouts.
-    "recorder",
+# The substages preceding it should also have no timeout, until we ensure that the recorder
-}
+# is not accidentally promoted as a dependency of any of the integrations in them.
-DISCOVERY_INTEGRATIONS = ("bluetooth", "dhcp", "ssdp", "usb", "zeroconf")
+# If we add timeouts to the frontend substages, we should make sure they don't apply in recovery mode.
+STAGE_0_INTEGRATIONS = (
+    # Load logging and http deps as soon as possible
+    ("logging, http deps", LOGGING_AND_HTTP_DEPS_INTEGRATIONS, None),
+    # Setup frontend
+    ("frontend", FRONTEND_INTEGRATIONS, None),
+    # Setup recorder
+    ("recorder", {"recorder"}, None),
+    # Start up debuggers. Start these first in case they want to wait.
+    ("debugger", {"debugpy"}, STAGE_0_SUBSTAGE_TIMEOUT),
+    # Zeroconf is used for mdns resolution in aiohttp client helper.
+    ("zeroconf", {"zeroconf"}, STAGE_0_SUBSTAGE_TIMEOUT),
+)
+
+DISCOVERY_INTEGRATIONS = ("bluetooth", "dhcp", "ssdp", "usb")
+# Stage 1 integrations are not to be preimported in bootstrap.
 STAGE_1_INTEGRATIONS = {
     # We need to make sure discovery integrations
     # update their deps before stage 2 integrations
@@ -189,9 +202,8 @@ STAGE_1_INTEGRATIONS = {
     "mqtt_eventstream",
     # To provide account link implementations
     "cloud",
-    # Ensure supervisor is available
-    "hassio",
 }
+
 DEFAULT_INTEGRATIONS = {
     # These integrations are set up unless recovery mode is activated.
     #
@@ -232,22 +244,12 @@ DEFAULT_INTEGRATIONS_SUPERVISOR = {
     # These integrations are set up if using the Supervisor
     "hassio",
 }
+
 CRITICAL_INTEGRATIONS = {
     # Recovery mode is activated if these integrations fail to set up
     "frontend",
 }
 
-SETUP_ORDER = (
-    # Load logging and http deps as soon as possible
-    ("logging, http deps", LOGGING_AND_HTTP_DEPS_INTEGRATIONS),
-    # Setup frontend
-    ("frontend", FRONTEND_INTEGRATIONS),
-    # Setup recorder
-    ("recorder", RECORDER_INTEGRATIONS),
-    # Start up debuggers. Start these first in case they want to wait.
-    ("debugger", DEBUGGER_INTEGRATIONS),
-)
-
 #
 # Storage keys we are likely to load during startup
 # in order of when we expect to load them.
@@ -694,7 +696,6 @@ async def async_mount_local_lib_path(config_dir: str) -> str:
     return deps_dir
 
 
-@core.callback
 def _get_domains(hass: core.HomeAssistant, config: dict[str, Any]) -> set[str]:
     """Get domains of components to set up."""
     # Filter out the repeating and common config section [homeassistant]
@@ -890,69 +891,48 @@ async def _async_set_up_integrations(
     domains_to_setup, integration_cache = await _async_resolve_domains_to_setup(
         hass, config
     )
+    stage_2_domains = domains_to_setup.copy()
 
     # Initialize recorder
     if "recorder" in domains_to_setup:
         recorder.async_initialize_recorder(hass)
 
-    pre_stage_domains = [
+    stage_0_and_1_domains: list[tuple[str, set[str], int | None]] = [
-        (name, domains_to_setup & domain_group) for name, domain_group in SETUP_ORDER
+        *(
+            (name, domain_group & domains_to_setup, timeout)
+            for name, domain_group, timeout in STAGE_0_INTEGRATIONS
+        ),
+        ("stage 1", STAGE_1_INTEGRATIONS & domains_to_setup, STAGE_1_TIMEOUT),
     ]
 
-    # calculate what components to setup in what stage
+    _LOGGER.info("Setting up stage 0 and 1")
-    stage_1_domains: set[str] = set()
+    for name, domain_group, timeout in stage_0_and_1_domains:
+        if not domain_group:
+            continue
 
-    # Find all dependencies of any dependency of any stage 1 integration that
+        _LOGGER.info("Setting up %s: %s", name, domain_group)
-    # we plan on loading and promote them to stage 1. This is done only to not
+        to_be_loaded = domain_group.copy()
-    # get misleading log messages
+        to_be_loaded.update(
-    deps_promotion: set[str] = STAGE_1_INTEGRATIONS
+            dep
-    while deps_promotion:
+            for domain in domain_group
-        old_deps_promotion = deps_promotion
+            if (integration := integration_cache.get(domain)) is not None
-        deps_promotion = set()
+            for dep in integration.all_dependencies
+        )
+        async_set_domains_to_be_loaded(hass, to_be_loaded)
+        stage_2_domains -= to_be_loaded
 
-        for domain in old_deps_promotion:
+        if timeout is None:
-            if domain not in domains_to_setup or domain in stage_1_domains:
-                continue
-
-            stage_1_domains.add(domain)
-
-            if (dep_itg := integration_cache.get(domain)) is None:
-                continue
-
-            deps_promotion.update(dep_itg.all_dependencies)
-
-    stage_2_domains = domains_to_setup - stage_1_domains
-
-    for name, domain_group in pre_stage_domains:
-        if domain_group:
-            stage_2_domains -= domain_group
-            _LOGGER.info("Setting up %s: %s", name, domain_group)
-            to_be_loaded = domain_group.copy()
-            to_be_loaded.update(
-                dep
-                for domain in domain_group
-                if (integration := integration_cache.get(domain)) is not None
-                for dep in integration.all_dependencies
-            )
-            async_set_domains_to_be_loaded(hass, to_be_loaded)
             await _async_setup_multi_components(hass, domain_group, config)
-
+        else:
-    # Enables after dependencies when setting up stage 1 domains
+            try:
-    async_set_domains_to_be_loaded(hass, stage_1_domains)
+                async with hass.timeout.async_timeout(timeout, cool_down=COOLDOWN_TIME):
-
+                    await _async_setup_multi_components(hass, domain_group, config)
-    # Start setup
+            except TimeoutError:
-    if stage_1_domains:
+                _LOGGER.warning(
-        _LOGGER.info("Setting up stage 1: %s", stage_1_domains)
+                    "Setup timed out for %s waiting on %s - moving forward",
-        try:
+                    name,
-            async with hass.timeout.async_timeout(
+                    hass._active_tasks,  # noqa: SLF001
-                STAGE_1_TIMEOUT, cool_down=COOLDOWN_TIME
+                )
-            ):
-                await _async_setup_multi_components(hass, stage_1_domains, config)
-        except TimeoutError:
-            _LOGGER.warning(
-                "Setup timed out for stage 1 waiting on %s - moving forward",
-                hass._active_tasks,  # noqa: SLF001
-            )
 
     # Add after dependencies when setting up stage 2 domains
     async_set_domains_to_be_loaded(hass, stage_2_domains)

tests/test_bootstrap.py

@@ -1090,7 +1090,7 @@ async def test_tasks_logged_that_block_stage_1(
         patch.object(bootstrap, "STAGE_1_TIMEOUT", 0),
         patch.object(bootstrap, "COOLDOWN_TIME", 0),
         patch.object(
-            bootstrap, "STAGE_1_INTEGRATIONS", [*original_stage_1, "normal_integration"]
+            bootstrap, "STAGE_1_INTEGRATIONS", {*original_stage_1, "normal_integration"}
         ),
     ):
         await bootstrap._async_set_up_integrations(hass, {"normal_integration": {}})
@@ -1373,11 +1373,11 @@ async def test_pre_import_no_requirements(hass: HomeAssistant) -> None:
 
 
 @pytest.mark.timeout(20)
-async def test_bootstrap_does_not_preload_stage_1_integrations() -> None:
+async def test_bootstrap_does_not_preimport_stage_1_integrations() -> None:
-    """Test that the bootstrap does not preload stage 1 integrations.
+    """Test that the bootstrap does not preimport stage 1 integrations.
 
     If this test fails it means that stage1 integrations are being
-    loaded too soon and will not get their requirements updated
+    imported too soon and will not get their requirements updated
     before they are loaded at runtime.
     """
 
@@ -1391,13 +1391,9 @@ async def test_bootstrap_does_not_preload_stage_1_integrations() -> None:
     assert process.returncode == 0
     decoded_stdout = stdout.decode()
 
-    disallowed_integrations = bootstrap.STAGE_1_INTEGRATIONS.copy()
-    # zeroconf is a top level dep now
-    disallowed_integrations.remove("zeroconf")
-
     # Ensure no stage1 integrations have been imported
     # as a side effect of importing the pre-imports
-    for integration in disallowed_integrations:
+    for integration in bootstrap.STAGE_1_INTEGRATIONS:
         assert f"homeassistant.components.{integration}" not in decoded_stdout
 
 

tests/test_circular_imports.py

@@ -7,11 +7,8 @@ import pytest
 
 from homeassistant.bootstrap import (
     CORE_INTEGRATIONS,
-    DEBUGGER_INTEGRATIONS,
     DEFAULT_INTEGRATIONS,
-    FRONTEND_INTEGRATIONS,
+    STAGE_0_INTEGRATIONS,
-    LOGGING_AND_HTTP_DEPS_INTEGRATIONS,
-    RECORDER_INTEGRATIONS,
     STAGE_1_INTEGRATIONS,
 )
 
@@ -21,11 +18,12 @@ from homeassistant.bootstrap import (
     "component",
     sorted(
         {
-            *DEBUGGER_INTEGRATIONS,
             *CORE_INTEGRATIONS,
-            *LOGGING_AND_HTTP_DEPS_INTEGRATIONS,
+            *(
-            *FRONTEND_INTEGRATIONS,
+                domain
-            *RECORDER_INTEGRATIONS,
+                for name, domains, timeout in STAGE_0_INTEGRATIONS
+                for domain in domains
+            ),
             *STAGE_1_INTEGRATIONS,
             *DEFAULT_INTEGRATIONS,
         }