RFC - Add a 3rd state to the HA shutdown sequence for writing… (#33358)

* add third stage to hass shutdown

* use 3rd stage in storage

* update core state

* add writing data to multi stop ignore

* update core test

* review comment

* update name based on feedback

homeassistant/const.py

@@ -184,6 +184,7 @@ EVENT_CORE_CONFIG_UPDATE = "core_config_updated"
 EVENT_HOMEASSISTANT_CLOSE = "homeassistant_close"
 EVENT_HOMEASSISTANT_START = "homeassistant_start"
 EVENT_HOMEASSISTANT_STOP = "homeassistant_stop"
+EVENT_HOMEASSISTANT_FINAL_WRITE = "homeassistant_final_write"
 EVENT_LOGBOOK_ENTRY = "logbook_entry"
 EVENT_PLATFORM_DISCOVERED = "platform_discovered"
 EVENT_SCRIPT_STARTED = "script_started"

homeassistant/core.py

@@ -47,6 +47,7 @@ from homeassistant.const import (
     EVENT_CALL_SERVICE,
     EVENT_CORE_CONFIG_UPDATE,
     EVENT_HOMEASSISTANT_CLOSE,
+    EVENT_HOMEASSISTANT_FINAL_WRITE,
     EVENT_HOMEASSISTANT_START,
     EVENT_HOMEASSISTANT_STOP,
     EVENT_SERVICE_REGISTERED,
@@ -151,6 +152,7 @@ class CoreState(enum.Enum):
     starting = "STARTING"
     running = "RUNNING"
     stopping = "STOPPING"
+    writing_data = "WRITING_DATA"
 
     def __str__(self) -> str:
         """Return the event."""
@@ -412,7 +414,7 @@ class HomeAssistant:
             # regardless of the state of the loop.
             if self.state == CoreState.not_running:  # just ignore
                 return
-            if self.state == CoreState.stopping:
+            if self.state == CoreState.stopping or self.state == CoreState.writing_data:
                 _LOGGER.info("async_stop called twice: ignored")
                 return
             if self.state == CoreState.starting:
@@ -426,6 +428,11 @@ class HomeAssistant:
         await self.async_block_till_done()
 
         # stage 2
+        self.state = CoreState.writing_data
+        self.bus.async_fire(EVENT_HOMEASSISTANT_FINAL_WRITE)
+        await self.async_block_till_done()
+
+        # stage 3
         self.state = CoreState.not_running
         self.bus.async_fire(EVENT_HOMEASSISTANT_CLOSE)
         await self.async_block_till_done()

homeassistant/helpers/restore_state.py

@@ -4,7 +4,10 @@ from datetime import datetime, timedelta
 import logging
 from typing import Any, Awaitable, Dict, List, Optional, Set, cast
 
-from homeassistant.const import EVENT_HOMEASSISTANT_START, EVENT_HOMEASSISTANT_STOP
+from homeassistant.const import (
+    EVENT_HOMEASSISTANT_FINAL_WRITE,
+    EVENT_HOMEASSISTANT_START,
+)
 from homeassistant.core import (
     CoreState,
     HomeAssistant,
@@ -184,7 +187,9 @@ class RestoreStateData:
         async_track_time_interval(self.hass, _async_dump_states, STATE_DUMP_INTERVAL)
 
         # Dump states when stopping hass
-        self.hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _async_dump_states)
+        self.hass.bus.async_listen_once(
+            EVENT_HOMEASSISTANT_FINAL_WRITE, _async_dump_states
+        )
 
     @callback
     def async_restore_entity_added(self, entity_id: str) -> None:

homeassistant/helpers/storage.py

@@ -5,7 +5,7 @@ import logging
 import os
 from typing import Any, Callable, Dict, List, Optional, Type, Union
 
-from homeassistant.const import EVENT_HOMEASSISTANT_STOP
+from homeassistant.const import EVENT_HOMEASSISTANT_FINAL_WRITE
 from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
 from homeassistant.helpers.event import async_call_later
 from homeassistant.loader import bind_hass
@@ -153,7 +153,7 @@ class Store:
         """Ensure that we write if we quit before delay has passed."""
         if self._unsub_stop_listener is None:
             self._unsub_stop_listener = self.hass.bus.async_listen_once(
-                EVENT_HOMEASSISTANT_STOP, self._async_callback_stop_write
+                EVENT_HOMEASSISTANT_FINAL_WRITE, self._async_callback_stop_write
             )
 
     @callback

tests/helpers/test_storage.py

@@ -6,7 +6,7 @@ from unittest.mock import Mock, patch
 
 import pytest
 
-from homeassistant.const import EVENT_HOMEASSISTANT_STOP
+from homeassistant.const import EVENT_HOMEASSISTANT_FINAL_WRITE
 from homeassistant.helpers import storage
 from homeassistant.util import dt
 
@@ -85,7 +85,7 @@ async def test_saving_on_stop(hass, hass_storage):
     store.async_delay_save(lambda: MOCK_DATA, 1)
     assert store.key not in hass_storage
 
-    hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP)
+    hass.bus.async_fire(EVENT_HOMEASSISTANT_FINAL_WRITE)
     await hass.async_block_till_done()
     assert hass_storage[store.key] == {
         "version": MOCK_VERSION,

tests/test_core.py

@@ -21,6 +21,7 @@ from homeassistant.const import (
     EVENT_CALL_SERVICE,
     EVENT_CORE_CONFIG_UPDATE,
     EVENT_HOMEASSISTANT_CLOSE,
+    EVENT_HOMEASSISTANT_FINAL_WRITE,
     EVENT_HOMEASSISTANT_STOP,
     EVENT_SERVICE_REGISTERED,
     EVENT_SERVICE_REMOVED,
@@ -151,10 +152,14 @@ def test_stage_shutdown():
     """Simulate a shutdown, test calling stuff."""
     hass = get_test_home_assistant()
     test_stop = []
+    test_final_write = []
     test_close = []
     test_all = []
 
     hass.bus.listen(EVENT_HOMEASSISTANT_STOP, lambda event: test_stop.append(event))
+    hass.bus.listen(
+        EVENT_HOMEASSISTANT_FINAL_WRITE, lambda event: test_final_write.append(event)
+    )
     hass.bus.listen(EVENT_HOMEASSISTANT_CLOSE, lambda event: test_close.append(event))
     hass.bus.listen("*", lambda event: test_all.append(event))
 
@@ -162,7 +167,8 @@ def test_stage_shutdown():
 
     assert len(test_stop) == 1
     assert len(test_close) == 1
-    assert len(test_all) == 1
+    assert len(test_final_write) == 1
+    assert len(test_all) == 2
 
 
 class TestHomeAssistant(unittest.TestCase):