evlli/hass-core
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Speed up initializing config flows (#124015)

Browse source
This commit is contained in:
J. Nick Koston 2024-08-16 16:48:03 -05:00 • committed by GitHub
parent 24680b731f
commit 91951ed734
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 28 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

43
homeassistant/config_entries.py
View file

@ -3,7 +3,7 @@
from __future__ import annotations
import asyncio
from collections import UserDict
from collections import UserDict, defaultdict
from collections.abc import (
Callable,
Coroutine,
@ -1224,8 +1224,12 @@ class ConfigEntriesFlowManager(data_entry_flow.FlowManager[ConfigFlowResult]):
super().__init__(hass)
self.config_entries = config_entries
self._hass_config = hass_config
self._pending_import_flows: dict[str, dict[str, asyncio.Future[None]]] = {}
self._initialize_futures: dict[str, list[asyncio.Future[None]]] = {}
self._pending_import_flows: defaultdict[
str, dict[str, asyncio.Future[None]]
] = defaultdict(dict)
self._initialize_futures: defaultdict[str, set[asyncio.Future[None]]] = (
defaultdict(set)
)
self._discovery_debouncer = Debouncer[None](
hass,
_LOGGER,
@ -1278,12 +1282,11 @@ class ConfigEntriesFlowManager(data_entry_flow.FlowManager[ConfigFlowResult]):
loop = self.hass.loop
if context["source"] == SOURCE_IMPORT:
self._pending_import_flows.setdefault(handler, {})[flow_id] = (
loop.create_future()
)
self._pending_import_flows[handler][flow_id] = loop.create_future()
cancel_init_future = loop.create_future()
self._initialize_futures.setdefault(handler, []).append(cancel_init_future)
handler_init_futures = self._initialize_futures[handler]
handler_init_futures.add(cancel_init_future)
try:
async with interrupt(
cancel_init_future,
@ -1294,8 +1297,13 @@ class ConfigEntriesFlowManager(data_entry_flow.FlowManager[ConfigFlowResult]):
except FlowCancelledError as ex:
raise asyncio.CancelledError from ex
finally:
self._initialize_futures[handler].remove(cancel_init_future)
self._pending_import_flows.get(handler, {}).pop(flow_id, None)
handler_init_futures.remove(cancel_init_future)
if not handler_init_futures:
del self._initialize_futures[handler]
if handler in self._pending_import_flows:
self._pending_import_flows[handler].pop(flow_id, None)
if not self._pending_import_flows[handler]:
del self._pending_import_flows[handler]
if result["type"] != data_entry_flow.FlowResultType.ABORT:
await self.async_post_init(flow, result)
@ -1322,11 +1330,18 @@ class ConfigEntriesFlowManager(data_entry_flow.FlowManager[ConfigFlowResult]):
try:
result = await self._async_handle_step(flow, flow.init_step, data)
finally:
init_done = self._pending_import_flows.get(handler, {}).get(flow_id)
if init_done and not init_done.done():
init_done.set_result(None)
self._set_pending_import_done(flow)
return flow, result
def _set_pending_import_done(self, flow: ConfigFlow) -> None:
"""Set pending import flow as done."""
if (
(handler_import_flows := self._pending_import_flows.get(flow.handler))
and (init_done := handler_import_flows.get(flow.flow_id))
and not init_done.done()
):
init_done.set_result(None)
@callback
def async_shutdown(self) -> None:
"""Cancel any initializing flows."""
@ -1347,9 +1362,7 @@ class ConfigEntriesFlowManager(data_entry_flow.FlowManager[ConfigFlowResult]):
# We do this to avoid a circular dependency where async_finish_flow sets up a
# new entry, which needs the integration to be set up, which is waiting for
# init to be done.
init_done = self._pending_import_flows.get(flow.handler, {}).get(flow.flow_id)
if init_done and not init_done.done():
init_done.set_result(None)
self._set_pending_import_done(flow)
# Remove notification if no other discovery config entries in progress
if not self._async_has_other_discovery_flows(flow.flow_id):