Improve performance of google assistant supported checks (#99454)

* Improve performance of google assistant supported checks

* tweak

* tweak

* split function

* tweak

homeassistant/components/google_assistant/helpers.py

@@ -5,6 +5,7 @@ from abc import ABC, abstractmethod
 from asyncio import gather
 from collections.abc import Callable, Mapping
 from datetime import datetime, timedelta
+from functools import lru_cache
 from http import HTTPStatus
 import logging
 import pprint
@@ -490,9 +491,34 @@ def get_google_type(domain, device_class):
     return typ if typ is not None else DOMAIN_TO_GOOGLE_TYPES[domain]
 
 
+@lru_cache(maxsize=4096)
+def supported_traits_for_state(state: State) -> list[type[trait._Trait]]:
+    """Return all supported traits for state."""
+    domain = state.domain
+    attributes = state.attributes
+    features = attributes.get(ATTR_SUPPORTED_FEATURES, 0)
+
+    if not isinstance(features, int):
+        _LOGGER.warning(
+            "Entity %s contains invalid supported_features value %s",
+            state.entity_id,
+            features,
+        )
+        return []
+
+    device_class = state.attributes.get(ATTR_DEVICE_CLASS)
+    return [
+        Trait
+        for Trait in trait.TRAITS
+        if Trait.supported(domain, features, device_class, attributes)
+    ]
+
+
 class GoogleEntity:
     """Adaptation of Entity expressed in Google's terms."""
 
+    __slots__ = ("hass", "config", "state", "_traits")
+
     def __init__(
         self, hass: HomeAssistant, config: AbstractConfig, state: State
     ) -> None:
@@ -502,6 +528,10 @@ class GoogleEntity:
         self.state = state
         self._traits: list[trait._Trait] | None = None
 
+    def __repr__(self) -> str:
+        """Return the representation."""
+        return f"<GoogleEntity {self.state.entity_id}: {self.state.name}>"
+
     @property
     def entity_id(self):
         """Return entity ID."""
@@ -512,26 +542,10 @@ class GoogleEntity:
         """Return traits for entity."""
         if self._traits is not None:
             return self._traits
-
         state = self.state
-        domain = state.domain
-        attributes = state.attributes
-        features = attributes.get(ATTR_SUPPORTED_FEATURES, 0)
-
-        if not isinstance(features, int):
-            _LOGGER.warning(
-                "Entity %s contains invalid supported_features value %s",
-                self.entity_id,
-                features,
-            )
-            return []
-
-        device_class = state.attributes.get(ATTR_DEVICE_CLASS)
-
         self._traits = [
             Trait(self.hass, state, self.config)
-            for Trait in trait.TRAITS
+            for Trait in supported_traits_for_state(state)
-            if Trait.supported(domain, features, device_class, attributes)
         ]
         return self._traits
 
@@ -554,18 +568,8 @@ class GoogleEntity:
 
     @callback
     def is_supported(self) -> bool:
-        """Return if the entity is supported by Google."""
+        """Return if entity is supported."""
-        features: int | None = self.state.attributes.get(ATTR_SUPPORTED_FEATURES)
+        return bool(self.traits())
-
-        result = self.config.is_supported_cache.get(self.entity_id)
-
-        if result is None or result[0] != features:
-            result = self.config.is_supported_cache[self.entity_id] = (
-                features,
-                bool(self.traits()),
-            )
-
-        return result[1]
 
     @callback
     def might_2fa(self) -> bool:
@@ -725,19 +729,64 @@ def deep_update(target, source):
     return target
 
 
+@callback
+def async_get_google_entity_if_supported_cached(
+    hass: HomeAssistant, config: AbstractConfig, state: State
+) -> GoogleEntity | None:
+    """Return a GoogleEntity if entity is supported checking the cache first.
+
+    This function will check the cache, and call async_get_google_entity_if_supported
+    if the entity is not in the cache, which will update the cache.
+    """
+    entity_id = state.entity_id
+    is_supported_cache = config.is_supported_cache
+    features: int | None = state.attributes.get(ATTR_SUPPORTED_FEATURES)
+    if result := is_supported_cache.get(entity_id):
+        cached_features, supported = result
+        if cached_features == features:
+            return GoogleEntity(hass, config, state) if supported else None
+    # Cache miss, check if entity is supported
+    return async_get_google_entity_if_supported(hass, config, state)
+
+
+@callback
+def async_get_google_entity_if_supported(
+    hass: HomeAssistant, config: AbstractConfig, state: State
+) -> GoogleEntity | None:
+    """Return a GoogleEntity if entity is supported.
+
+    This function will update the cache, but it does not check the cache first.
+    """
+    features: int | None = state.attributes.get(ATTR_SUPPORTED_FEATURES)
+    entity = GoogleEntity(hass, config, state)
+    is_supported = bool(entity.traits())
+    config.is_supported_cache[state.entity_id] = (features, is_supported)
+    return entity if is_supported else None
+
+
 @callback
 def async_get_entities(
     hass: HomeAssistant, config: AbstractConfig
 ) -> list[GoogleEntity]:
     """Return all entities that are supported by Google."""
-    entities = []
+    entities: list[GoogleEntity] = []
+    is_supported_cache = config.is_supported_cache
     for state in hass.states.async_all():
-        if state.entity_id in CLOUD_NEVER_EXPOSED_ENTITIES:
+        entity_id = state.entity_id
+        if entity_id in CLOUD_NEVER_EXPOSED_ENTITIES:
             continue
-
+        # Check check inlined for performance to avoid
-        entity = GoogleEntity(hass, config, state)
+        # function calls for every entity since we enumerate
-
+        # the entire state machine here
-        if entity.is_supported():
+        features: int | None = state.attributes.get(ATTR_SUPPORTED_FEATURES)
+        if result := is_supported_cache.get(entity_id):
+            cached_features, supported = result
+            if cached_features == features:
+                if supported:
+                    entities.append(GoogleEntity(hass, config, state))
+                continue
+            # Cached features don't match, fall through to check
+            # if the entity is supported and update the cache.
+        if entity := async_get_google_entity_if_supported(hass, config, state):
             entities.append(entity)
-
     return entities

homeassistant/components/google_assistant/report_state.py

@@ -6,13 +6,17 @@ import logging
 from typing import Any
 
 from homeassistant.const import MATCH_ALL
-from homeassistant.core import CALLBACK_TYPE, HassJob, HomeAssistant, callback
+from homeassistant.core import CALLBACK_TYPE, HassJob, HomeAssistant, State, callback
 from homeassistant.helpers.event import async_call_later, async_track_state_change
 from homeassistant.helpers.significant_change import create_checker
 
 from .const import DOMAIN
 from .error import SmartHomeError
-from .helpers import AbstractConfig, GoogleEntity, async_get_entities
+from .helpers import (
+    AbstractConfig,
+    async_get_entities,
+    async_get_google_entity_if_supported_cached,
+)
 
 # Time to wait until the homegraph updates
 # https://github.com/actions-on-google/smart-home-nodejs/issues/196#issuecomment-439156639
@@ -54,8 +58,10 @@ def async_enable_report_state(hass: HomeAssistant, google_config: AbstractConfig
 
     report_states_job = HassJob(report_states)
 
-    async def async_entity_state_listener(changed_entity, old_state, new_state):
+    async def async_entity_state_listener(
-        nonlocal unsub_pending
+        changed_entity: str, old_state: State | None, new_state: State | None
+    ) -> None:
+        nonlocal unsub_pending, checker
 
         if not hass.is_running:
             return
@@ -66,9 +72,11 @@ def async_enable_report_state(hass: HomeAssistant, google_config: AbstractConfig
         if not google_config.should_expose(new_state):
             return
 
-        entity = GoogleEntity(hass, google_config, new_state)
+        if not (
-
+            entity := async_get_google_entity_if_supported_cached(
-        if not entity.is_supported():
+                hass, google_config, new_state
+            )
+        ):
             return
 
         try:
@@ -77,6 +85,7 @@ def async_enable_report_state(hass: HomeAssistant, google_config: AbstractConfig
             _LOGGER.debug("Not reporting state for %s: %s", changed_entity, err.code)
             return
 
+        assert checker is not None
         if not checker.async_is_significant_change(new_state, extra_arg=entity_data):
             return
 

tests/components/google_assistant/test_helpers.py

@@ -447,32 +447,53 @@ async def test_config_local_sdk_warn_version(
     ) in caplog.text
 
 
-def test_is_supported_cached() -> None:
+def test_async_get_entities_cached(hass: HomeAssistant) -> None:
-    """Test is_supported is cached."""
+    """Test async_get_entities is cached."""
     config = MockConfig()
 
-    def entity(features: int):
+    hass.states.async_set("light.ceiling_lights", "off")
-        return helpers.GoogleEntity(
+    hass.states.async_set("light.bed_light", "off")
-            None,
+    hass.states.async_set("not_supported.not_supported", "off")
-            config,
+
-            State("test.entity_id", "on", {"supported_features": features}),
+    google_entities = helpers.async_get_entities(hass, config)
-        )
+    assert len(google_entities) == 2
+    assert config.is_supported_cache == {
+        "light.bed_light": (None, True),
+        "light.ceiling_lights": (None, True),
+        "not_supported.not_supported": (None, False),
+    }
 
     with patch(
         "homeassistant.components.google_assistant.helpers.GoogleEntity.traits",
-        return_value=[1],
+        return_value=RuntimeError("Should not be called"),
-    ) as mock_traits:
+    ):
-        assert entity(1).is_supported() is True
+        google_entities = helpers.async_get_entities(hass, config)
-        assert len(mock_traits.mock_calls) == 1
 
-        # Supported feature changes, so we calculate again
+    assert len(google_entities) == 2
-        assert entity(2).is_supported() is True
+    assert config.is_supported_cache == {
-        assert len(mock_traits.mock_calls) == 2
+        "light.bed_light": (None, True),
+        "light.ceiling_lights": (None, True),
+        "not_supported.not_supported": (None, False),
+    }
 
-        mock_traits.reset_mock()
+    hass.states.async_set("light.new", "on")
+    google_entities = helpers.async_get_entities(hass, config)
 
-        # Supported feature is same, so we do not calculate again
+    assert len(google_entities) == 3
-        mock_traits.side_effect = ValueError
+    assert config.is_supported_cache == {
+        "light.bed_light": (None, True),
+        "light.new": (None, True),
+        "light.ceiling_lights": (None, True),
+        "not_supported.not_supported": (None, False),
+    }
 
-        assert entity(2).is_supported() is True
+    hass.states.async_set("light.new", "on", {"supported_features": 1})
-        assert len(mock_traits.mock_calls) == 0
+    google_entities = helpers.async_get_entities(hass, config)
+
+    assert len(google_entities) == 3
+    assert config.is_supported_cache == {
+        "light.bed_light": (None, True),
+        "light.new": (1, True),
+        "light.ceiling_lights": (None, True),
+        "not_supported.not_supported": (None, False),
+    }

tests/components/google_assistant/test_report_state.py

@@ -69,7 +69,7 @@ async def test_report_state(
 
     # Test that if serialize returns same value, we don't send
     with patch(
-        "homeassistant.components.google_assistant.report_state.GoogleEntity.query_serialize",
+        "homeassistant.components.google_assistant.helpers.GoogleEntity.query_serialize",
         return_value={"same": "info"},
     ), patch.object(BASIC_CONFIG, "async_report_state_all", AsyncMock()) as mock_report:
         # New state, so reported
@@ -104,7 +104,7 @@ async def test_report_state(
     with patch.object(
         BASIC_CONFIG, "async_report_state_all", AsyncMock()
     ) as mock_report, patch(
-        "homeassistant.components.google_assistant.report_state.GoogleEntity.query_serialize",
+        "homeassistant.components.google_assistant.helpers.GoogleEntity.query_serialize",
         side_effect=error.SmartHomeError("mock-error", "mock-msg"),
     ):
         hass.states.async_set("light.kitchen", "off")