Use wait_shared_future in setup (#183893)

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

homeassistant/setup.py

@@ -30,7 +30,7 @@ from .exceptions import DependencyError, HomeAssistantError
 from .helpers import issue_registry as ir, singleton, translation
 from .helpers.issue_registry import IssueSeverity, async_create_issue
 from .helpers.typing import ConfigType
-from .util.async_ import create_eager_task
+from .util.async_ import create_eager_task, wait_shared_future
 from .util.hass_dict import HassKey
 
 current_setup_group: contextvars.ContextVar[tuple[str, str | None] | None] = (
@@ -159,7 +159,7 @@ async def async_setup_component(
     setup_done_futures = hass.data.setdefault(_DATA_SETUP_DONE, {})
 
     if existing_setup_future := setup_futures.get(domain):
-        return await existing_setup_future
+        return await wait_shared_future(existing_setup_future)
 
     setup_future = hass.loop.create_future()
     setup_futures[domain] = setup_future
@@ -199,13 +199,15 @@ async def _async_process_dependencies(
     """
     setup_futures = hass.data.setdefault(_DATA_SETUP, {})
 
-    dependencies_tasks: dict[str, asyncio.Future[bool]] = {}
+    dependencies_tasks: dict[str, Awaitable[bool]] = {}
+    fut: Awaitable[bool]
 
     for dep in integration.dependencies:
-        fut = setup_futures.get(dep)
-        if fut is None:
-            if dep in hass.config.components:
-                continue
+        if (shared_fut := setup_futures.get(dep)) is not None:
+            fut = wait_shared_future(shared_fut)
+        elif dep in hass.config.components:
+            continue
+        else:
             fut = create_eager_task(
                 async_setup_component(hass, dep, config),
                 name=f"setup {dep} as dependency of {integration.domain}",
@@ -221,10 +223,11 @@ async def _async_process_dependencies(
     for dep in integration.after_dependencies:
         if dep not in to_be_loaded or dep in dependencies_tasks:
             continue
-        fut = setup_futures.get(dep)
-        if fut is None:
-            if dep in hass.config.components:
-                continue
+        if (shared_fut := setup_futures.get(dep)) is not None:
+            fut = wait_shared_future(shared_fut)
+        elif dep in hass.config.components:
+            continue
+        else:
             fut = create_eager_task(
                 async_setup_component(hass, dep, config),
                 name=f"setup {dep} as after dependency of {integration.domain}",
@@ -837,5 +840,5 @@ async def async_wait_component(hass: HomeAssistant, domain: str) -> bool:
     """Wait until a component is set up if pending, then return if it is set up."""
     setup_done = hass.data.get(_DATA_SETUP_DONE, {})
     if setup_future := setup_done.get(domain):
-        await setup_future
+        await wait_shared_future(setup_future)
     return domain in hass.config.components

tests/test_setup.py

@@ -337,6 +337,63 @@ async def test_component_not_setup_twice_if_loaded_during_other_setup(
     assert len(result) == 1
 
 
+async def test_component_setup_concurrent_waiter_cancelled(
+    hass: HomeAssistant,
+) -> None:
+    """Test cancelling a concurrent setup caller does not break the setup."""
+    setup_started = asyncio.Event()
+    setup_stall = asyncio.Event()
+
+    async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+        setup_started.set()
+        await setup_stall.wait()
+        return True
+
+    mock_integration(hass, MockModule("comp", async_setup=async_setup))
+
+    setup_task1 = asyncio.create_task(setup.async_setup_component(hass, "comp", {}))
+    await setup_started.wait()
+    setup_future = hass.data[setup._DATA_SETUP]["comp"]
+    setup_task2 = asyncio.create_task(setup.async_setup_component(hass, "comp", {}))
+    await asyncio.sleep(0)
+    setup_task2.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await setup_task2
+    setup_stall.set()
+
+    assert await setup_task1 is True
+    assert setup_future.result() is True
+
+
+async def test_component_setup_dependency_waiter_cancelled(
+    hass: HomeAssistant,
+) -> None:
+    """Test cancelling a component setup does not break its dependency setup."""
+    setup_started = asyncio.Event()
+    setup_stall = asyncio.Event()
+
+    async def async_setup_dep(hass: HomeAssistant, config: ConfigType) -> bool:
+        setup_started.set()
+        await setup_stall.wait()
+        return True
+
+    mock_integration(hass, MockModule("dep", async_setup=async_setup_dep))
+    mock_integration(hass, MockModule("comp", dependencies=["dep"]))
+
+    dep_task = asyncio.create_task(setup.async_setup_component(hass, "dep", {}))
+    await setup_started.wait()
+    setup_future = hass.data[setup._DATA_SETUP]["dep"]
+    comp_task = asyncio.create_task(setup.async_setup_component(hass, "comp", {}))
+    await hass.async_block_till_done()
+    comp_task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await comp_task
+    setup_stall.set()
+
+    assert await dep_task is True
+    assert setup_future.result() is True
+
+
 async def test_component_not_setup_missing_dependencies(hass: HomeAssistant) -> None:
     """Test we do not set up a component if not all dependencies loaded."""
     deps = ["maybe_existing"]
@@ -1449,3 +1506,30 @@ async def test_async_wait_component(hass: HomeAssistant) -> None:
     # Clear the event, then call again to make sure we don't block
     setup_stall.clear()
     assert await setup.async_wait_component(hass, "test") is True
+
+
+async def test_async_wait_component_waiter_cancelled(hass: HomeAssistant) -> None:
+    """Test cancelling async_wait_component does not break the setup."""
+    setup_started = asyncio.Event()
+    setup_stall = asyncio.Event()
+
+    async def mock_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+        setup_started.set()
+        await setup_stall.wait()
+        return True
+
+    mock_integration(hass, MockModule("test", async_setup=mock_setup))
+    setup.async_set_domains_to_be_loaded(hass, {"test"})
+    setup_done_future = hass.data[setup._DATA_SETUP_DONE]["test"]
+
+    setup_task = asyncio.create_task(setup.async_setup_component(hass, "test", {}))
+    await setup_started.wait()
+    wait_task = asyncio.create_task(setup.async_wait_component(hass, "test"))
+    await asyncio.sleep(0)
+    wait_task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await wait_task
+    setup_stall.set()
+
+    assert await setup_task is True
+    assert setup_done_future.result() is True