Add background tasks to config entries (#88335)

* Use a set for config entries task tracking

* Allow adding background tasks to config entries

* Add tests for config entry add tasks

* Update docstrings on core create task

* Migrate roon and august

* Use in more places

* Guard for None

homeassistant/components/august/__init__.py

@@ -191,8 +191,8 @@ class AugustData(AugustSubscriberMixin):
             # Do not prevent setup as the sync can timeout
             # but it is not a fatal error as the lock
             # will recover automatically when it comes back online.
-            self._config_entry.async_on_unload(
+            self._config_entry.async_create_background_task(
-                asyncio.create_task(self._async_initial_sync()).cancel
+                self._hass, self._async_initial_sync(), "august-initial-sync"
             )
 
     async def _async_initial_sync(self):

homeassistant/components/google/calendar.py

@@ -245,6 +245,8 @@ async def async_setup_entry(
 class CalendarSyncUpdateCoordinator(DataUpdateCoordinator[Timeline]):
     """Coordinator for calendar RPC calls that use an efficient sync."""
 
+    config_entry: ConfigEntry
+
     def __init__(
         self,
         hass: HomeAssistant,
@@ -299,6 +301,8 @@ class CalendarQueryUpdateCoordinator(DataUpdateCoordinator[list[Event]]):
     for limitations in the calendar API for supporting search.
     """
 
+    config_entry: ConfigEntry
+
     def __init__(
         self,
         hass: HomeAssistant,
@@ -434,7 +438,9 @@ class GoogleCalendarEntity(
             await self.coordinator.async_request_refresh()
             self._apply_coordinator_update()
 
-        self.hass.async_create_background_task(refresh(), "google.calendar-refresh")
+        self.coordinator.config_entry.async_create_background_task(
+            self.hass, refresh(), "google.calendar-refresh"
+        )
 
     async def async_get_events(
         self, hass: HomeAssistant, start_date: datetime, end_date: datetime

homeassistant/components/qnap_qsw/entity.py

@@ -145,7 +145,7 @@ class QswFirmwareEntity(CoordinatorEntity[QswFirmwareCoordinator]):
     def get_device_value(self, key: str, subkey: str) -> Any:
         """Return device value by key."""
         value = None
-        if key in self.coordinator.data:
+        if self.coordinator.data is not None and key in self.coordinator.data:
             data = self.coordinator.data[key]
             if subkey in data:
                 value = data[subkey]

homeassistant/components/roon/server.py

@@ -66,8 +66,8 @@ class RoonServer:
         )
 
         # Initialize Roon background polling
-        self.config_entry.async_on_unload(
+        self.config_entry.async_create_background_task(
-            asyncio.create_task(self.async_do_loop()).cancel
+            self.hass, self.async_do_loop(), "roon.server-do-loop"
         )
 
         return True

homeassistant/components/sonos/__init__.py

@@ -483,7 +483,8 @@ class SonosDiscoveryManager:
         if uid not in self.data.discovery_known:
             _LOGGER.debug("New %s discovery uid=%s: %s", source, uid, info)
             self.data.discovery_known.add(uid)
-        self.hass.async_create_background_task(
+        self.entry.async_create_background_task(
+            self.hass,
             self._async_handle_discovery_message(
                 uid,
                 discovered_ip,

homeassistant/components/zha/core/gateway.py

@@ -254,8 +254,8 @@ class ZHAGateway:
             )
 
         # background the fetching of state for mains powered devices
-        self._hass.async_create_background_task(
+        self.config_entry.async_create_background_task(
-            fetch_updated_state(), "zha.gateway-fetch_updated_state"
+            self._hass, fetch_updated_state(), "zha.gateway-fetch_updated_state"
         )
 
     def device_joined(self, device: zigpy.device.Device) -> None:

homeassistant/components/zha/light.py

@@ -612,7 +612,9 @@ class BaseLight(LogMixin, light.LightEntity):
             )
             if self._debounced_member_refresh is not None:
                 self.debug("transition complete - refreshing group member states")
-                self.hass.async_create_background_task(
+                assert self.platform and self.platform.config_entry
+                self.platform.config_entry.async_create_background_task(
+                    self.hass,
                     self._debounced_member_refresh.async_call(),
                     "zha.light-refresh-debounced-member",
                 )

homeassistant/config_entries.py

@@ -220,7 +220,8 @@ class ConfigEntry:
         "_async_cancel_retry_setup",
         "_on_unload",
         "reload_lock",
-        "_pending_tasks",
+        "_tasks",
+        "_background_tasks",
     )
 
     def __init__(
@@ -315,7 +316,8 @@ class ConfigEntry:
         # Reload lock to prevent conflicting reloads
         self.reload_lock = asyncio.Lock()
 
-        self._pending_tasks: list[asyncio.Future[Any]] = []
+        self._tasks: set[asyncio.Future[Any]] = set()
+        self._background_tasks: set[asyncio.Future[Any]] = set()
 
     async def async_setup(
         self,
@@ -681,11 +683,23 @@ class ConfigEntry:
             while self._on_unload:
                 self._on_unload.pop()()
 
-        while self._pending_tasks:
+        if not self._tasks and not self._background_tasks:
-            pending = [task for task in self._pending_tasks if not task.done()]
+            return
-            self._pending_tasks.clear()
+
-            if pending:
+        for task in self._background_tasks:
-                await asyncio.gather(*pending)
+            task.cancel()
+
+        _, pending = await asyncio.wait(
+            [*self._tasks, *self._background_tasks], timeout=10
+        )
+
+        for task in pending:
+            _LOGGER.warning(
+                "Unloading %s (%s) config entry. Task %s did not complete in time",
+                self.title,
+                self.domain,
+                task,
+            )
 
     @callback
     def async_start_reauth(
@@ -736,9 +750,24 @@ class ConfigEntry:
         target: target to call.
         """
         task = hass.async_create_task(target)
+        self._tasks.add(task)
+        task.add_done_callback(self._tasks.remove)
 
-        self._pending_tasks.append(task)
+        return task
 
+    @callback
+    def async_create_background_task(
+        self, hass: HomeAssistant, target: Coroutine[Any, Any, _R], name: str
+    ) -> asyncio.Task[_R]:
+        """Create a background task tied to the config entry lifecycle.
+
+        Background tasks are automatically canceled when config entry is unloaded.
+
+        target: target to call.
+        """
+        task = hass.async_create_background_task(target, name)
+        self._background_tasks.add(task)
+        task.add_done_callback(self._background_tasks.remove)
         return task
 
 

homeassistant/core.py

@@ -514,7 +514,8 @@ class HomeAssistant:
     def async_create_task(self, target: Coroutine[Any, Any, _R]) -> asyncio.Task[_R]:
         """Create a task from within the eventloop.
 
-        This method must be run in the event loop.
+        This method must be run in the event loop. If you are using this in your
+        integration, use the create task methods on the config entry instead.
 
         target: target to call.
         """
@@ -533,8 +534,7 @@ class HomeAssistant:
 
         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
+        integration, use the create task methods on the config entry instead.
-        your task belongs to is unloaded.
 
         This method must be run in the event loop.
         """

tests/test_config_entries.py

@@ -3573,3 +3573,26 @@ async def test_initializing_flows_canceled_on_shutdown(hass: HomeAssistant, mana
 
         with pytest.raises(asyncio.exceptions.CancelledError):
             await task
+
+
+async def test_task_tracking(hass):
+    """Test task tracking for a config entry."""
+    entry = MockConfigEntry(title="test_title", domain="test")
+
+    event = asyncio.Event()
+    results = []
+
+    async def test_task():
+        try:
+            await event.wait()
+            results.append("normal")
+        except asyncio.CancelledError:
+            results.append("background")
+            raise
+
+    entry.async_create_task(hass, test_task())
+    entry.async_create_background_task(hass, test_task(), "background-task-name")
+    await asyncio.sleep(0)
+    hass.loop.call_soon(event.set)
+    await entry._async_process_on_unload()
+    assert results == ["background", "normal"]