From b51dcb600ea7f8abd123c4ababc2a816e36c77c9 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Tue, 20 Jun 2023 14:48:28 +0100
Subject: [PATCH] Avoid enumerating the whole state machine to find zone
 entities (#94866)

---
 homeassistant/components/zone/__init__.py | 37 +++++++++++++++++++++--
 1 file changed, 34 insertions(+), 3 deletions(-)

diff --git a/homeassistant/components/zone/__init__.py b/homeassistant/components/zone/__init__.py
index b7212e91091..35d835c8f16 100644
--- a/homeassistant/components/zone/__init__.py
+++ b/homeassistant/components/zone/__init__.py
@@ -99,6 +99,8 @@ STORAGE_VERSION = 1
 
 ENTITY_ID_SORTER = attrgetter("entity_id")
 
+ZONE_ENTITY_IDS = "zone_entity_ids"
+
 
 @bind_hass
 def async_active_zone(
@@ -111,9 +113,15 @@ def async_active_zone(
     # Sort entity IDs so that we are deterministic if equal distance to 2 zones
     min_dist = None
     closest = None
-
-    for zone in sorted(hass.states.async_all(DOMAIN), key=ENTITY_ID_SORTER):
-        if zone.state == STATE_UNAVAILABLE or zone.attributes.get(ATTR_PASSIVE):
+    # This can be called before async_setup by device tracker
+    zone_entity_ids: list[str] = hass.data.get(ZONE_ENTITY_IDS, [])
+    for entity_id in zone_entity_ids:
+        zone = hass.states.get(entity_id)
+        if (
+            not zone
+            or zone.state == STATE_UNAVAILABLE
+            or zone.attributes.get(ATTR_PASSIVE)
+        ):
             continue
 
         zone_dist = distance(
@@ -141,6 +149,27 @@ def async_active_zone(
     return closest
 
 
+@callback
+def async_setup_track_zone_entity_ids(hass: HomeAssistant) -> None:
+    """Set up track of entity IDs for zones."""
+    zone_entity_ids: list[str] = hass.states.async_entity_ids(DOMAIN)
+    hass.data[ZONE_ENTITY_IDS] = zone_entity_ids
+
+    @callback
+    def _async_add_zone_entity_id(event_: Event) -> None:
+        """Add zone entity ID."""
+        zone_entity_ids.append(event_.data[ATTR_ENTITY_ID])
+        zone_entity_ids.sort()
+
+    @callback
+    def _async_remove_zone_entity_id(event_: Event) -> None:
+        """Remove zone entity ID."""
+        zone_entity_ids.remove(event_.data[ATTR_ENTITY_ID])
+
+    event.async_track_state_added_domain(hass, DOMAIN, _async_add_zone_entity_id)
+    event.async_track_state_removed_domain(hass, DOMAIN, _async_remove_zone_entity_id)
+
+
 def in_zone(zone: State, latitude: float, longitude: float, radius: float = 0) -> bool:
     """Test if given latitude, longitude is in given zone.
 
@@ -184,6 +213,8 @@ class ZoneStorageCollection(collection.DictStorageCollection):
 
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up configured zones as well as Home Assistant zone if necessary."""
+    async_setup_track_zone_entity_ids(hass)
+
     component = entity_component.EntityComponent[Zone](_LOGGER, DOMAIN, hass)
     id_manager = collection.IDManager()