Use an eager task to setup entity platforms (#111574)

* Use an eager task to setup entity platforms

Ideally we would have awaited this function instead, but we want
to shield it from cancellation so we wrap it in asyncio.shield
which schedules it as a task. Since we have integrations that
never suspend in async_setup_entry, we can avoid scheduling on
the evnet loop with an eager task for this case

* its an executor future

* its an executor future

* fix

* doc string lied
This commit is contained in:
J. Nick Koston 2024-02-27 03:45:45 -10:00 committed by GitHub
parent 224ca122fd
commit 08e0008d31
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -260,7 +260,7 @@ class EntityPlatform:
return return
@callback @callback
def async_create_setup_task() -> ( def async_create_setup_awaitable() -> (
Coroutine[Any, Any, None] | asyncio.Future[None] Coroutine[Any, Any, None] | asyncio.Future[None]
): ):
"""Get task to set up platform.""" """Get task to set up platform."""
@ -283,7 +283,7 @@ class EntityPlatform:
discovery_info, discovery_info,
) )
await self._async_setup_platform(async_create_setup_task) await self._async_setup_platform(async_create_setup_awaitable)
@callback @callback
def async_shutdown(self) -> None: def async_shutdown(self) -> None:
@ -305,7 +305,7 @@ class EntityPlatform:
platform = self.platform platform = self.platform
@callback @callback
def async_create_setup_task() -> Coroutine[Any, Any, None]: def async_create_setup_awaitable() -> Coroutine[Any, Any, None]:
"""Get task to set up platform.""" """Get task to set up platform."""
config_entries.current_entry.set(config_entry) config_entries.current_entry.set(config_entry)
@ -313,14 +313,16 @@ class EntityPlatform:
self.hass, config_entry, self._async_schedule_add_entities_for_entry self.hass, config_entry, self._async_schedule_add_entities_for_entry
) )
return await self._async_setup_platform(async_create_setup_task) return await self._async_setup_platform(async_create_setup_awaitable)
async def _async_setup_platform( async def _async_setup_platform(
self, async_create_setup_task: Callable[[], Awaitable[None]], tries: int = 0 self,
async_create_setup_awaitable: Callable[[], Awaitable[None]],
tries: int = 0,
) -> bool: ) -> bool:
"""Set up a platform via config file or config entry. """Set up a platform via config file or config entry.
async_create_setup_task creates a coroutine that sets up platform. async_create_setup_awaitable creates an awaitable that sets up platform.
""" """
current_platform.set(self) current_platform.set(self)
logger = self.logger logger = self.logger
@ -340,10 +342,12 @@ class EntityPlatform:
) )
with async_start_setup(hass, [full_name]): with async_start_setup(hass, [full_name]):
try: try:
task = async_create_setup_task() awaitable = async_create_setup_awaitable()
if asyncio.iscoroutine(awaitable):
awaitable = create_eager_task(awaitable)
async with hass.timeout.async_timeout(SLOW_SETUP_MAX_WAIT, self.domain): async with hass.timeout.async_timeout(SLOW_SETUP_MAX_WAIT, self.domain):
await asyncio.shield(task) await asyncio.shield(awaitable)
# Block till all entities are done # Block till all entities are done
while self._tasks: while self._tasks:
@ -379,7 +383,9 @@ class EntityPlatform:
async def setup_again(*_args: Any) -> None: async def setup_again(*_args: Any) -> None:
"""Run setup again.""" """Run setup again."""
self._async_cancel_retry_setup = None self._async_cancel_retry_setup = None
await self._async_setup_platform(async_create_setup_task, tries) await self._async_setup_platform(
async_create_setup_awaitable, tries
)
if hass.state is CoreState.running: if hass.state is CoreState.running:
self._async_cancel_retry_setup = async_call_later( self._async_cancel_retry_setup = async_call_later(