Reduce websocket event and state JSON construction overhead (#101974)

homeassistant/components/websocket_api/messages.py

@@ -32,9 +32,6 @@ MINIMAL_MESSAGE_SCHEMA: Final = vol.Schema(
 # Base schema to extend by message handlers
 BASE_COMMAND_MESSAGE_SCHEMA: Final = vol.Schema({vol.Required("id"): cv.positive_int})
 
-IDEN_TEMPLATE: Final = "__IDEN__"
-IDEN_JSON_TEMPLATE: Final = '"__IDEN__"'
-
 STATE_DIFF_ADDITIONS = "+"
 STATE_DIFF_REMOVALS = "-"
 
@@ -42,6 +39,21 @@ ENTITY_EVENT_ADD = "a"
 ENTITY_EVENT_REMOVE = "r"
 ENTITY_EVENT_CHANGE = "c"
 
+BASE_ERROR_MESSAGE = {
+    "type": const.TYPE_RESULT,
+    "success": False,
+}
+
+INVALID_JSON_PARTIAL_MESSAGE = JSON_DUMP(
+    {
+        **BASE_ERROR_MESSAGE,
+        "error": {
+            "code": const.ERR_UNKNOWN_ERROR,
+            "message": "Invalid JSON in response",
+        },
+    }
+)
+
 
 def result_message(iden: int, result: Any = None) -> dict[str, Any]:
     """Return a success result message."""
@@ -50,24 +62,21 @@ def result_message(iden: int, result: Any = None) -> dict[str, Any]:
 
 def construct_result_message(iden: int, payload: str) -> str:
     """Construct a success result message JSON."""
-    iden_str = str(iden)
-    return f'{{"id":{iden_str},"type":"result","success":true,"result":{payload}}}'
+    return f'{{"id":{iden},"type":"result","success":true,"result":{payload}}}'
 
 
 def error_message(iden: int | None, code: str, message: str) -> dict[str, Any]:
     """Return an error result message."""
     return {
         "id": iden,
-        "type": const.TYPE_RESULT,
-        "success": False,
+        **BASE_ERROR_MESSAGE,
         "error": {"code": code, "message": message},
     }
 
 
 def construct_event_message(iden: int, payload: str) -> str:
     """Construct an event message JSON."""
-    iden_str = str(iden)
-    return f'{{"id":{iden_str},"type":"event","event":{payload}}}'
+    return f'{{"id":{iden},"type":"event","event":{payload}}}'
 
 
 def event_message(iden: int, event: Any) -> dict[str, Any]:
@@ -84,18 +93,19 @@ def cached_event_message(iden: int, event: Event) -> str:
     all getting many of the same events (mostly state changed)
     we can avoid serializing the same data for each connection.
     """
-    return _cached_event_message(event).replace(IDEN_JSON_TEMPLATE, str(iden), 1)
+    return f'{_partial_cached_event_message(event)[:-1]},"id":{iden}}}'
 
 
 @lru_cache(maxsize=128)
-def _cached_event_message(event: Event) -> str:
+def _partial_cached_event_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
+    The message is constructed without the id which appended
+    in cached_event_message.
     """
-    return message_to_json(
-        {"id": IDEN_TEMPLATE, "type": "event", "event": event.as_dict()}
+    return (
+        _message_to_json_or_none({"type": "event", "event": event.as_dict()})
+        or INVALID_JSON_PARTIAL_MESSAGE
     )
 
 
@@ -108,18 +118,19 @@ def cached_state_diff_message(iden: int, event: Event) -> str:
     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)
+    return f'{_partial_cached_state_diff_message(event)[:-1]},"id":{iden}}}'
 
 
 @lru_cache(maxsize=128)
-def _cached_state_diff_message(event: Event) -> str:
+def _partial_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
+    The message is constructed without the id which
+    will be appended in cached_state_diff_message
     """
-    return message_to_json(
-        {"id": IDEN_TEMPLATE, "type": "event", "event": _state_diff_event(event)}
+    return (
+        _message_to_json_or_none({"type": "event", "event": _state_diff_event(event)})
+        or INVALID_JSON_PARTIAL_MESSAGE
     )
 
 
@@ -189,8 +200,8 @@ def _state_diff(
     return {ENTITY_EVENT_CHANGE: {new_state.entity_id: diff}}
 
 
-def message_to_json(message: dict[str, Any]) -> str:
-    """Serialize a websocket message to json."""
+def _message_to_json_or_none(message: dict[str, Any]) -> str | None:
+    """Serialize a websocket message to json or return None."""
     try:
         return JSON_DUMP(message)
     except (ValueError, TypeError):
@@ -200,8 +211,13 @@ def message_to_json(message: dict[str, Any]) -> str:
                 find_paths_unserializable_data(message, dump=JSON_DUMP)
             ),
         )
-        return JSON_DUMP(
-            error_message(
-                message["id"], const.ERR_UNKNOWN_ERROR, "Invalid JSON in response"
-            )
+    return None
+
+
+def message_to_json(message: dict[str, Any]) -> str:
+    """Serialize a websocket message to json or return an error."""
+    return _message_to_json_or_none(message) or JSON_DUMP(
+        error_message(
+            message["id"], const.ERR_UNKNOWN_ERROR, "Invalid JSON in response"
         )
+    )

tests/components/websocket_api/test_messages.py

@@ -2,7 +2,7 @@
 import pytest
 
 from homeassistant.components.websocket_api.messages import (
-    _cached_event_message as lru_event_cache,
+    _partial_cached_event_message as lru_event_cache,
     _state_diff_event,
     cached_event_message,
     message_to_json,