RFC - Add a 3rd state to the HA shutdown sequence for writing data

homeassistant/const.py

@@ -184,6 +184,7 @@
 EVENT_HOMEASSISTANT_CLOSE = "homeassistant_close"
 EVENT_HOMEASSISTANT_START = "homeassistant_start"
 EVENT_HOMEASSISTANT_STOP = "homeassistant_stop"
+EVENT_HOMEASSISTANT_WRITE_DATA = "homeassistant_write_data"
 EVENT_LOGBOOK_ENTRY = "logbook_entry"
 EVENT_PLATFORM_DISCOVERED = "platform_discovered"
 EVENT_SCRIPT_STARTED = "script_started"

homeassistant/core.py

@@ -49,6 +49,7 @@
     EVENT_HOMEASSISTANT_CLOSE,
     EVENT_HOMEASSISTANT_START,
     EVENT_HOMEASSISTANT_STOP,
+    EVENT_HOMEASSISTANT_WRITE_DATA,
     EVENT_SERVICE_REGISTERED,
     EVENT_SERVICE_REMOVED,
     EVENT_STATE_CHANGED,
@@ -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 @@ async def async_stop(self, exit_code: int = 0, *, force: bool = False) -> None:
             # 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,12 @@ async def async_stop(self, exit_code: int = 0, *, force: bool = False) -> None:
         await self.async_block_till_done()
 
         # stage 2
+        self.state = CoreState.writing_data
+        self.async_track_tasks()
+        self.bus.async_fire(EVENT_HOMEASSISTANT_WRITE_DATA)
+        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 @@
 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_START,
+    EVENT_HOMEASSISTANT_WRITE_DATA,
+)
 from homeassistant.core import (
     CoreState,
     HomeAssistant,
@@ -184,7 +187,9 @@ def _async_dump_states(*_: Any) -> None:
         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_WRITE_DATA, _async_dump_states
+        )
 
     @callback
     def async_restore_entity_added(self, entity_id: str) -> None:

homeassistant/helpers/storage.py

@@ -5,7 +5,7 @@
 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_WRITE_DATA
 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 @@ def _async_ensure_stop_listener(self):
         """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_WRITE_DATA, self._async_callback_stop_write
             )
 
     @callback

tests/helpers/test_storage.py

@@ -6,7 +6,7 @@
 
 import pytest
 
-from homeassistant.const import EVENT_HOMEASSISTANT_STOP
+from homeassistant.const import EVENT_HOMEASSISTANT_WRITE_DATA
 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_WRITE_DATA)
     await hass.async_block_till_done()
     assert hass_storage[store.key] == {
         "version": MOCK_VERSION,

tests/test_core.py

@@ -22,6 +22,7 @@
     EVENT_CORE_CONFIG_UPDATE,
     EVENT_HOMEASSISTANT_CLOSE,
     EVENT_HOMEASSISTANT_STOP,
+    EVENT_HOMEASSISTANT_WRITE_DATA,
     EVENT_SERVICE_REGISTERED,
     EVENT_SERVICE_REMOVED,
     EVENT_STATE_CHANGED,
@@ -151,18 +152,23 @@ def test_stage_shutdown():
     """Simulate a shutdown, test calling stuff."""
     hass = get_test_home_assistant()
     test_stop = []
+    test_write_data = []
     test_close = []
     test_all = []
 
     hass.bus.listen(EVENT_HOMEASSISTANT_STOP, lambda event: test_stop.append(event))
+    hass.bus.listen(
+        EVENT_HOMEASSISTANT_WRITE_DATA, lambda event: test_write_data.append(event)
+    )
     hass.bus.listen(EVENT_HOMEASSISTANT_CLOSE, lambda event: test_close.append(event))
     hass.bus.listen("*", lambda event: test_all.append(event))
 
     hass.stop()
 
     assert len(test_stop) == 1
     assert len(test_close) == 1
-    assert len(test_all) == 1
+    assert len(test_write_data) == 1
+    assert len(test_all) == 2
 
 
 class TestHomeAssistant(unittest.TestCase):