Improve websocket throughput and reduce latency

homeassistant/components/websocket_api/commands.py

@@ -715,7 +715,7 @@ def handle_supported_features(
     hass: HomeAssistant, connection: ActiveConnection, msg: dict[str, Any]
 ) -> None:
     """Handle setting supported features."""
-    connection.supported_features = msg["features"]
+    connection.set_supported_features(msg["features"])
     connection.send_result(msg["id"])
 
 

homeassistant/components/websocket_api/connection.py

@@ -48,13 +48,19 @@ def __init__(
         self.refresh_token_id = refresh_token.id
         self.subscriptions: dict[Hashable, Callable[[], Any]] = {}
         self.last_id = 0
+        self.can_coalesce = False
         self.supported_features: dict[str, float] = {}
         self.handlers: dict[str, tuple[MessageHandler, vol.Schema]] = self.hass.data[
             const.DOMAIN
         ]
         self.binary_handlers: list[BinaryHandler | None] = []
         current_connection.set(self)
 
+    def set_supported_features(self, features: dict[str, float]) -> None:
+        """Set supported features."""
+        self.supported_features = features
+        self.can_coalesce = const.FEATURE_COALESCE_MESSAGES in features
+
     def get_description(self, request: web.Request | None) -> str:
         """Return a description of the connection."""
         description = self.user.name or ""

homeassistant/components/websocket_api/http.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import asyncio
+from collections import deque
 from collections.abc import Callable
 from contextlib import suppress
 import datetime as dt
@@ -22,7 +23,6 @@
 from .const import (
     CANCELLATION_ERRORS,
     DATA_CONNECTIONS,
-    FEATURE_COALESCE_MESSAGES,
     MAX_PENDING_MSG,
     PENDING_MSG_PEAK,
     PENDING_MSG_PEAK_TIME,
@@ -71,7 +71,8 @@ def __init__(self, hass: HomeAssistant, request: web.Request) -> None:
         self.hass = hass
         self.request = request
         self.wsock = web.WebSocketResponse(heartbeat=55)
-        self._to_write: asyncio.Queue = asyncio.Queue(maxsize=MAX_PENDING_MSG)
+        self._message_queue: deque = deque()
+        self._queue_ready: asyncio.Event = asyncio.Event()
         self._handle_task: asyncio.Task | None = None
         self._writer_task: asyncio.Task | None = None
         self._closing: bool = False
@@ -89,33 +90,40 @@ def description(self) -> str:
     async def _writer(self) -> None:
         """Write outgoing messages."""
         # Exceptions if Socket disconnected or cancelled by connection handler
-        to_write = self._to_write
+        message_queue = self._message_queue
+        queue_ready = self._queue_ready
         logger = self._logger
         wsock = self.wsock
         try:
             with suppress(RuntimeError, ConnectionResetError, *CANCELLATION_ERRORS):
                 while not self.wsock.closed:
-                    if (process := await to_write.get()) is None:
+                    await queue_ready.wait()
+                    if (process := message_queue.popleft()) is None:
                         return
+
                     message = process if isinstance(process, str) else process()
+                    no_remaining_messages = len(message_queue) == 0
+
                     if (
-                        to_write.empty()
+                        no_remaining_messages
                         or not self.connection
-                        or FEATURE_COALESCE_MESSAGES
-                        not in self.connection.supported_features
+                        or not self.connection.can_coalesce
                     ):
+                        if no_remaining_messages:
+                            queue_ready.clear()
                         logger.debug("Sending %s", message)
                         await wsock.send_str(message)
                         continue
 
                     messages: list[str] = [message]
-                    while not to_write.empty():
-                        if (process := to_write.get_nowait()) is None:
+                    while len(message_queue):
+                        if (process := message_queue.popleft()) is None:
                             return
                         messages.append(
                             process if isinstance(process, str) else process()
                         )
 
+                    queue_ready.clear()
                     coalesced_messages = "[" + ",".join(messages) + "]"
                     logger.debug("Sending %s", coalesced_messages)
                     await wsock.send_str(coalesced_messages)
@@ -146,11 +154,9 @@ def _send_message(self, message: str | dict[str, Any] | Callable[[], str]) -> No
         if isinstance(message, dict):
             message = message_to_json(message)
 
-        to_write = self._to_write
-
-        try:
-            to_write.put_nowait(message)
-        except asyncio.QueueFull:
+        message_queue = self._message_queue
+        queue_size_before_add = len(message_queue)
+        if queue_size_before_add >= MAX_PENDING_MSG:
             self._logger.error(
                 (
                     "%s: Client unable to keep up with pending messages. Reached %s pending"
@@ -163,9 +169,12 @@ def _send_message(self, message: str | dict[str, Any] | Callable[[], str]) -> No
             )
             self._cancel()
 
+        message_queue.append(message)
+        self._queue_ready.set()
+
         peak_checker_active = self._peak_checker_unsub is not None
 
-        if to_write.qsize() < PENDING_MSG_PEAK:
+        if queue_size_before_add <= PENDING_MSG_PEAK:
             if peak_checker_active:
                 self._cancel_peak_checker()
             return
@@ -180,7 +189,7 @@ def _check_write_peak(self, _utc_time: dt.datetime) -> None:
         """Check that we are no longer above the write peak."""
         self._peak_checker_unsub = None
 
-        if self._to_write.qsize() < PENDING_MSG_PEAK:
+        if len(self._message_queue) < PENDING_MSG_PEAK:
             return
 
         self._logger.error(
@@ -357,7 +366,8 @@ def handle_hass_stop(event: Event) -> None:
             self._closing = True
 
             try:
-                self._to_write.put_nowait(None)
+                self._message_queue.append(None)
+                self._queue_ready.set()
                 # Make sure all error messages are written before closing
                 await self._writer_task
                 await wsock.close()