From 653da6e31f815af2df376ae6d03cab2793376145 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sun, 15 Oct 2023 11:39:09 -1000 Subject: [PATCH] Reduce websocket event and state JSON construction overhead (#101974) --- .../components/websocket_api/messages.py | 70 ++++++++++++------- .../components/websocket_api/test_messages.py | 2 +- 2 files changed, 44 insertions(+), 28 deletions(-) diff --git a/homeassistant/components/websocket_api/messages.py b/homeassistant/components/websocket_api/messages.py index 6e88c36c328..e1b038f4222 100644 --- a/homeassistant/components/websocket_api/messages.py +++ b/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" ) + ) diff --git a/tests/components/websocket_api/test_messages.py b/tests/components/websocket_api/test_messages.py index 6aafb9f2685..35ed55183d4 100644 --- a/tests/components/websocket_api/test_messages.py +++ b/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,