diff --git a/custom_components/heating_mpc/__init__.py b/custom_components/heating_mpc/__init__.py index 0eb32f4..39b01c2 100644 --- a/custom_components/heating_mpc/__init__.py +++ b/custom_components/heating_mpc/__init__.py @@ -12,7 +12,7 @@ from homeassistant.helpers.event import async_track_state_change_event from .const import CONF_HEATER, CONF_ROOMS, CONF_WEIGHTS, DOMAIN, CONF_OUTSIDE, WEATHER_SOURCES from .weather_sources import sources_for -from .coordinator import HeatingMPCCoordinator, storage_lock +from .coordinator import HeatingMPCCoordinator, mark_model_removed, storage_lock from .zones import SUBENTRY_ZONE, ZONE_KEY, ZONE_NAME, legacy_zone_data PLATFORMS: list[Platform] = [Platform.SENSOR, Platform.BINARY_SENSOR] @@ -29,9 +29,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) except BaseException: - # HA skips unload callbacks for unexpected setup errors, so clean up here. + # HA skips unload callbacks for unexpected setup errors, so clean up here, + # including a refresh scheduled by entities that were already added. unsub_relay() hass.data.get(DOMAIN, {}).pop(entry.entry_id, None) + await coordinator.async_shutdown() raise entry.async_on_unload(unsub_relay) # Remove only obsolete zone forecasts belonging to this entry/platform. Zones @@ -53,17 +55,25 @@ async def async_reload_entry(hass: HomeAssistant, entry: ConfigEntry) -> None: async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: + coordinator = hass.data.get(DOMAIN, {}).get(entry.entry_id) + if coordinator is not None: + # Before anything is unloaded: an update still running must not publish or save. + coordinator.async_begin_shutdown() ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) if ok: hass.data[DOMAIN].pop(entry.entry_id, None) + elif coordinator is not None: + coordinator.async_cancel_shutdown() return ok async def async_remove_entry(hass: HomeAssistant, entry: ConfigEntry) -> None: """Remove the learned model when this integration instance is deleted.""" store = Store(hass, 1, f"{DOMAIN}.{entry.entry_id}") - # Wait for a save still in progress, which would otherwise recreate the file. + # Wait for a save still in progress, which would otherwise recreate the file; + # a save still waiting for the lock sees the mark and does not write. async with storage_lock(hass, entry.entry_id): + mark_model_removed(hass, entry.entry_id) await store.async_remove() hass.data.get(f"{DOMAIN}_storage_locks", {}).pop(entry.entry_id, None) diff --git a/custom_components/heating_mpc/coordinator.py b/custom_components/heating_mpc/coordinator.py index 1d510fa..13ffa47 100644 --- a/custom_components/heating_mpc/coordinator.py +++ b/custom_components/heating_mpc/coordinator.py @@ -56,6 +56,15 @@ def storage_lock(hass: HomeAssistant, entry_id: str) -> asyncio.Lock: return hass.data.setdefault(f"{DOMAIN}_storage_locks", {}).setdefault(entry_id, asyncio.Lock()) +def mark_model_removed(hass: HomeAssistant, entry_id: str) -> None: + """Record, under the storage lock, that the entry's model file was deleted for good.""" + hass.data.setdefault(f"{DOMAIN}_removed", set()).add(entry_id) + + +def model_removed(hass: HomeAssistant, entry_id: str) -> bool: + return entry_id in hass.data.get(f"{DOMAIN}_removed", ()) + + def entry_zone_list(entry: ConfigEntry, config: dict[str, Any]) -> list[Zone]: """Zones from the entry's zone subentries; one zone per legacy room sensor otherwise.""" subentries = getattr(entry, "subentries", None) or {} @@ -148,18 +157,57 @@ class HeatingMPCCoordinator(DataUpdateCoordinator[dict[str, Any]]): self._off_samples = 0 # Increases with every update; an older update finishing late never publishes. self._run = 0 + self._closing = False async def async_restore_model(self) -> None: async with storage_lock(self.hass, self.entry.entry_id): await self._async_restore_model() - async def _async_save(self, outside_source: str | None) -> None: + async def _async_save(self, outside_source: str | None) -> bool: + """Write the model unless unloading began or the model was deleted. + + Shielded: unloading cancels a running scheduled update, but the file write + itself continues in a worker thread. The lock must stay held until that + write has finished, or a deletion could run first and be undone by it. + """ + write = self.hass.async_create_task(self._async_write(outside_source), f"{DOMAIN} model save") + return await asyncio.shield(write) + + async def _async_write(self, outside_source: str | None) -> bool: async with storage_lock(self.hass, self.entry.entry_id): + # Checked after acquiring the lock, not before: a save that waited for + # a deletion to finish must not write the file again. + if self.closing or model_removed(self.hass, self.entry.entry_id): + return False await self._store.async_save(self._payload(outside_source)) + return True + + @property + def closing(self) -> bool: + """True from the start of unloading: nothing is published, saved or recorded.""" + return self._closing or getattr(self, "_shutdown_requested", False) + + @callback + def async_begin_shutdown(self) -> None: + """Mark unloading before anything is unloaded. + + Home Assistant's own shutdown of the coordinator runs only after platforms + are unloaded; an update still running in that window must not publish or save. + """ + self._closing = True + + @callback + def async_cancel_shutdown(self) -> None: + """Unloading failed and the entry stays loaded: carry on.""" + self._closing = False + + async def async_shutdown(self) -> None: + self._closing = True + await super().async_shutdown() def _stale(self, run: int) -> bool: """True when a newer update started or the entry is unloading.""" - return run != self._run or getattr(self, "_shutdown_requested", False) + return run != self._run or self.closing async def _async_restore_model(self) -> None: stored = await self._store.async_load() or {} @@ -303,6 +351,8 @@ class HeatingMPCCoordinator(DataUpdateCoordinator[dict[str, Any]]): @callback def async_relay_changed(self, event: Event) -> None: """Record relay state changes as they happen; read-only state listener.""" + if self.closing: + return # a callback queued before the listener was removed old, new = event.data.get("old_state"), event.data.get("new_state") if new is None: self.record_relay_state(event.time_fired, "removed") @@ -527,6 +577,8 @@ class HeatingMPCCoordinator(DataUpdateCoordinator[dict[str, Any]]): snapshot = {entity: self.hass.states.get(entity) for entity in set(self.sensors) | weather_entity_ids} self._run += 1 + if self.closing: + return self._superseded({"status": "inputs_changed"}) run = self._run inputs_before = self._input_snapshot() sensor_quality = {} diff --git a/tests/ha/test_concurrency.py b/tests/ha/test_concurrency.py new file mode 100644 index 0000000..572c2c3 --- /dev/null +++ b/tests/ha/test_concurrency.py @@ -0,0 +1,134 @@ +"""Unload, reload and removal while work is still in flight, in the real HA lifecycle.""" +import asyncio +import threading +from datetime import timedelta +from unittest.mock import patch + +from homeassistant.helpers import entity_registry as er +from homeassistant.helpers.storage import Store +from homeassistant.util import dt as dt_util +from pytest_homeassistant_custom_component.common import async_fire_time_changed + +from custom_components.heating_mpc import coordinator as coordinator_module +from custom_components.heating_mpc.const import DOMAIN + +REAL_OPTIMIZE = coordinator_module.optimize + + +async def setup(hass, entry): + entry.add_to_hass(hass) + assert await hass.config_entries.async_setup(entry.entry_id) + await hass.async_block_till_done() + return hass.data[DOMAIN][entry.entry_id] + + +def trained(coordinator): + model = coordinator.learners["sensor.room"] + model.count, model.on_count, model.off_count = 250, 125, 125 + coordinator._update_counts() + + +class BlockedOptimizer: + """Holds the optimizer in its worker thread until released.""" + + def __init__(self): + self.started, self.release = threading.Event(), threading.Event() + + def __call__(self, *args, **kwargs): + self.started.set() + self.release.wait(5) + return REAL_OPTIMIZE(*args, **kwargs) + + async def wait_started(self): + # Polled on the event loop: the test executor may run one job at a time. + for _ in range(500): + if self.started.is_set(): + return True + await asyncio.sleep(0.01) + return False + + +def status_id(hass, entry): + return er.async_get(hass).async_get_entity_id("sensor", DOMAIN, f"{entry.entry_id}_status") + + +async def test_unload_during_optimization_publishes_nothing(hass, entry, inputs): + coordinator = await setup(hass, entry) + trained(coordinator) + shown = coordinator.data + blocked = BlockedOptimizer() + with patch.object(coordinator_module, "optimize", blocked): + refresh = hass.async_create_task(coordinator.async_refresh()) + assert await blocked.wait_started() + assert await hass.config_entries.async_unload(entry.entry_id) + state_after_unload = hass.states.get(status_id(hass, entry)) + blocked.release.set() + await refresh + await hass.async_block_till_done() + # The plan finished after unloading began: kept out of the data and the states. + assert coordinator.data is shown + assert coordinator.data["status"] == "learning" + assert hass.states.get(status_id(hass, entry)) == state_after_unload + + +async def test_reload_during_optimization_shows_only_the_new_coordinator(hass, entry, inputs): + old = await setup(hass, entry) + trained(old) # only in memory: the reloaded coordinator is not ready + blocked = BlockedOptimizer() + with patch.object(coordinator_module, "optimize", blocked): + refresh = hass.async_create_task(old.async_refresh()) + assert await blocked.wait_started() + assert await hass.config_entries.async_reload(entry.entry_id) + await hass.async_block_till_done() + new = hass.data[DOMAIN][entry.entry_id] + blocked.release.set() + await refresh + await hass.async_block_till_done() + assert new is not old and old.closing + assert old.data["status"] != "ready" + assert hass.states.get(status_id(hass, entry)).state == new.data["status"] == "learning" + + +async def test_removal_during_a_save_leaves_no_model_behind(hass, entry, inputs, freezer): + coordinator = await setup(hass, entry) + started, release = asyncio.Event(), asyncio.Event() + write = coordinator._store.async_save + + async def slow_save(data): + started.set() + await release.wait() + await write(data) + coordinator._store.async_save = slow_save + future = dt_util.utcnow() + timedelta(minutes=5, seconds=1) + freezer.move_to(future) + async_fire_time_changed(hass, future) # a learning poll, which saves + await asyncio.wait_for(started.wait(), 5) + removal = hass.async_create_task(hass.config_entries.async_remove(entry.entry_id)) + for _ in range(5): + await asyncio.sleep(0) + assert not removal.done() # deletion waits for the save in progress + release.set() + assert await removal + await hass.async_block_till_done() + assert await Store(hass, 1, f"{DOMAIN}.{entry.entry_id}").async_load() is None + + +async def test_repeated_setup_and_unload_leave_nothing_running(hass, entry, inputs, freezer): + entry.add_to_hass(hass) + coordinators = [] + for _ in range(3): + assert await hass.config_entries.async_setup(entry.entry_id) + await hass.async_block_till_done() + coordinators.append(hass.data[DOMAIN][entry.entry_id]) + assert await hass.config_entries.async_unload(entry.entry_id) + await hass.async_block_till_done() + assert entry.entry_id not in hass.data[DOMAIN] + assert all(c.closing and c._unsub_refresh is None and not c._listeners for c in coordinators) + counts = [c.learners["sensor.room"].count for c in coordinators] + hass.states.async_set("switch.boiler", "on") + future = dt_util.utcnow() + timedelta(minutes=10) + freezer.move_to(future) + async_fire_time_changed(hass, future) + await hass.async_block_till_done() + assert not any(c._relay_events for c in coordinators) + assert [c.learners["sensor.room"].count for c in coordinators] == counts diff --git a/tests/test_coordinator.py b/tests/test_coordinator.py index 799cc99..b14689b 100644 --- a/tests/test_coordinator.py +++ b/tests/test_coordinator.py @@ -24,6 +24,9 @@ class CoordinatorBase: async def async_config_entry_first_refresh(self): self.data = await self._async_update_data() + async def async_shutdown(self): + self._shutdown_requested = True + class Store: def __init__(self, *args): @@ -72,6 +75,7 @@ class BoundaryTests(unittest.IsolatedAsyncioTestCase): config=types.SimpleNamespace(units=types.SimpleNamespace(temperature_unit='°C')), services=types.SimpleNamespace(async_call=Mock(side_effect=AssertionError('No service calls permitted'))), track_state=Mock(return_value=Mock(name='unsub_relay')), + async_create_task=lambda coro, name=None, eager_start=True: asyncio.ensure_future(coro), registry=types.SimpleNamespace(entries=[], async_remove=Mock()), config_entries=types.SimpleNamespace( async_forward_entry_setups=AsyncMock(), async_unload_platforms=AsyncMock(return_value=True), @@ -1308,6 +1312,82 @@ class BoundaryTests(unittest.IsolatedAsyncioTestCase): self.assertIs(await self.sample(), previous) worker.assert_not_called() + async def test_unload_marks_shutdown_before_platforms_unload(self): + await integration.async_setup_entry(self.hass, self.entry) + c = self.hass.data['heating_mpc']['entry'] + seen = [] + async def unload_platforms(entry, platforms): + seen.append(c.closing) + return True + self.hass.config_entries.async_unload_platforms.side_effect = unload_platforms + self.assertTrue(await integration.async_unload_entry(self.hass, self.entry)) + self.assertEqual(seen, [True]) + # A failed unload leaves the entry loaded, so the coordinator carries on. + await integration.async_setup_entry(self.hass, self.entry) + c = self.hass.data['heating_mpc']['entry'] + self.hass.config_entries.async_unload_platforms.side_effect = None + self.hass.config_entries.async_unload_platforms.return_value = False + self.assertFalse(await integration.async_unload_entry(self.hass, self.entry)) + self.assertFalse(c.closing) + + async def test_nothing_is_learned_saved_or_recorded_once_unloading_began(self): + self.ready_office() + await self.sample() + shown = self.c.data = {'status': 'learning', 'proposal': None} + self.fresh('sensor.room', 'sensor.outside', 'sensor.ups_power') + self.c.async_begin_shutdown() + count = self.c.learners['office'].count + self.assertIs(await self.sample(), shown) + self.assertEqual(self.c.learners['office'].count, count) + self.c._store.async_save.assert_not_awaited() + self.assertFalse(await self.c._async_save('sensor.outside')) + # A relay callback queued before the listener was removed is ignored. + new = types.SimpleNamespace(state='on', last_changed=self.now) + self.c.async_relay_changed(types.SimpleNamespace( + data={'old_state': types.SimpleNamespace(state='off'), 'new_state': new}, time_fired=self.now)) + self.assertFalse(self.c._relay_events) + + async def test_failed_setup_shuts_the_coordinator_down(self): + self.hass.config_entries.async_forward_entry_setups.side_effect = RuntimeError('test') + with patch.object(coordinator.HeatingMPCCoordinator, 'async_shutdown', autospec=True) as shutdown: + with self.assertRaises(RuntimeError): + await integration.async_setup_entry(self.hass, self.entry) + shutdown.assert_awaited_once() + + async def test_simultaneous_saves_never_overlap(self): + self.ready_office() + order = [] + async def slow_save(payload): + order.append('start') + await asyncio.sleep(0) + await asyncio.sleep(0) + order.append('end') + self.c._store.async_save.side_effect = slow_save + results = await asyncio.gather(*(self.c._async_save('sensor.outside') for _ in range(3))) + self.assertEqual(results, [True, True, True]) + self.assertEqual(order, ['start', 'end'] * 3) + + async def test_a_save_waiting_for_a_deletion_does_not_recreate_the_model(self): + self.ready_office() + order, release = [], asyncio.Event() + store = Store() + async def slow_remove(): + order.append('delete started') + await release.wait() + order.append('deleted') + store.async_remove.side_effect = slow_remove + with patch.object(integration, 'Store', return_value=store): + removal = asyncio.ensure_future(integration.async_remove_entry(self.hass, self.entry)) + await asyncio.sleep(0) + saving = asyncio.ensure_future(self.c._async_save('sensor.outside')) + await asyncio.sleep(0) + self.assertEqual(order, ['delete started']) # the save is waiting for the lock + release.set() + await removal + self.assertFalse(await saving) + self.assertEqual(order, ['delete started', 'deleted']) + self.c._store.async_save.assert_not_awaited() + async def test_restore_waits_for_a_save_still_in_progress(self): self.ready_office() order = []