diff --git a/tests/test_gateway.py b/tests/test_gateway.py index 17792ba40..4fa14b860 100644 --- a/tests/test_gateway.py +++ b/tests/test_gateway.py @@ -41,6 +41,7 @@ from zha.application.helpers import ZHAData from zha.application.platforms import GroupEntity from zha.application.platforms.light.const import EFFECT_OFF, LightEntityFeature +from zha.async_ import AsyncUtilMixin from zha.quirks import DeviceMatch, DeviceRegistry, ModelInfo, QuirkRegistryEntry from zha.zigbee.device import Device, DeviceEntityAddedEvent, DeviceEntityRemovedEvent from zha.zigbee.group import Group, GroupMemberReference @@ -1252,6 +1253,440 @@ async def test_gateway_shutdown_group_on_remove_failure( assert "Group removal failed" in caplog.text +async def test_gateway_shutdown_controller_failure_cleans_up_and_retries( + zha_gateway: Gateway, +) -> None: + """Test local cleanup completes before a failed controller is retried.""" + application_controller = zha_gateway.application_controller + original_controller_shutdown = application_controller.shutdown + controller_shutdown_error = RuntimeError("Controller shutdown failed") + controller_shutdown_attempts = 0 + background_task_cancelled = asyncio.Event() + + async def shutdown_controller() -> None: + nonlocal controller_shutdown_attempts + controller_shutdown_attempts += 1 + if controller_shutdown_attempts == 1: + raise controller_shutdown_error + await original_controller_shutdown() + + async def wait_until_cancelled() -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + background_task_cancelled.set() + raise + + background_task = zha_gateway.async_create_background_task( + wait_until_cancelled(), + name="test pending task", + untracked=True, + ) + await asyncio.sleep(0) + + zha_device = next(iter(zha_gateway.devices.values())) + zha_group = next(iter(zha_gateway.groups.values())) + with ( + patch.object( + application_controller, + "shutdown", + side_effect=shutdown_controller, + ) as controller_shutdown, + patch.object( + application_controller, + "remove_listener", + wraps=application_controller.remove_listener, + ) as remove_controller_listener, + patch.object( + application_controller.groups, + "remove_listener", + wraps=application_controller.groups.remove_listener, + ) as remove_group_listener, + patch.object( + zha_gateway.global_updater, + "stop", + wraps=zha_gateway.global_updater.stop, + ) as stop_global_updater, + patch.object( + zha_gateway._device_availability_checker, + "stop", + wraps=zha_gateway._device_availability_checker.stop, + ) as stop_availability_checker, + patch.object( + zha_device, "on_remove", wraps=zha_device.on_remove + ) as remove_device, + patch.object(zha_group, "on_remove", wraps=zha_group.on_remove) as remove_group, + ): + with pytest.raises(RuntimeError) as exc_info: + await zha_gateway.shutdown() + + assert exc_info.value is controller_shutdown_error + assert background_task_cancelled.is_set() + assert background_task.cancelled() + assert not zha_gateway._background_tasks + assert not zha_gateway._tracked_completable_tasks + assert not zha_gateway._device_init_tasks + assert not zha_gateway._untracked_background_tasks + assert not zha_gateway.devices + assert not zha_gateway.groups + assert zha_gateway.application_controller is application_controller + assert zha_gateway.shutting_down + + await zha_gateway.shutdown() + + assert zha_gateway.application_controller is None + assert controller_shutdown.await_count == 2 + remove_controller_listener.assert_called_once_with(zha_gateway) + remove_group_listener.assert_called_once_with(zha_gateway) + stop_global_updater.assert_called_once_with() + stop_availability_checker.assert_called_once_with() + remove_device.assert_awaited_once_with() + remove_group.assert_awaited_once_with() + + +async def test_gateway_shutdown_retries_interrupted_device_cleanup( + zha_gateway: Gateway, +) -> None: + """Test interrupted wrapper cleanup retains the wrapper for a retry.""" + application_controller = zha_gateway.application_controller + zha_device = next(iter(zha_gateway.devices.values())) + original_device_on_remove = zha_device.on_remove + original_base_shutdown = AsyncUtilMixin.shutdown + device_cleanup_attempts = 0 + base_shutdown_attempts = 0 + + async def remove_device() -> None: + nonlocal device_cleanup_attempts + device_cleanup_attempts += 1 + if device_cleanup_attempts == 1: + raise asyncio.CancelledError + await original_device_on_remove() + + async def shutdown_base(instance: AsyncUtilMixin) -> None: + nonlocal base_shutdown_attempts + base_shutdown_attempts += 1 + await original_base_shutdown(instance) + + with ( + patch.object( + application_controller, + "shutdown", + wraps=application_controller.shutdown, + ) as controller_shutdown, + patch.object(zha_device, "on_remove", side_effect=remove_device) as on_remove, + patch.object( + AsyncUtilMixin, + "shutdown", + autospec=True, + side_effect=shutdown_base, + ), + ): + with pytest.raises(asyncio.CancelledError): + await zha_gateway.shutdown() + + assert not zha_gateway.devices + assert not zha_gateway.groups + assert zha_gateway.application_controller is application_controller + + await zha_gateway.shutdown() + + assert zha_gateway.application_controller is None + assert on_remove.await_count == 2 + controller_shutdown.assert_awaited_once_with() + assert base_shutdown_attempts == 1 + + +async def test_gateway_initialize_cannot_replace_controller_during_failed_shutdown( + zha_gateway: Gateway, +) -> None: + """Test initialization preserves a controller whose shutdown is incomplete.""" + application_controller = zha_gateway.application_controller + original_controller_shutdown = application_controller.shutdown + controller_shutdown_error = RuntimeError("Controller shutdown failed") + controller_shutdown_attempts = 0 + + async def shutdown_controller() -> None: + nonlocal controller_shutdown_attempts + controller_shutdown_attempts += 1 + if controller_shutdown_attempts == 1: + raise controller_shutdown_error + await original_controller_shutdown() + + with ( + patch.object( + application_controller, + "shutdown", + side_effect=shutdown_controller, + ) as controller_shutdown, + patch.object( + zha_gateway, + "get_application_controller_data", + side_effect=AssertionError("Controller factory must not be queried"), + ) as get_application_controller_data, + ): + with pytest.raises(RuntimeError) as shutdown_exc_info: + await zha_gateway.shutdown() + assert shutdown_exc_info.value is controller_shutdown_error + + with pytest.raises( + RuntimeError, + match="Cannot initialize ZHA while shutdown work remains incomplete", + ): + await zha_gateway.async_initialize() + + get_application_controller_data.assert_not_called() + assert controller_shutdown.await_count == 2 + assert zha_gateway.application_controller is None + + +async def test_gateway_shutdown_cancellation_during_controller_delay_cleans_up( + zha_gateway: Gateway, +) -> None: + """Test cancellation during the controller delay cannot skip local cleanup.""" + application_controller = zha_gateway.application_controller + original_sleep = asyncio.sleep + controlled_shutdown_delay = 1234.5 + controller_delay_started = asyncio.Event() + hold_controller_delay = asyncio.Event() + background_task_cancelled = asyncio.Event() + + async def controlled_sleep(delay: float) -> None: + if delay == controlled_shutdown_delay: + controller_delay_started.set() + await hold_controller_delay.wait() + return + await original_sleep(delay) + + async def wait_until_cancelled() -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + background_task_cancelled.set() + raise + + background_task = zha_gateway.async_create_background_task( + wait_until_cancelled(), + name="test pending task", + untracked=True, + ) + await asyncio.sleep(0) + + shutdown_task: asyncio.Task[None] | None = None + with ( + patch.object( + application_controller, + "shutdown", + wraps=application_controller.shutdown, + ) as controller_shutdown, + patch("zha.application.gateway.SHUT_DOWN_DELAY_S", controlled_shutdown_delay), + patch("zha.application.gateway.asyncio.sleep", side_effect=controlled_sleep), + ): + try: + shutdown_task = asyncio.create_task(zha_gateway.shutdown()) + async with asyncio.timeout(5): + await controller_delay_started.wait() + + shutdown_task.cancel() + with pytest.raises(asyncio.CancelledError): + await shutdown_task + + assert background_task_cancelled.is_set() + assert background_task.cancelled() + assert not zha_gateway._background_tasks + assert not zha_gateway._tracked_completable_tasks + assert not zha_gateway._device_init_tasks + assert not zha_gateway._untracked_background_tasks + assert not zha_gateway.devices + assert not zha_gateway.groups + assert zha_gateway.application_controller is None + + await zha_gateway.shutdown() + controller_shutdown.assert_awaited_once_with() + finally: + hold_controller_delay.set() + if shutdown_task is not None and not shutdown_task.done(): + shutdown_task.cancel() + if shutdown_task is not None: + await asyncio.gather(shutdown_task, return_exceptions=True) + if not background_task.done(): + background_task.cancel() + await asyncio.gather(background_task, return_exceptions=True) + + +async def test_gateway_shutdown_base_cleanup_failure_clears_state_and_retries( + zha_gateway: Gateway, +) -> None: + """Test a late base cleanup failure leaves a deterministic retry path.""" + application_controller = zha_gateway.application_controller + base_cleanup_error = RuntimeError("Base cleanup failed") + original_base_shutdown = AsyncUtilMixin.shutdown + base_shutdown_attempts = 0 + background_task_cancelled = asyncio.Event() + + async def shutdown_base(instance: AsyncUtilMixin) -> None: + nonlocal base_shutdown_attempts + base_shutdown_attempts += 1 + if base_shutdown_attempts == 1: + raise base_cleanup_error + await original_base_shutdown(instance) + + async def wait_until_cancelled() -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + background_task_cancelled.set() + raise + + background_task = zha_gateway.async_create_background_task( + wait_until_cancelled(), + name="test pending task", + untracked=True, + ) + await asyncio.sleep(0) + + with ( + patch.object( + application_controller, + "shutdown", + wraps=application_controller.shutdown, + ) as controller_shutdown, + patch.object( + AsyncUtilMixin, + "shutdown", + autospec=True, + side_effect=shutdown_base, + ) as base_shutdown, + ): + with pytest.raises(RuntimeError) as exc_info: + await zha_gateway.shutdown() + + assert exc_info.value is base_cleanup_error + assert not zha_gateway.devices + assert not zha_gateway.groups + assert zha_gateway.application_controller is None + assert not background_task.done() + assert background_task in zha_gateway._untracked_background_tasks + + await zha_gateway.shutdown() + + controller_shutdown.assert_awaited_once_with() + assert base_shutdown.await_count == 2 + assert background_task_cancelled.is_set() + assert background_task.cancelled() + assert not zha_gateway._background_tasks + assert not zha_gateway._tracked_completable_tasks + assert not zha_gateway._device_init_tasks + assert not zha_gateway._untracked_background_tasks + + +async def test_gateway_shutdown_preserves_controller_error_when_cleanup_fails( + zha_gateway: Gateway, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test a local cleanup error cannot replace the controller failure.""" + application_controller = zha_gateway.application_controller + original_controller_shutdown = application_controller.shutdown + controller_shutdown_error = RuntimeError("Controller shutdown failed") + base_cleanup_error = RuntimeError("Base cleanup failed") + controller_shutdown_attempts = 0 + original_base_shutdown = AsyncUtilMixin.shutdown + base_shutdown_attempts = 0 + + async def shutdown_controller() -> None: + nonlocal controller_shutdown_attempts + controller_shutdown_attempts += 1 + if controller_shutdown_attempts == 1: + raise controller_shutdown_error + await original_controller_shutdown() + + async def shutdown_base(instance: AsyncUtilMixin) -> None: + nonlocal base_shutdown_attempts + base_shutdown_attempts += 1 + await original_base_shutdown(instance) + if base_shutdown_attempts == 1: + raise base_cleanup_error + + with ( + patch.object( + application_controller, + "shutdown", + side_effect=shutdown_controller, + ) as controller_shutdown, + patch.object( + AsyncUtilMixin, + "shutdown", + autospec=True, + side_effect=shutdown_base, + ) as base_shutdown, + ): + with pytest.raises(RuntimeError) as exc_info: + await zha_gateway.shutdown() + + assert exc_info.value is controller_shutdown_error + assert "Local cleanup failed after an earlier shutdown error" in caplog.text + assert "Base cleanup failed" in caplog.text + assert not zha_gateway.devices + assert not zha_gateway.groups + assert zha_gateway.application_controller is application_controller + + await zha_gateway.shutdown() + + assert zha_gateway.application_controller is None + assert controller_shutdown.await_count == 2 + assert base_shutdown.await_count == 2 + + +async def test_gateway_concurrent_shutdown_waits_for_active_attempt( + zha_gateway: Gateway, +) -> None: + """Test concurrent shutdown callers share one serialized attempt.""" + application_controller = zha_gateway.application_controller + original_controller_shutdown = application_controller.shutdown + controller_shutdown_started = asyncio.Event() + finish_controller_shutdown = asyncio.Event() + + async def controlled_controller_shutdown() -> None: + controller_shutdown_started.set() + await finish_controller_shutdown.wait() + await original_controller_shutdown() + + first_shutdown_task: asyncio.Task[None] | None = None + second_shutdown_task: asyncio.Task[None] | None = None + with patch.object( + application_controller, + "shutdown", + side_effect=controlled_controller_shutdown, + ) as controller_shutdown: + try: + first_shutdown_task = asyncio.create_task(zha_gateway.shutdown()) + async with asyncio.timeout(5): + await controller_shutdown_started.wait() + + second_shutdown_task = asyncio.create_task(zha_gateway.shutdown()) + await asyncio.sleep(0) + assert not second_shutdown_task.done() + controller_shutdown.assert_awaited_once_with() + + finish_controller_shutdown.set() + await asyncio.gather(first_shutdown_task, second_shutdown_task) + finally: + finish_controller_shutdown.set() + pending_shutdown_tasks = [ + task + for task in (first_shutdown_task, second_shutdown_task) + if task is not None + ] + for task in pending_shutdown_tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*pending_shutdown_tasks, return_exceptions=True) + + controller_shutdown.assert_awaited_once_with() + assert zha_gateway.application_controller is None + assert not zha_gateway.devices + assert not zha_gateway.groups + + async def test_group_on_remove_entity_failure( zha_gateway: Gateway, caplog: pytest.LogCaptureFixture, diff --git a/zha/application/gateway.py b/zha/application/gateway.py index ecff916d1..62dad700c 100644 --- a/zha/application/gateway.py +++ b/zha/application/gateway.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +from collections import deque from collections.abc import AsyncGenerator from contextlib import asynccontextmanager, suppress from dataclasses import dataclass @@ -182,6 +183,13 @@ def __init__(self, config: ZHAData) -> None: self.coordinator_zha_device: Device | None = None self.shutting_down: bool = False + # Controller shutdown and local cleanup can fail independently. Track the + # phases separately so a later call retries only incomplete work. + self._shutdown_lock: asyncio.Lock = asyncio.Lock() + self._local_shutdown_preparation_complete: bool = False + self._local_shutdown_finalization_complete: bool = False + self._devices_pending_shutdown: deque[Device] | None = None + self._groups_pending_shutdown: deque[Group] | None = None self._reload_task: asyncio.Task | None = None self.global_updater: GlobalUpdater = GlobalUpdater(self) @@ -195,6 +203,15 @@ def radio_type(self) -> RadioType: """Get the current radio type.""" return RadioType[self.config.config.coordinator_configuration.radio_type] + @property + def _shutdown_complete(self) -> bool: + """Return whether every shutdown phase reached its terminal state.""" + return ( + self._local_shutdown_preparation_complete + and self._local_shutdown_finalization_complete + and self.application_controller is None + ) + def get_application_controller_data(self) -> tuple[ControllerApplication, dict]: """Get an uninitialized instance of a zigpy `ControllerApplication`.""" app_config = self.config.zigpy_config @@ -242,32 +259,42 @@ async def async_from_config(cls, config: ZHAData) -> Self: async def _async_initialize(self) -> None: """Initialize controller and connect radio.""" - self.shutting_down = False - - app_controller_cls, app_config = self.get_application_controller_data() - self.application_controller = await app_controller_cls.new( - config=app_config, - auto_form=False, - start_radio=False, - device_resolver=DEVICE_REGISTRY.resolve, - uninitialized_packet_handler=( - self.config.config.quirks_configuration.uninitialized_packet_handler - ), - ) + async with self._shutdown_lock: + if self.shutting_down and not self._shutdown_complete: + raise RuntimeError( + "Cannot initialize ZHA while shutdown work remains incomplete" + ) - await self.application_controller.startup(auto_form=True) + self.shutting_down = False + self._local_shutdown_preparation_complete = False + self._local_shutdown_finalization_complete = False + self._devices_pending_shutdown = None + self._groups_pending_shutdown = None + + app_controller_cls, app_config = self.get_application_controller_data() + self.application_controller = await app_controller_cls.new( + config=app_config, + auto_form=False, + start_radio=False, + device_resolver=DEVICE_REGISTRY.resolve, + uninitialized_packet_handler=( + self.config.config.quirks_configuration.uninitialized_packet_handler + ), + ) - self.coordinator_zha_device = self.get_or_create_device( - self._find_coordinator_device() - ) + await self.application_controller.startup(auto_form=True) - await self.load_devices() - self.load_groups() + self.coordinator_zha_device = self.get_or_create_device( + self._find_coordinator_device() + ) - self.application_controller.add_listener(self) - self.application_controller.groups.add_listener(self) - self.global_updater.start() - self._device_availability_checker.start() + await self.load_devices() + self.load_groups() + + self.application_controller.add_listener(self) + self.application_controller.groups.add_listener(self) + self.global_updater.start() + self._device_availability_checker.start() async def async_initialize(self) -> None: """Initialize controller and connect radio.""" @@ -883,18 +910,24 @@ async def async_remove_zigpy_group(self, group_id: int) -> None: await asyncio.gather(*tasks) self.application_controller.groups.pop(group_id) - async def shutdown(self) -> None: - """Stop ZHA Controller Application.""" - if self.shutting_down: - _LOGGER.debug("Ignoring duplicate shutdown event") - return - - self.shutting_down = True + async def _async_prepare_local_shutdown(self) -> None: + """Stop local producers and tear down device and group state.""" + if self._devices_pending_shutdown is None: + self._devices_pending_shutdown = deque(self._devices.values()) + if self._groups_pending_shutdown is None: + self._groups_pending_shutdown = deque(self._groups.values()) self.global_updater.stop() self._device_availability_checker.stop() - for device in self._devices.values(): + if self.application_controller is not None: + # A controller retained after a failed shutdown must not emit callbacks + # into local state that has already been cleared. + self.application_controller.remove_listener(self) + self.application_controller.groups.remove_listener(self) + + while self._devices_pending_shutdown: + device = self._devices_pending_shutdown[0] try: await device.on_remove() except Exception: @@ -903,8 +936,10 @@ async def shutdown(self) -> None: device, exc_info=True, ) + self._devices_pending_shutdown.popleft() - for group in self._groups.values(): + while self._groups_pending_shutdown: + group = self._groups_pending_shutdown[0] try: await group.on_remove() except Exception: @@ -913,18 +948,62 @@ async def shutdown(self) -> None: group, exc_info=True, ) + self._groups_pending_shutdown.popleft() - _LOGGER.debug("Shutting down ZHA ControllerApplication") - if self.application_controller is not None: - await self.application_controller.shutdown() - self.application_controller = None - # give bellows thread callback a chance to run - await asyncio.sleep(SHUT_DOWN_DELAY_S) + self._devices_pending_shutdown = None + self._groups_pending_shutdown = None - await super().shutdown() + async def _async_finalize_local_shutdown(self) -> None: + """Cancel owned work and clear local device and group state.""" + try: + await super().shutdown() + finally: + self._devices.clear() + self._groups.clear() + + async def shutdown(self) -> None: + """Stop ZHA Controller Application.""" + async with self._shutdown_lock: + if self._shutdown_complete: + _LOGGER.debug("Ignoring duplicate shutdown event") + return + + self.shutting_down = True + # Defer cancellation and controller failures until local finalization; + # a secondary cleanup failure must not hide the original cause. + primary_shutdown_error: BaseException | None = None + local_cleanup_error: BaseException | None = None - self._devices.clear() - self._groups.clear() + try: + if not self._local_shutdown_preparation_complete: + await self._async_prepare_local_shutdown() + self._local_shutdown_preparation_complete = True + + if self.application_controller is not None: + _LOGGER.debug("Shutting down ZHA ControllerApplication") + await self.application_controller.shutdown() + self.application_controller = None + # Give the bellows thread callback a chance to run. + await asyncio.sleep(SHUT_DOWN_DELAY_S) + except BaseException as shutdown_exception: + primary_shutdown_error = shutdown_exception + finally: + if not self._local_shutdown_finalization_complete: + try: + await self._async_finalize_local_shutdown() + except BaseException as local_cleanup_exception: + local_cleanup_error = local_cleanup_exception + if primary_shutdown_error is not None: + _LOGGER.exception( + "Local cleanup failed after an earlier shutdown error" + ) + else: + self._local_shutdown_finalization_complete = True + + if primary_shutdown_error is not None: + raise primary_shutdown_error + if local_cleanup_error is not None: + raise local_cleanup_error def handle_message( # pylint: disable=unused-argument self,