Add support for background tasks in HA (#88265)

* Add support for background tasks

* make name mandatory for background tasks

* Update docstring

* async_create_background_task

* Grammar

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
 

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)

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