From 78e546b35a433b92ed6648a10d87508efb5ecab1 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Wed, 1 Nov 2023 04:25:02 -0500
Subject: [PATCH] Avoid enumerating the whole state machine on api service
 calls (#103147)

---
 homeassistant/components/api/__init__.py | 27 ++++++++++++++++++------
 1 file changed, 20 insertions(+), 7 deletions(-)

diff --git a/homeassistant/components/api/__init__.py b/homeassistant/components/api/__init__.py
index 0cade0f81ca..077e5ec9093 100644
--- a/homeassistant/components/api/__init__.py
+++ b/homeassistant/components/api/__init__.py
@@ -1,9 +1,11 @@
 """Rest API for Home Assistant."""
 import asyncio
 from asyncio import timeout
+from collections.abc import Collection
 from functools import lru_cache
 from http import HTTPStatus
 import logging
+from typing import Any
 
 from aiohttp import web
 from aiohttp.web_exceptions import HTTPBadRequest
@@ -16,6 +18,7 @@ from homeassistant.components.http import HomeAssistantView, require_admin
 from homeassistant.const import (
     CONTENT_TYPE_JSON,
     EVENT_HOMEASSISTANT_STOP,
+    EVENT_STATE_CHANGED,
     MATCH_ALL,
     URL_API,
     URL_API_COMPONENTS,
@@ -38,10 +41,12 @@ from homeassistant.exceptions import (
     Unauthorized,
 )
 from homeassistant.helpers import config_validation as cv, template
+from homeassistant.helpers.event import EventStateChangedData
 from homeassistant.helpers.json import json_dumps
 from homeassistant.helpers.service import async_get_all_descriptions
-from homeassistant.helpers.typing import ConfigType
+from homeassistant.helpers.typing import ConfigType, EventType
 from homeassistant.util.json import json_loads
+from homeassistant.util.read_only_dict import ReadOnlyDict
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -369,6 +374,18 @@ class APIDomainServicesView(HomeAssistantView):
             )
 
         context = self.context(request)
+        changed_states: list[ReadOnlyDict[str, Collection[Any]]] = []
+
+        @ha.callback
+        def _async_save_changed_entities(
+            event: EventType[EventStateChangedData],
+        ) -> None:
+            if event.context == context and (state := event.data["new_state"]):
+                changed_states.append(state.as_dict())
+
+        cancel_listen = hass.bus.async_listen(
+            EVENT_STATE_CHANGED, _async_save_changed_entities, run_immediately=True
+        )
 
         try:
             await hass.services.async_call(
@@ -376,12 +393,8 @@ class APIDomainServicesView(HomeAssistantView):
             )
         except (vol.Invalid, ServiceNotFound) as ex:
             raise HTTPBadRequest() from ex
-
-        changed_states = []
-
-        for state in hass.states.async_all():
-            if state.context is context:
-                changed_states.append(state)
+        finally:
+            cancel_listen()
 
         return self.json(changed_states)