Stop publishing and saving from the start of unload; keep deletions final

Home Assistant marks a coordinator shut down only after the platforms are
unloaded, and the model save checked for shutdown before waiting for the
storage lock. A deletion could therefore be followed by a save that had
been waiting for it, recreating the file. Also, unloading cancels a running
scheduled update, which released the lock while the file write continued
in a worker thread.

- The coordinator is marked closing before platforms unload (restored if
  unloading fails); closing updates, relay callbacks and saves do nothing.
- The save checks closing and a removed-model mark after acquiring the
  lock, and runs shielded so the lock is held until the write finishes.
- Removal marks the model removed under the lock before deleting it.
- A failed setup also shuts the coordinator down.

Tests: unload and reload during a running optimization, removal during a
save, a save waiting for a deletion, simultaneous saves, repeated
setup/unload, shutdown ordering and failed setup.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

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)
 

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 = {}

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

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 = []