Websocket api to subscribe to entities (payloads reduced by ~80%+ vs state_changed events) (#67891)

homeassistant/components/websocket_api/messages.py

@@ -7,7 +7,7 @@ from typing import Any, Final
 
 import voluptuous as vol
 
-from homeassistant.core import Event
+from homeassistant.core import Event, State
 from homeassistant.helpers import config_validation as cv
 from homeassistant.util.json import (
     find_paths_unserializable_data,
@@ -31,6 +31,19 @@ BASE_COMMAND_MESSAGE_SCHEMA: Final = vol.Schema({vol.Required("id"): cv.positive
 IDEN_TEMPLATE: Final = "__IDEN__"
 IDEN_JSON_TEMPLATE: Final = '"__IDEN__"'
 
+COMPRESSED_STATE_STATE = "s"
+COMPRESSED_STATE_ATTRIBUTES = "a"
+COMPRESSED_STATE_CONTEXT = "c"
+COMPRESSED_STATE_LAST_CHANGED = "lc"
+COMPRESSED_STATE_LAST_UPDATED = "lu"
+
+STATE_DIFF_ADDITIONS = "+"
+STATE_DIFF_REMOVALS = "-"
+
+ENTITY_EVENT_ADD = "a"
+ENTITY_EVENT_REMOVE = "r"
+ENTITY_EVENT_CHANGE = "c"
+
 
 def result_message(iden: int, result: Any = None) -> dict[str, Any]:
     """Return a success result message."""
@@ -74,6 +87,110 @@ def _cached_event_message(event: Event) -> str:
     return message_to_json(event_message(IDEN_TEMPLATE, event))
 
 
+def cached_state_diff_message(iden: int, event: Event) -> str:
+    """Return an event message.
+
+    Serialize to json once per message.
+
+    Since we can have many clients connected that are
+    all getting many of the same events (mostly state changed)
+    we can avoid serializing the same data for each connection.
+    """
+    return _cached_state_diff_message(event).replace(IDEN_JSON_TEMPLATE, str(iden), 1)
+
+
+@lru_cache(maxsize=128)
+def _cached_state_diff_message(event: Event) -> str:
+    """Cache and serialize the event to json.
+
+    The IDEN_TEMPLATE is used which will be replaced
+    with the actual iden in cached_event_message
+    """
+    return message_to_json(event_message(IDEN_TEMPLATE, _state_diff_event(event)))
+
+
+def _state_diff_event(event: Event) -> dict:
+    """Convert a state_changed event to the minimal version.
+
+    State update example
+
+    {
+        "a": {entity_id: compressed_state,…}
+        "c": {entity_id: diff,…}
+        "r": [entity_id,…]
+    }
+    """
+    if (event_new_state := event.data["new_state"]) is None:
+        return {ENTITY_EVENT_REMOVE: [event.data["entity_id"]]}
+    assert isinstance(event_new_state, State)
+    if (event_old_state := event.data["old_state"]) is None:
+        return {
+            ENTITY_EVENT_ADD: {
+                event_new_state.entity_id: compressed_state_dict_add(event_new_state)
+            }
+        }
+    assert isinstance(event_old_state, State)
+    return _state_diff(event_old_state, event_new_state)
+
+
+def _state_diff(
+    old_state: State, new_state: State
+) -> dict[str, dict[str, dict[str, dict[str, str | list[str]]]]]:
+    """Create a diff dict that can be used to overlay changes."""
+    diff: dict = {STATE_DIFF_ADDITIONS: {}}
+    additions = diff[STATE_DIFF_ADDITIONS]
+    if old_state.state != new_state.state:
+        additions[COMPRESSED_STATE_STATE] = new_state.state
+    if old_state.last_changed != new_state.last_changed:
+        additions[COMPRESSED_STATE_LAST_CHANGED] = new_state.last_changed.timestamp()
+    elif old_state.last_updated != new_state.last_updated:
+        additions[COMPRESSED_STATE_LAST_UPDATED] = new_state.last_updated.timestamp()
+    if old_state.context.parent_id != new_state.context.parent_id:
+        additions.setdefault(COMPRESSED_STATE_CONTEXT, {})[
+            "parent_id"
+        ] = new_state.context.parent_id
+    if old_state.context.user_id != new_state.context.user_id:
+        additions.setdefault(COMPRESSED_STATE_CONTEXT, {})[
+            "user_id"
+        ] = new_state.context.user_id
+    if old_state.context.id != new_state.context.id:
+        if COMPRESSED_STATE_CONTEXT in additions:
+            additions[COMPRESSED_STATE_CONTEXT]["id"] = new_state.context.id
+        else:
+            additions[COMPRESSED_STATE_CONTEXT] = new_state.context.id
+    old_attributes = old_state.attributes
+    for key, value in new_state.attributes.items():
+        if old_attributes.get(key) != value:
+            additions.setdefault(COMPRESSED_STATE_ATTRIBUTES, {})[key] = value
+    if removed := set(old_attributes).difference(new_state.attributes):
+        diff[STATE_DIFF_REMOVALS] = {COMPRESSED_STATE_ATTRIBUTES: removed}
+    return {ENTITY_EVENT_CHANGE: {new_state.entity_id: diff}}
+
+
+def compressed_state_dict_add(state: State) -> dict[str, Any]:
+    """Build a compressed dict of a state for adds.
+
+    Omits the lu (last_updated) if it matches (lc) last_changed.
+
+    Sends c (context) as a string if it only contains an id.
+    """
+    if state.context.parent_id is None and state.context.user_id is None:
+        context: dict[str, Any] | str = state.context.id  # type: ignore[unreachable]
+    else:
+        context = state.context.as_dict()
+    compressed_state: dict[str, Any] = {
+        COMPRESSED_STATE_STATE: state.state,
+        COMPRESSED_STATE_ATTRIBUTES: state.attributes,
+        COMPRESSED_STATE_CONTEXT: context,
+    }
+    if state.last_changed == state.last_updated:
+        compressed_state[COMPRESSED_STATE_LAST_CHANGED] = state.last_changed.timestamp()
+    else:
+        compressed_state[COMPRESSED_STATE_LAST_CHANGED] = state.last_changed.timestamp()
+        compressed_state[COMPRESSED_STATE_LAST_UPDATED] = state.last_updated.timestamp()
+    return compressed_state
+
+
 def message_to_json(message: dict[str, Any]) -> str:
     """Serialize a websocket message to json."""
     try: