Drop unused restore state snapshots from memory (#183342)

Co-authored-by: Claude <noreply@anthropic.com>

homeassistant/components/utility_meter/diagnostics.py

@@ -13,22 +13,18 @@ async def async_get_config_entry_diagnostics(
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
 
-    tariff_sensors = []
-
-    for sensor in hass.data[DATA_UTILITY][entry.entry_id][DATA_TARIFF_SENSORS]:
-        restored_last_extra_data = await sensor.async_get_last_extra_data()
-
-        tariff_sensors.append(
-            {
-                "name": sensor.name,
-                "entity_id": sensor.entity_id,
-                "extra_attributes": sensor.extra_state_attributes,
-                "last_sensor_data": restored_last_extra_data,
-                "period": sensor._period,  # noqa: SLF001
-                "cron": sensor._cron_pattern,  # noqa: SLF001
-                "source": sensor._sensor_source_id,  # noqa: SLF001
-            }
-        )
+    tariff_sensors = [
+        {
+            "name": sensor.name,
+            "entity_id": sensor.entity_id,
+            "extra_attributes": sensor.extra_state_attributes,
+            "last_sensor_data": sensor.extra_restore_state_data.as_dict(),
+            "period": sensor._period,  # noqa: SLF001
+            "cron": sensor._cron_pattern,  # noqa: SLF001
+            "source": sensor._sensor_source_id,  # noqa: SLF001
+        }
+        for sensor in hass.data[DATA_UTILITY][entry.entry_id][DATA_TARIFF_SENSORS]
+    ]
 
     return {
         "config_entry": entry,

homeassistant/helpers/restore_state.py

@@ -169,6 +169,8 @@ class RestoreStateData:
         This includes the states of all registered entities, as well as the
         stored states from the previous run, which have not been created as
         entities on this run, and have not expired.
+
+        Stored states that will not be saved are dropped from memory too.
         """
         now = dt_util.utcnow()
         all_states = self.hass.states.async_all()
@@ -199,6 +201,7 @@ class RestoreStateData:
                 )
             )
         expiration_time = now - STATE_EXPIRATION
+        last_states: dict[str, StoredState] = {}
 
         for entity_id, stored_state in self.last_states.items():
             # Don't save old states that have entities in the current run
@@ -212,6 +215,9 @@ class RestoreStateData:
                 continue
 
             stored_states.append(stored_state)
+            last_states[entity_id] = stored_state
+
+        self.last_states = last_states
 
         return stored_states
 

tests/components/utility_meter/snapshots/test_diagnostics.ambr

@@ -44,9 +44,10 @@
           'tariff': 'tariff0',
         }),
         'last_sensor_data': dict({
+          'input_device_class': 'None',
           'last_period': '0',
           'last_reset': '2024-04-05T00:00:00+00:00',
-          'last_valid_state': 3,
+          'last_valid_state': None,
           'native_unit_of_measurement': 'kWh',
           'native_value': dict({
             '__type': "<class 'decimal.Decimal'>",
@@ -70,9 +71,10 @@
           'tariff': 'tariff1',
         }),
         'last_sensor_data': dict({
+          'input_device_class': 'None',
           'last_period': '0',
           'last_reset': '2024-04-05T00:00:00+00:00',
-          'last_valid_state': 7,
+          'last_valid_state': None,
           'native_unit_of_measurement': 'kWh',
           'native_value': dict({
             '__type': "<class 'decimal.Decimal'>",

tests/helpers/test_restore_state.py

@@ -321,6 +321,8 @@ async def test_dump_data(hass: HomeAssistant) -> None:
     assert state1["state"]["state"] == "off"
     assert state2["state"]["entity_id"] == "input_boolean.b5"
     assert state2["state"]["state"] == "off"
+    # States that are not written anymore are dropped from memory as well
+    assert list(data.last_states) == ["input_boolean.b3", "input_boolean.b5"]
 
     # Test that removed entities are not persisted
     await entity.async_remove()