diff --git a/homeassistant/components/websocket_api/decorators.py b/homeassistant/components/websocket_api/decorators.py index 296271c7cfd..9afffd9fb28 100644 --- a/homeassistant/components/websocket_api/decorators.py +++ b/homeassistant/components/websocket_api/decorators.py @@ -1,7 +1,6 @@ """Decorators for the Websocket API.""" from __future__ import annotations -import asyncio from collections.abc import Callable from functools import wraps from typing import Any @@ -33,6 +32,7 @@ def async_response( func: const.AsyncWebSocketCommandHandler, ) -> const.WebSocketCommandHandler: """Decorate an async function to handle WebSocket API messages.""" + task_name = f"websocket_api.async:{func.__name__}" @callback @wraps(func) @@ -42,7 +42,10 @@ def async_response( """Schedule the handler.""" # As the webserver is now started before the start # event we do not want to block for websocket responders - asyncio.create_task(_handle_async_response(func, hass, connection, msg)) + hass.async_create_background_task( + _handle_async_response(func, hass, connection, msg), + task_name, + ) return schedule_handler diff --git a/homeassistant/core.py b/homeassistant/core.py index 15f05ff608c..99b7218a9ff 100644 --- a/homeassistant/core.py +++ b/homeassistant/core.py @@ -279,6 +279,7 @@ class HomeAssistant: """Initialize new Home Assistant object.""" self.loop = asyncio.get_running_loop() self._tasks: set[asyncio.Future[Any]] = set() + self._background_tasks: set[asyncio.Future[Any]] = set() self.bus = EventBus(self) self.services = ServiceRegistry(self) self.states = StateMachine(self.bus, self.loop) @@ -520,7 +521,26 @@ class HomeAssistant: task = self.loop.create_task(target) self._tasks.add(task) task.add_done_callback(self._tasks.remove) + return task + @callback + def async_create_background_task( + self, + target: Coroutine[Any, Any, _R], + name: str, + ) -> asyncio.Task[_R]: + """Create a task from within the eventloop. + + 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, make sure you also cancel the task when the config entry + your task belongs to is unloaded. + + This method must be run in the event loop. + """ + task = self.loop.create_task(target, name=name) + self._background_tasks.add(task) + task.add_done_callback(self._background_tasks.remove) return task @callback @@ -687,6 +707,12 @@ class HomeAssistant: "Stopping Home Assistant before startup has completed may fail" ) + # Cancel all background tasks + for task in self._background_tasks: + self._tasks.add(task) + task.add_done_callback(self._tasks.remove) + task.cancel() + # stage 1 self.state = CoreState.stopping self.bus.async_fire(EVENT_HOMEASSISTANT_STOP) diff --git a/tests/test_core.py b/tests/test_core.py index 1e0e69cc3af..6268f7f4ac1 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -1933,3 +1933,21 @@ async def test_state_changed_events_to_not_leak_contexts(hass: HomeAssistant) -> gc.collect() assert len(_get_by_type("homeassistant.core.Context")) == init_count + + +async def test_background_task(hass): + """Test background tasks being quit.""" + result = asyncio.Future() + + async def test_task(): + try: + await asyncio.sleep(1) + except asyncio.CancelledError: + result.set_result(hass.state) + raise + + task = hass.async_create_background_task(test_task(), "happy task") + assert "happy task" in str(task) + await asyncio.sleep(0) + await hass.async_stop() + assert result.result() == ha.CoreState.stopping