From 9df00bd285a656814c8073ad92e88d1778efdf25 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sat, 18 Feb 2023 23:02:36 -0600
Subject: [PATCH] Adjust recorder LRU cache to align with the number of
 entities (#88350)

* Adjust size of recorder LRU based on number of entities

If there are a large number of entities the cache would
get thrashed as there were more state attributes being
recorded than the size of the cache. This meant we had
to go back to the database to do lookups frequently when
an instance has more than 2048 entities that change
frequently

* add a test

* do not actually record 4096 states

* patch target

* patch target
---
 homeassistant/components/recorder/core.py | 23 +++++++++++++++++++++--
 tests/components/recorder/test_init.py    | 15 +++++++++++++++
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/recorder/core.py b/homeassistant/components/recorder/core.py
index 665028b074d..7cd0ad56f79 100644
--- a/homeassistant/components/recorder/core.py
+++ b/homeassistant/components/recorder/core.py
@@ -411,6 +411,7 @@ class Recorder(threading.Thread):
     @callback
     def _async_hass_started(self, hass: HomeAssistant) -> None:
         """Notify that hass has started."""
+        self.async_adjust_lru()
         self._hass_started.set_result(None)
 
     @callback
@@ -475,7 +476,25 @@ class Recorder(threading.Thread):
             self.queue_task(PerodicCleanupTask())
 
     @callback
-    def async_periodic_statistics(self, now: datetime) -> None:
+    def _async_five_minute_tasks(self, now: datetime) -> None:
+        """Run tasks every five minutes."""
+        self.async_adjust_lru()
+        self.async_periodic_statistics()
+
+    @callback
+    def async_adjust_lru(self) -> None:
+        """Trigger the LRU adjustment.
+
+        If the number of entities has increased, increase the size of the LRU
+        cache to avoid thrashing.
+        """
+        current_size = self._state_attributes_ids.get_size()
+        new_size = self.hass.states.async_entity_ids_count() * 2
+        if new_size > current_size:
+            self._state_attributes_ids.set_size(new_size)
+
+    @callback
+    def async_periodic_statistics(self) -> None:
         """Trigger the statistics run.
 
         Short term statistics run every 5 minutes
@@ -570,7 +589,7 @@ class Recorder(threading.Thread):
 
         # Compile short term statistics every 5 minutes
         self._periodic_listener = async_track_utc_time_change(
-            self.hass, self.async_periodic_statistics, minute=range(0, 60, 5), second=10
+            self.hass, self._async_five_minute_tasks, minute=range(0, 60, 5), second=10
         )
 
     async def _async_wait_for_started(self) -> object | None:
diff --git a/tests/components/recorder/test_init.py b/tests/components/recorder/test_init.py
index 5c542405376..11233af1462 100644
--- a/tests/components/recorder/test_init.py
+++ b/tests/components/recorder/test_init.py
@@ -2075,3 +2075,18 @@ async def test_excluding_attributes_by_integration(
     expected = _state_with_context(hass, entity_id)
     expected.attributes = {"test_attr": 5}
     assert state.as_dict() == expected.as_dict()
+
+
+async def test_lru_increases_with_many_entities(
+    recorder_mock: Recorder, hass: HomeAssistant
+) -> None:
+    """Test that the recorder's internal LRU cache increases with many entities."""
+    # We do not actually want to record 4096 entities so we mock the entity count
+    mock_entity_count = 4096
+    with patch.object(
+        hass.states, "async_entity_ids_count", return_value=mock_entity_count
+    ):
+        async_fire_time_changed(hass, dt_util.utcnow() + timedelta(minutes=10))
+        await hass.async_block_till_done()
+
+    assert recorder_mock._state_attributes_ids.get_size() == mock_entity_count * 2