Don't write storage to disk while stopping (#33456)

* Don't write storage to disk while stopping

* rework change

* lint

* remove delay save and schedule final write at stop

* update tests

* fix test component using private methods

* cleanup

* always listen

* use stop in restore state again

* whitelist JSON exceptions for later

* review comment

* make zwave tests use mock storage

homeassistant/core.py

@@ -152,7 +152,7 @@ class CoreState(enum.Enum):
     starting = "STARTING"
     running = "RUNNING"
     stopping = "STOPPING"
-    writing_data = "WRITING_DATA"
+    final_write = "FINAL_WRITE"
 
     def __str__(self) -> str:
         """Return the event."""
@@ -414,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 or self.state == CoreState.writing_data:
+            if self.state == CoreState.stopping or self.state == CoreState.final_write:
                 _LOGGER.info("async_stop called twice: ignored")
                 return
             if self.state == CoreState.starting:
@@ -428,7 +428,7 @@ class HomeAssistant:
         await self.async_block_till_done()
 
         # stage 2
-        self.state = CoreState.writing_data
+        self.state = CoreState.final_write
         self.bus.async_fire(EVENT_HOMEASSISTANT_FINAL_WRITE)
         await self.async_block_till_done()
 

homeassistant/helpers/restore_state.py

@@ -4,10 +4,7 @@ from datetime import datetime, timedelta
 import logging
 from typing import Any, Awaitable, Dict, List, Optional, Set, cast
 
-from homeassistant.const import (
+from homeassistant.const import EVENT_HOMEASSISTANT_START, EVENT_HOMEASSISTANT_STOP
-    EVENT_HOMEASSISTANT_FINAL_WRITE,
-    EVENT_HOMEASSISTANT_START,
-)
 from homeassistant.core import (
     CoreState,
     HomeAssistant,
@@ -187,9 +184,7 @@ 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(
+        self.hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _async_dump_states)
-            EVENT_HOMEASSISTANT_FINAL_WRITE, _async_dump_states
-        )
 
     @callback
     def async_restore_entity_added(self, entity_id: str) -> None:

homeassistant/helpers/storage.py

@@ -6,7 +6,7 @@ import os
 from typing import Any, Callable, Dict, List, Optional, Type, Union
 
 from homeassistant.const import EVENT_HOMEASSISTANT_FINAL_WRITE
-from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
+from homeassistant.core import CALLBACK_TYPE, CoreState, HomeAssistant, callback
 from homeassistant.helpers.event import async_call_later
 from homeassistant.loader import bind_hass
 from homeassistant.util import json as json_util
@@ -72,7 +72,7 @@ class Store:
         self._private = private
         self._data: Optional[Dict[str, Any]] = None
         self._unsub_delay_listener: Optional[CALLBACK_TYPE] = None
-        self._unsub_stop_listener: Optional[CALLBACK_TYPE] = None
+        self._unsub_final_write_listener: Optional[CALLBACK_TYPE] = None
         self._write_lock = asyncio.Lock()
         self._load_task: Optional[asyncio.Future] = None
         self._encoder = encoder
@@ -132,7 +132,12 @@ class Store:
         self._data = {"version": self.version, "key": self.key, "data": data}
 
         self._async_cleanup_delay_listener()
-        self._async_cleanup_stop_listener()
+        self._async_cleanup_final_write_listener()
+
+        if self.hass.state == CoreState.stopping:
+            self._async_ensure_final_write_listener()
+            return
+
         await self._async_handle_write_data()
 
     @callback
@@ -141,27 +146,31 @@ class Store:
         self._data = {"version": self.version, "key": self.key, "data_func": data_func}
 
         self._async_cleanup_delay_listener()
+        self._async_cleanup_final_write_listener()
+
+        if self.hass.state == CoreState.stopping:
+            self._async_ensure_final_write_listener()
+            return
 
         self._unsub_delay_listener = async_call_later(
             self.hass, delay, self._async_callback_delayed_write
         )
-
+        self._async_ensure_final_write_listener()
-        self._async_ensure_stop_listener()
 
     @callback
-    def _async_ensure_stop_listener(self):
+    def _async_ensure_final_write_listener(self):
         """Ensure that we write if we quit before delay has passed."""
-        if self._unsub_stop_listener is None:
+        if self._unsub_final_write_listener is None:
-            self._unsub_stop_listener = self.hass.bus.async_listen_once(
+            self._unsub_final_write_listener = self.hass.bus.async_listen_once(
-                EVENT_HOMEASSISTANT_FINAL_WRITE, self._async_callback_stop_write
+                EVENT_HOMEASSISTANT_FINAL_WRITE, self._async_callback_final_write
             )
 
     @callback
-    def _async_cleanup_stop_listener(self):
+    def _async_cleanup_final_write_listener(self):
         """Clean up a stop listener."""
-        if self._unsub_stop_listener is not None:
+        if self._unsub_final_write_listener is not None:
-            self._unsub_stop_listener()
+            self._unsub_final_write_listener()
-            self._unsub_stop_listener = None
+            self._unsub_final_write_listener = None
 
     @callback
     def _async_cleanup_delay_listener(self):
@@ -172,13 +181,17 @@ class Store:
 
     async def _async_callback_delayed_write(self, _now):
         """Handle a delayed write callback."""
+        # catch the case where a call is scheduled and then we stop Home Assistant
+        if self.hass.state == CoreState.stopping:
+            self._async_ensure_final_write_listener()
+            return
         self._unsub_delay_listener = None
-        self._async_cleanup_stop_listener()
+        self._async_cleanup_final_write_listener()
         await self._async_handle_write_data()
 
-    async def _async_callback_stop_write(self, _event):
+    async def _async_callback_final_write(self, _event):
-        """Handle a write because Home Assistant is stopping."""
+        """Handle a write because Home Assistant is in final write state."""
-        self._unsub_stop_listener = None
+        self._unsub_final_write_listener = None
         self._async_cleanup_delay_listener()
         await self._async_handle_write_data()
 

tests/common.py

@@ -1005,7 +1005,7 @@ async def flush_store(store):
     if store._data is None:
         return
 
-    store._async_cleanup_stop_listener()
+    store._async_cleanup_final_write_listener()
     store._async_cleanup_delay_listener()
     await store._async_handle_write_data()
 
@@ -1018,7 +1018,7 @@ async def get_system_health_info(hass, domain):
 def mock_integration(hass, module):
     """Mock an integration."""
     integration = loader.Integration(
-        hass, f"homeassistant.components.{module.DOMAIN}", None, module.mock_manifest(),
+        hass, f"homeassistant.components.{module.DOMAIN}", None, module.mock_manifest()
     )
 
     _LOGGER.info("Adding mock integration: %s", module.DOMAIN)

tests/components/zwave/test_init.py

@@ -28,6 +28,7 @@ from tests.common import (
     get_test_home_assistant,
     mock_coro,
     mock_registry,
+    mock_storage,
 )
 from tests.mock.zwave import MockEntityValues, MockNetwork, MockNode, MockValue
 
@@ -827,6 +828,8 @@ class TestZWaveDeviceEntityValues(unittest.TestCase):
     def setUp(self):
         """Initialize values for this testcase class."""
         self.hass = get_test_home_assistant()
+        self.mock_storage = mock_storage()
+        self.mock_storage.__enter__()
         self.hass.start()
         self.registry = mock_registry(self.hass)
 
@@ -862,6 +865,7 @@ class TestZWaveDeviceEntityValues(unittest.TestCase):
     def tearDown(self):  # pylint: disable=invalid-name
         """Stop everything that was started."""
         self.hass.stop()
+        self.mock_storage.__exit__(None, None, None)
 
     @patch.object(zwave, "import_module")
     @patch.object(zwave, "discovery")
@@ -1194,6 +1198,8 @@ class TestZWaveServices(unittest.TestCase):
     def setUp(self):
         """Initialize values for this testcase class."""
         self.hass = get_test_home_assistant()
+        self.mock_storage = mock_storage()
+        self.mock_storage.__enter__()
         self.hass.start()
 
         # Initialize zwave
@@ -1209,6 +1215,7 @@ class TestZWaveServices(unittest.TestCase):
         self.hass.services.call("zwave", "stop_network", {})
         self.hass.block_till_done()
         self.hass.stop()
+        self.mock_storage.__exit__(None, None, None)
 
     def test_add_node(self):
         """Test zwave add_node service."""

tests/conftest.py

@@ -19,7 +19,10 @@ from homeassistant.exceptions import ServiceNotFound
 from homeassistant.setup import async_setup_component
 from homeassistant.util import location
 
-from tests.ignore_uncaught_exceptions import IGNORE_UNCAUGHT_EXCEPTIONS
+from tests.ignore_uncaught_exceptions import (
+    IGNORE_UNCAUGHT_EXCEPTIONS,
+    IGNORE_UNCAUGHT_JSON_EXCEPTIONS,
+)
 
 pytest.register_assert_rewrite("tests.common")
 
@@ -104,6 +107,13 @@ def hass(loop, hass_storage, request):
             continue
         if isinstance(ex, ServiceNotFound):
             continue
+        if (
+            isinstance(ex, TypeError)
+            and "is not JSON serializable" in str(ex)
+            and (request.module.__name__, request.function.__name__)
+            in IGNORE_UNCAUGHT_JSON_EXCEPTIONS
+        ):
+            continue
         raise ex
 
 
@@ -211,7 +221,7 @@ def hass_client(hass, aiohttp_client, hass_access_token):
     async def auth_client():
         """Return an authenticated client."""
         return await aiohttp_client(
-            hass.http.app, headers={"Authorization": f"Bearer {hass_access_token}"},
+            hass.http.app, headers={"Authorization": f"Bearer {hass_access_token}"}
         )
 
     return auth_client

tests/helpers/test_storage.py

@@ -6,7 +6,11 @@ from unittest.mock import Mock, patch
 
 import pytest
 
-from homeassistant.const import EVENT_HOMEASSISTANT_FINAL_WRITE
+from homeassistant.const import (
+    EVENT_HOMEASSISTANT_FINAL_WRITE,
+    EVENT_HOMEASSISTANT_STOP,
+)
+from homeassistant.core import CoreState
 from homeassistant.helpers import storage
 from homeassistant.util import dt
 
@@ -79,10 +83,18 @@ async def test_saving_with_delay(hass, store, hass_storage):
     }
 
 
-async def test_saving_on_stop(hass, hass_storage):
+async def test_saving_on_final_write(hass, hass_storage):
     """Test delayed saves trigger when we quit Home Assistant."""
     store = storage.Store(hass, MOCK_VERSION, MOCK_KEY)
-    store.async_delay_save(lambda: MOCK_DATA, 1)
+    store.async_delay_save(lambda: MOCK_DATA, 5)
+    assert store.key not in hass_storage
+
+    hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP)
+    hass.state = CoreState.stopping
+    await hass.async_block_till_done()
+
+    async_fire_time_changed(hass, dt.utcnow() + timedelta(seconds=10))
+    await hass.async_block_till_done()
     assert store.key not in hass_storage
 
     hass.bus.async_fire(EVENT_HOMEASSISTANT_FINAL_WRITE)
@@ -94,6 +106,43 @@ async def test_saving_on_stop(hass, hass_storage):
     }
 
 
+async def test_not_delayed_saving_while_stopping(hass, hass_storage):
+    """Test delayed saves don't write after the stop event has fired."""
+    store = storage.Store(hass, MOCK_VERSION, MOCK_KEY)
+    hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP)
+    await hass.async_block_till_done()
+    hass.state = CoreState.stopping
+
+    store.async_delay_save(lambda: MOCK_DATA, 1)
+    async_fire_time_changed(hass, dt.utcnow() + timedelta(seconds=2))
+    await hass.async_block_till_done()
+    assert store.key not in hass_storage
+
+
+async def test_not_delayed_saving_after_stopping(hass, hass_storage):
+    """Test delayed saves don't write after stop if issued before stopping Home Assistant."""
+    store = storage.Store(hass, MOCK_VERSION, MOCK_KEY)
+    store.async_delay_save(lambda: MOCK_DATA, 10)
+    assert store.key not in hass_storage
+
+    hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP)
+    hass.state = CoreState.stopping
+    await hass.async_block_till_done()
+    assert store.key not in hass_storage
+
+    async_fire_time_changed(hass, dt.utcnow() + timedelta(seconds=15))
+    await hass.async_block_till_done()
+    assert store.key not in hass_storage
+
+
+async def test_not_saving_while_stopping(hass, hass_storage):
+    """Test saves don't write when stopping Home Assistant."""
+    store = storage.Store(hass, MOCK_VERSION, MOCK_KEY)
+    hass.state = CoreState.stopping
+    await store.async_save(MOCK_DATA)
+    assert store.key not in hass_storage
+
+
 async def test_loading_while_delay(hass, store, hass_storage):
     """Test we load new data even if not written yet."""
     await store.async_save({"delay": "no"})

tests/ignore_uncaught_exceptions.py

@@ -89,3 +89,42 @@ IGNORE_UNCAUGHT_EXCEPTIONS = [
     ("tests.components.yr.test_sensor", "test_forecast_setup"),
     ("tests.components.zwave.test_init", "test_power_schemes"),
 ]
+
+IGNORE_UNCAUGHT_JSON_EXCEPTIONS = [
+    ("tests.components.spotify.test_config_flow", "test_full_flow"),
+    ("tests.components.smartthings.test_init", "test_config_entry_loads_platforms"),
+    (
+        "tests.components.smartthings.test_init",
+        "test_scenes_unauthorized_loads_platforms",
+    ),
+    (
+        "tests.components.smartthings.test_init",
+        "test_config_entry_loads_unconnected_cloud",
+    ),
+    ("tests.components.samsungtv.test_config_flow", "test_ssdp"),
+    ("tests.components.samsungtv.test_config_flow", "test_user_websocket"),
+    ("tests.components.samsungtv.test_config_flow", "test_user_already_configured"),
+    ("tests.components.samsungtv.test_config_flow", "test_autodetect_websocket"),
+    ("tests.components.samsungtv.test_config_flow", "test_autodetect_websocket_ssl"),
+    ("tests.components.samsungtv.test_config_flow", "test_ssdp_already_configured"),
+    ("tests.components.samsungtv.test_config_flow", "test_ssdp_noprefix"),
+    ("tests.components.samsungtv.test_config_flow", "test_user_legacy"),
+    ("tests.components.samsungtv.test_config_flow", "test_autodetect_legacy"),
+    (
+        "tests.components.samsungtv.test_media_player",
+        "test_select_source_invalid_source",
+    ),
+    (
+        "tests.components.samsungtv.test_media_player",
+        "test_play_media_channel_as_string",
+    ),
+    (
+        "tests.components.samsungtv.test_media_player",
+        "test_play_media_channel_as_non_positive",
+    ),
+    ("tests.components.samsungtv.test_media_player", "test_turn_off_websocket"),
+    ("tests.components.samsungtv.test_media_player", "test_play_media_invalid_type"),
+    ("tests.components.harmony.test_config_flow", "test_form_import"),
+    ("tests.components.harmony.test_config_flow", "test_form_ssdp"),
+    ("tests.components.harmony.test_config_flow", "test_user_form"),
+]