Avoid creating tasks for dependencies already being setup (#111034)

Co-authored-by: Paulus Schoutsen <balloob@gmail.com>

homeassistant/config_entries.py

@@ -1756,7 +1756,10 @@ class ConfigEntries:
         Config entries which are created after Home Assistant is started can't be waited
         for, the function will just return if the config entry is loaded or not.
         """
-        if setup_future := self.hass.data.get(DATA_SETUP_DONE, {}).get(entry.domain):
+        setup_done: dict[str, asyncio.Future[bool]] = self.hass.data.get(
+            DATA_SETUP_DONE, {}
+        )
+        if setup_future := setup_done.get(entry.domain):
             await setup_future
         # The component was not loaded.
         if entry.domain not in self.hass.config.components:

homeassistant/setup.py

@@ -35,7 +35,7 @@ ATTR_COMPONENT: Final = "component"
 
 BASE_PLATFORMS = {platform.value for platform in Platform}
 
-# DATA_SETUP is a dict[str, asyncio.Task[bool]], indicating domains which are currently
+# DATA_SETUP is a dict[str, asyncio.Future[bool]], indicating domains which are currently
 # being setup or which failed to setup:
 # - Tasks are added to DATA_SETUP by `async_setup_component`, the key is the domain
 #   being setup and the Task is the `_async_setup_component` helper.
@@ -43,7 +43,7 @@ BASE_PLATFORMS = {platform.value for platform in Platform}
 #   the task returned True.
 DATA_SETUP = "setup_tasks"
 
-# DATA_SETUP_DONE is a dict [str, asyncio.Future], indicating components which
+# DATA_SETUP_DONE is a dict [str, asyncio.Future[bool]], indicating components which
 # will be setup:
 # - Events are added to DATA_SETUP_DONE during bootstrap by
 #   async_set_domains_to_be_loaded, the key is the domain which will be loaded.
@@ -51,7 +51,7 @@ DATA_SETUP = "setup_tasks"
 #   is finished, regardless of if the setup was successful or not.
 DATA_SETUP_DONE = "setup_done"
 
-# DATA_SETUP_DONE is a dict [str, datetime], indicating when an attempt
+# DATA_SETUP_STARTED is a dict [str, float], indicating when an attempt
 # to setup a component started.
 DATA_SETUP_STARTED = "setup_started"
 
@@ -116,10 +116,10 @@ def async_set_domains_to_be_loaded(hass: core.HomeAssistant, domains: set[str])
      - Properly handle after_dependencies.
      - Keep track of domains which will load but have not yet finished loading
     """
-    hass.data.setdefault(DATA_SETUP_DONE, {})
+    setup_done_futures: dict[str, asyncio.Future[bool]] = hass.data.setdefault(
-    hass.data[DATA_SETUP_DONE].update(
+        DATA_SETUP_DONE, {}
-        {domain: hass.loop.create_future() for domain in domains}
     )
+    setup_done_futures.update({domain: hass.loop.create_future() for domain in domains})
 
 
 def setup_component(hass: core.HomeAssistant, domain: str, config: ConfigType) -> bool:
@@ -142,23 +142,40 @@ async def async_setup_component(
     setup_futures: dict[str, asyncio.Future[bool]] = hass.data.setdefault(
         DATA_SETUP, {}
     )
+    setup_done_futures: dict[str, asyncio.Future[bool]] = hass.data.setdefault(
+        DATA_SETUP_DONE, {}
+    )
 
-    if existing_future := setup_futures.get(domain):
+    if existing_setup_future := setup_futures.get(domain):
-        return await existing_future
+        return await existing_setup_future
 
-    future = hass.loop.create_future()
+    setup_future = hass.loop.create_future()
-    setup_futures[domain] = future
+    setup_futures[domain] = setup_future
 
     try:
         result = await _async_setup_component(hass, domain, config)
-        future.set_result(result)
+        setup_future.set_result(result)
+        if setup_done_future := setup_done_futures.pop(domain, None):
+            setup_done_future.set_result(result)
         return result
     except BaseException as err:  # pylint: disable=broad-except
+        futures = [setup_future]
+        if setup_done_future := setup_done_futures.pop(domain, None):
+            futures.append(setup_done_future)
+        for future in futures:
+            # If the setup call is cancelled it likely means
+            # Home Assistant is shutting down so the future might
+            # already be done which will cause this to raise
+            # an InvalidStateError which is appropriate because
+            # the component setup was cancelled and is in an
+            # indeterminate state.
             future.set_exception(err)
+            with contextlib.suppress(BaseException):
+                # Clear the flag as its normal that nothing
+                # will wait for this future to be resolved
+                # if there are no concurrent setup attempts
+                await future
         raise
-    finally:
-        if future := hass.data.get(DATA_SETUP_DONE, {}).pop(domain, None):
-            future.set_result(None)
 
 
 async def _async_process_dependencies(
@@ -168,14 +185,22 @@ async def _async_process_dependencies(
 
     Returns a list of dependencies which failed to set up.
     """
+    setup_futures: dict[str, asyncio.Future[bool]] = hass.data.setdefault(
+        DATA_SETUP, {}
+    )
+
     dependencies_tasks = {
-        dep: hass.loop.create_task(async_setup_component(hass, dep, config))
+        dep: setup_futures.get(dep)
+        or hass.loop.create_task(
+            async_setup_component(hass, dep, config),
+            name=f"setup {dep} as dependency of {integration.domain}",
+        )
         for dep in integration.dependencies
         if dep not in hass.config.components
     }
 
-    after_dependencies_tasks: dict[str, asyncio.Future[None]] = {}
+    after_dependencies_tasks: dict[str, asyncio.Future[bool]] = {}
-    to_be_loaded: dict[str, asyncio.Future[None]] = hass.data.get(DATA_SETUP_DONE, {})
+    to_be_loaded: dict[str, asyncio.Future[bool]] = hass.data.get(DATA_SETUP_DONE, {})
     for dep in integration.after_dependencies:
         if (
             dep not in dependencies_tasks
@@ -191,13 +216,13 @@ async def _async_process_dependencies(
         _LOGGER.debug(
             "Dependency %s will wait for dependencies %s",
             integration.domain,
-            list(dependencies_tasks),
+            dependencies_tasks.keys(),
         )
     if after_dependencies_tasks:
         _LOGGER.debug(
             "Dependency %s will wait for after dependencies %s",
             integration.domain,
-            list(after_dependencies_tasks),
+            after_dependencies_tasks.keys(),
         )
 
     async with hass.timeout.async_freeze(integration.domain):
@@ -213,7 +238,7 @@ async def _async_process_dependencies(
         _LOGGER.error(
             "Unable to set up dependencies of '%s'. Setup failed for dependencies: %s",
             integration.domain,
-            ", ".join(failed),
+            failed,
         )
 
     return failed

tests/test_bootstrap.py

@@ -926,7 +926,7 @@ async def test_bootstrap_dependencies(
         """Assert the mqtt config entry was set up."""
         calls.append("mqtt")
         # assert the integration is not yet set up
-        assertions.append(hass.data["setup_done"][integration].is_set() is False)
+        assertions.append(hass.data["setup_done"][integration].done() is False)
         assertions.append(
             all(
                 dependency in hass.config.components
@@ -942,7 +942,7 @@ async def test_bootstrap_dependencies(
         # assert mqtt was already set up
         assertions.append(
             "mqtt" not in hass.data["setup_done"]
-            or hass.data["setup_done"]["mqtt"].is_set()
+            or hass.data["setup_done"]["mqtt"].done()
         )
         assertions.append("mqtt" in hass.config.components)
         return True
@@ -1029,5 +1029,6 @@ async def test_bootstrap_dependencies(
     assert calls == ["mqtt", integration]
 
     assert (
-        f"Dependency {integration} will wait for dependencies ['mqtt']" in caplog.text
+        f"Dependency {integration} will wait for dependencies dict_keys(['mqtt'])"
+        in caplog.text
     )

tests/test_setup.py

@@ -319,6 +319,7 @@ async def test_component_failing_setup(hass: HomeAssistant) -> None:
 
 async def test_component_exception_setup(hass: HomeAssistant) -> None:
     """Test component that raises exception during setup."""
+    setup.async_set_domains_to_be_loaded(hass, {"comp"})
 
     def exception_setup(hass, config):
         """Raise exception."""
@@ -330,6 +331,22 @@ async def test_component_exception_setup(hass: HomeAssistant) -> None:
     assert "comp" not in hass.config.components
 
 
+async def test_component_base_exception_setup(hass: HomeAssistant) -> None:
+    """Test component that raises exception during setup."""
+    setup.async_set_domains_to_be_loaded(hass, {"comp"})
+
+    def exception_setup(hass, config):
+        """Raise exception."""
+        raise BaseException("fail!")
+
+    mock_integration(hass, MockModule("comp", setup=exception_setup))
+
+    with pytest.raises(BaseException):
+        await setup.async_setup_component(hass, "comp", {})
+
+    assert "comp" not in hass.config.components
+
+
 async def test_component_setup_with_validation_and_dependency(
     hass: HomeAssistant,
 ) -> None: