diff --git a/homeassistant/bootstrap.py b/homeassistant/bootstrap.py index 446fe73f9263..f08f1f5ecac2 100644 --- a/homeassistant/bootstrap.py +++ b/homeassistant/bootstrap.py @@ -868,7 +868,7 @@ async def _async_set_up_integrations( _LOGGER.debug("Waiting for startup to wrap up") try: async with hass.timeout.async_timeout(WRAP_UP_TIMEOUT, cool_down=COOLDOWN_TIME): - await hass.async_block_till_done() + await hass.async_block_till_done(wait_periodic_tasks=False) except TimeoutError: _LOGGER.warning("Setup timed out for bootstrap - moving forward") diff --git a/homeassistant/core.py b/homeassistant/core.py index f0b107602288..14b8c7f2f3a5 100644 --- a/homeassistant/core.py +++ b/homeassistant/core.py @@ -23,6 +23,7 @@ import datetime import enum import functools import inspect +from itertools import chain import logging import os import pathlib @@ -382,6 +383,7 @@ class HomeAssistant: self.loop = asyncio.get_running_loop() self._tasks: set[asyncio.Future[Any]] = set() self._background_tasks: set[asyncio.Future[Any]] = set() + self._periodic_tasks: set[asyncio.Future[Any]] = set() self.bus = EventBus(self) self.services = ServiceRegistry(self) self.states = StateMachine(self.bus, self.loop) @@ -669,9 +671,17 @@ class HomeAssistant: ) -> asyncio.Task[_R]: """Create a task from within the event loop. - This is a background task which will not block startup and will be - automatically cancelled on shutdown. If you are using this in your - integration, use the create task methods on the config entry instead. + This type of task is for background tasks that usually run for + the lifetime of Home Assistant or an integration's setup. + + This is a background task which is different from a normal task: + + - Will not block startup + - Will be automatically cancelled on shutdown + - Calls to async_block_till_done will not wait for completion + + If you are using this in your integration, use the create task + methods on the config entry instead. This method must be run in the event loop. """ @@ -687,6 +697,38 @@ class HomeAssistant: task.add_done_callback(self._background_tasks.remove) return task + @callback + def async_create_periodic_task( + self, target: Coroutine[Any, Any, _R], name: str, eager_start: bool = False + ) -> asyncio.Task[_R]: + """Create a task from within the event loop. + + This type of tasks is for periodic updates such as polling + entities. + + This is a periodic task which is different from a normal task: + + - Will not block startup + - Will be automatically cancelled on shutdown + - Calls to async_block_till_done will wait for completion by default + + If you are using this in your integration, use the create task + methods on the config entry instead. + + This method must be run in the event loop. + """ + if eager_start: + task = create_eager_task(target, name=name, loop=self.loop) + if task.done(): + return task + else: + # Use loop.create_task + # to avoid the extra function call in asyncio.create_task. + task = self.loop.create_task(target, name=name) + self._periodic_tasks.add(task) + task.add_done_callback(self._periodic_tasks.remove) + return task + @callback def async_add_executor_job( self, target: Callable[..., _T], *args: Any @@ -796,16 +838,19 @@ class HomeAssistant: self.async_block_till_done(), self.loop ).result() - async def async_block_till_done(self) -> None: + async def async_block_till_done(self, wait_periodic_tasks: bool = True) -> None: """Block until all pending work is done.""" # To flush out any call_soon_threadsafe await asyncio.sleep(0) start_time: float | None = None current_task = asyncio.current_task() + to_wait: Iterable[asyncio.Future[Any]] = self._tasks + if wait_periodic_tasks: + to_wait = chain(self._tasks, self._periodic_tasks) while tasks := [ task - for task in self._tasks + for task in to_wait if task is not current_task and not cancelling(task) ]: await self._await_and_log_pending(tasks) @@ -936,7 +981,7 @@ class HomeAssistant: self._tasks = set() # Cancel all background tasks - for task in self._background_tasks: + for task in chain(self._background_tasks, self._periodic_tasks): self._tasks.add(task) task.add_done_callback(self._tasks.remove) task.cancel("Home Assistant is stopping") @@ -948,7 +993,7 @@ class HomeAssistant: self.bus.async_fire(EVENT_HOMEASSISTANT_STOP) try: async with self.timeout.async_timeout(STOP_STAGE_SHUTDOWN_TIMEOUT): - await self.async_block_till_done() + await self.async_block_till_done(wait_periodic_tasks=False) except TimeoutError: _LOGGER.warning( "Timed out waiting for integrations to stop, the shutdown will" @@ -961,7 +1006,7 @@ class HomeAssistant: self.bus.async_fire(EVENT_HOMEASSISTANT_FINAL_WRITE) try: async with self.timeout.async_timeout(FINAL_WRITE_STAGE_SHUTDOWN_TIMEOUT): - await self.async_block_till_done() + await self.async_block_till_done(wait_periodic_tasks=False) except TimeoutError: _LOGGER.warning( "Timed out waiting for final writes to complete, the shutdown will" @@ -1013,7 +1058,7 @@ class HomeAssistant: try: async with self.timeout.async_timeout(CLOSE_STAGE_SHUTDOWN_TIMEOUT): - await self.async_block_till_done() + await self.async_block_till_done(wait_periodic_tasks=False) except TimeoutError: _LOGGER.warning( "Timed out waiting for close event to be processed, the shutdown will" diff --git a/tests/test_core.py b/tests/test_core.py index 75d06a7c61fe..0e3f42db9155 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -514,7 +514,7 @@ async def test_shutdown_calls_block_till_done_after_shutdown_run_callback_thread """Ensure shutdown_run_callback_threadsafe is called before the final async_block_till_done.""" stop_calls = [] - async def _record_block_till_done(): + async def _record_block_till_done(wait_periodic_tasks: bool = True): nonlocal stop_calls stop_calls.append("async_block_till_done") @@ -2098,9 +2098,9 @@ async def test_chained_logging_hits_log_timeout( return hass.async_create_task(_task_chain_1()) - with patch.object(ha, "BLOCK_LOG_TIMEOUT", 0.0001): + with patch.object(ha, "BLOCK_LOG_TIMEOUT", 0.0): hass.async_create_task(_task_chain_1()) - await hass.async_block_till_done() + await hass.async_block_till_done(wait_periodic_tasks=False) assert "_task_chain_" in caplog.text