diff --git a/homeassistant/config_entries.py b/homeassistant/config_entries.py index cb99d6f944a..fc37b23d497 100644 --- a/homeassistant/config_entries.py +++ b/homeassistant/config_entries.py @@ -22,6 +22,8 @@ from random import randint from types import MappingProxyType from typing import TYPE_CHECKING, Any, Self, TypeVar, cast +from async_interrupt import interrupt + from . import data_entry_flow, loader from .components import persistent_notification from .const import EVENT_HOMEASSISTANT_STARTED, EVENT_HOMEASSISTANT_STOP, Platform @@ -948,6 +950,10 @@ current_entry: ContextVar[ConfigEntry | None] = ContextVar( ) +class FlowCancelledError(Exception): + """Error to indicate that a flow has been cancelled.""" + + class ConfigEntriesFlowManager(data_entry_flow.FlowManager): """Manage all the config entry flows that are in progress.""" @@ -962,7 +968,7 @@ class ConfigEntriesFlowManager(data_entry_flow.FlowManager): self.config_entries = config_entries self._hass_config = hass_config self._pending_import_flows: dict[str, dict[str, asyncio.Future[None]]] = {} - self._initialize_tasks: dict[str, list[asyncio.Task]] = {} + self._initialize_futures: dict[str, list[asyncio.Future[None]]] = {} self._discovery_debouncer = Debouncer( hass, _LOGGER, @@ -994,20 +1000,26 @@ class ConfigEntriesFlowManager(data_entry_flow.FlowManager): raise KeyError("Context not set or doesn't have a source set") flow_id = uuid_util.random_uuid_hex() + loop = self.hass.loop + if context["source"] == SOURCE_IMPORT: - init_done: asyncio.Future[None] = self.hass.loop.create_future() - self._pending_import_flows.setdefault(handler, {})[flow_id] = init_done - - task = asyncio.create_task( - self._async_init(flow_id, handler, context, data), - name=f"config entry flow {handler} {flow_id}", - ) - self._initialize_tasks.setdefault(handler, []).append(task) + self._pending_import_flows.setdefault(handler, {})[ + flow_id + ] = loop.create_future() + cancel_init_future = loop.create_future() + self._initialize_futures.setdefault(handler, []).append(cancel_init_future) try: - flow, result = await task + async with interrupt( + cancel_init_future, + FlowCancelledError, + "Config entry initialize canceled: Home Assistant is shutting down", + ): + flow, result = await self._async_init(flow_id, handler, context, data) + except FlowCancelledError as ex: + raise asyncio.CancelledError from ex finally: - self._initialize_tasks[handler].remove(task) + self._initialize_futures[handler].remove(cancel_init_future) self._pending_import_flows.get(handler, {}).pop(flow_id, None) if result["type"] != data_entry_flow.FlowResultType.ABORT: @@ -1042,11 +1054,9 @@ class ConfigEntriesFlowManager(data_entry_flow.FlowManager): async def async_shutdown(self) -> None: """Cancel any initializing flows.""" - for task_list in self._initialize_tasks.values(): - for task in task_list: - task.cancel( - "Config entry initialize canceled: Home Assistant is shutting down" - ) + for future_list in self._initialize_futures.values(): + for future in future_list: + future.set_result(None) await self._discovery_debouncer.async_shutdown() async def async_finish_flow( diff --git a/homeassistant/package_constraints.txt b/homeassistant/package_constraints.txt index 747ec5c9a43..610845c7b29 100644 --- a/homeassistant/package_constraints.txt +++ b/homeassistant/package_constraints.txt @@ -7,6 +7,7 @@ aiohttp-zlib-ng==0.3.1 aiohttp==3.9.3 aiohttp_cors==0.7.0 astral==2.2 +async-interrupt==1.1.1 async-upnp-client==0.38.2 atomicwrites-homeassistant==1.4.1 attrs==23.2.0 diff --git a/pyproject.toml b/pyproject.toml index 6cef9b36ebf..c3e4b17e712 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ dependencies = [ "aiohttp-fast-url-dispatcher==0.3.0", "aiohttp-zlib-ng==0.3.1", "astral==2.2", + "async-interrupt==1.1.1", "attrs==23.2.0", "atomicwrites-homeassistant==1.4.1", "awesomeversion==24.2.0", diff --git a/requirements.txt b/requirements.txt index 9bdc71b0388..1fd6e5f9f08 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,6 +8,7 @@ aiohttp_cors==0.7.0 aiohttp-fast-url-dispatcher==0.3.0 aiohttp-zlib-ng==0.3.1 astral==2.2 +async-interrupt==1.1.1 attrs==23.2.0 atomicwrites-homeassistant==1.4.1 awesomeversion==24.2.0