Better handle large amounts of data being sent over WS (#23842)

* Better handle large amounts of data being sent over WS

* Lint

homeassistant/components/websocket_api/connection.py

@@ -36,6 +36,13 @@ class ActiveConnection:
         """Send a result message."""
         self.send_message(messages.result_message(msg_id, result))
 
+    async def send_big_result(self, msg_id, result):
+        """Send a result message that would be expensive to JSON serialize."""
+        content = await self.hass.async_add_executor_job(
+            const.JSON_DUMP, messages.result_message(msg_id, result)
+        )
+        self.send_message(content)
+
     @callback
     def send_error(self, msg_id, code, message):
         """Send a error message."""

homeassistant/components/websocket_api/const.py

@@ -1,6 +1,9 @@
 """Websocket constants."""
 import asyncio
 from concurrent import futures
+from functools import partial
+import json
+from homeassistant.helpers.json import JSONEncoder
 
 DOMAIN = 'websocket_api'
 URL = '/api/websocket'
@@ -27,3 +30,5 @@ SIGNAL_WEBSOCKET_DISCONNECTED = 'websocket_disconnected'
 
 # Data used to store the current connection list
 DATA_CONNECTIONS = DOMAIN + '.connections'
+
+JSON_DUMP = partial(json.dumps, cls=JSONEncoder, allow_nan=False)

homeassistant/components/websocket_api/http.py

@@ -1,8 +1,6 @@
 """View to accept incoming websocket connection."""
 import asyncio
 from contextlib import suppress
-from functools import partial
-import json
 import logging
 
 from aiohttp import web, WSMsgType
@@ -11,18 +9,15 @@ import async_timeout
 from homeassistant.const import EVENT_HOMEASSISTANT_STOP
 from homeassistant.core import callback
 from homeassistant.components.http import HomeAssistantView
-from homeassistant.helpers.json import JSONEncoder
 
 from .const import (
     MAX_PENDING_MSG, CANCELLATION_ERRORS, URL, ERR_UNKNOWN_ERROR,
     SIGNAL_WEBSOCKET_CONNECTED, SIGNAL_WEBSOCKET_DISCONNECTED,
-    DATA_CONNECTIONS)
+    DATA_CONNECTIONS, JSON_DUMP)
 from .auth import AuthPhase, auth_required_message
 from .error import Disconnect
 from .messages import error_message
 
-JSON_DUMP = partial(json.dumps, cls=JSONEncoder, allow_nan=False)
-
 
 class WebsocketAPIView(HomeAssistantView):
     """View to serve a websockets endpoint."""
@@ -62,7 +57,10 @@ class WebSocketHandler:
                     break
                 self._logger.debug("Sending %s", message)
                 try:
-                    await self.wsock.send_json(message, dumps=JSON_DUMP)
+                    if isinstance(message, str):
+                        await self.wsock.send_str(message)
+                    else:
+                        await self.wsock.send_json(message, dumps=JSON_DUMP)
                 except (ValueError, TypeError) as err:
                     self._logger.error('Unable to serialize to JSON: %s\n%s',
                                        err, message)