diff --git a/homeassistant/components/august/__init__.py b/homeassistant/components/august/__init__.py index 0f91879a23e..7839d879901 100644 --- a/homeassistant/components/august/__init__.py +++ b/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( - asyncio.create_task(self._async_initial_sync()).cancel + self._config_entry.async_create_background_task( + self._hass, self._async_initial_sync(), "august-initial-sync" ) async def _async_initial_sync(self): diff --git a/homeassistant/components/google/calendar.py b/homeassistant/components/google/calendar.py index 26c4a303dae..d20155ad909 100644 --- a/homeassistant/components/google/calendar.py +++ b/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 diff --git a/homeassistant/components/qnap_qsw/entity.py b/homeassistant/components/qnap_qsw/entity.py index 7d1ec33ee71..288c184984d 100644 --- a/homeassistant/components/qnap_qsw/entity.py +++ b/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] diff --git a/homeassistant/components/roon/server.py b/homeassistant/components/roon/server.py index d30843a7c7b..32d909ff00f 100644 --- a/homeassistant/components/roon/server.py +++ b/homeassistant/components/roon/server.py @@ -66,8 +66,8 @@ class RoonServer: ) # Initialize Roon background polling - self.config_entry.async_on_unload( - asyncio.create_task(self.async_do_loop()).cancel + self.config_entry.async_create_background_task( + self.hass, self.async_do_loop(), "roon.server-do-loop" ) return True diff --git a/homeassistant/components/sonos/__init__.py b/homeassistant/components/sonos/__init__.py index 9d62f1d63ba..57aed48c466 100644 --- a/homeassistant/components/sonos/__init__.py +++ b/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, diff --git a/homeassistant/components/zha/core/gateway.py b/homeassistant/components/zha/core/gateway.py index af2fec6c6a8..128e3b145f1 100644 --- a/homeassistant/components/zha/core/gateway.py +++ b/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( - fetch_updated_state(), "zha.gateway-fetch_updated_state" + self.config_entry.async_create_background_task( + self._hass, fetch_updated_state(), "zha.gateway-fetch_updated_state" ) def device_joined(self, device: zigpy.device.Device) -> None: diff --git a/homeassistant/components/zha/light.py b/homeassistant/components/zha/light.py index 832104892fb..e7bc059054a 100644 --- a/homeassistant/components/zha/light.py +++ b/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", ) diff --git a/homeassistant/config_entries.py b/homeassistant/config_entries.py index 9ffd5afa478..df4cb651528 100644 --- a/homeassistant/config_entries.py +++ b/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: - pending = [task for task in self._pending_tasks if not task.done()] - self._pending_tasks.clear() - if pending: - await asyncio.gather(*pending) + if not self._tasks and not self._background_tasks: + return + + for task in self._background_tasks: + 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 diff --git a/homeassistant/core.py b/homeassistant/core.py index 99b7218a9ff..ef9436dd5bc 100644 --- a/homeassistant/core.py +++ b/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 - your task belongs to is unloaded. + integration, use the create task methods on the config entry instead. This method must be run in the event loop. """ diff --git a/tests/test_config_entries.py b/tests/test_config_entries.py index f63ad9020d3..be508543bd8 100644 --- a/tests/test_config_entries.py +++ b/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"]