From 45085dd97f670a1daf56a425be5a9330b0025975 Mon Sep 17 00:00:00 2001 From: Paulus Schoutsen Date: Tue, 14 May 2019 05:57:47 +0200 Subject: [PATCH] Better handle large amounts of data being sent over WS (#23842) * Better handle large amounts of data being sent over WS * Lint --- homeassistant/components/camera/__init__.py | 10 +++---- homeassistant/components/lovelace/__init__.py | 11 +++---- .../components/media_player/__init__.py | 4 +-- .../components/websocket_api/connection.py | 7 +++++ .../components/websocket_api/const.py | 5 ++++ .../components/websocket_api/http.py | 12 ++++---- .../websocket_api/test_connection.py | 30 +++++++++++++++++++ 7 files changed, 59 insertions(+), 20 deletions(-) create mode 100644 tests/components/websocket_api/test_connection.py diff --git a/homeassistant/components/camera/__init__.py b/homeassistant/components/camera/__init__.py index 7a37dffe3b8..7098d8bcb75 100644 --- a/homeassistant/components/camera/__init__.py +++ b/homeassistant/components/camera/__init__.py @@ -522,12 +522,10 @@ async def websocket_camera_thumbnail(hass, connection, msg): """ try: image = await async_get_image(hass, msg['entity_id']) - connection.send_message(websocket_api.result_message( - msg['id'], { - 'content_type': image.content_type, - 'content': base64.b64encode(image.content).decode('utf-8') - } - )) + await connection.send_big_result(msg['id'], { + 'content_type': image.content_type, + 'content': base64.b64encode(image.content).decode('utf-8') + }) except HomeAssistantError: connection.send_message(websocket_api.error_message( msg['id'], 'image_fetch_failed', 'Unable to fetch image')) diff --git a/homeassistant/components/lovelace/__init__.py b/homeassistant/components/lovelace/__init__.py index 03b1cf06d68..996e3f7b296 100644 --- a/homeassistant/components/lovelace/__init__.py +++ b/homeassistant/components/lovelace/__init__.py @@ -176,18 +176,19 @@ def handle_yaml_errors(func): error = None try: result = await func(hass, connection, msg) - message = websocket_api.result_message( - msg['id'], result - ) except ConfigNotFound: error = 'config_not_found', 'No config found.' except HomeAssistantError as err: error = 'error', str(err) if error is not None: - message = websocket_api.error_message(msg['id'], *error) + connection.send_error(msg['id'], *error) + return - connection.send_message(message) + if msg is not None: + await connection.send_big_result(msg['id'], result) + else: + connection.send_result(msg['id'], result) return send_with_error_handling diff --git a/homeassistant/components/media_player/__init__.py b/homeassistant/components/media_player/__init__.py index ccfa968fa9a..b433a90f329 100644 --- a/homeassistant/components/media_player/__init__.py +++ b/homeassistant/components/media_player/__init__.py @@ -869,8 +869,8 @@ async def websocket_handle_thumbnail(hass, connection, msg): 'Failed to fetch thumbnail')) return - connection.send_message(websocket_api.result_message( + await connection.send_big_result( msg['id'], { 'content_type': content_type, 'content': base64.b64encode(data).decode('utf-8') - })) + }) diff --git a/homeassistant/components/websocket_api/connection.py b/homeassistant/components/websocket_api/connection.py index c09e8c4c6e2..1aa1efc0eca 100644 --- a/homeassistant/components/websocket_api/connection.py +++ b/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.""" diff --git a/homeassistant/components/websocket_api/const.py b/homeassistant/components/websocket_api/const.py index 53ca680c4c9..9c776e3b949 100644 --- a/homeassistant/components/websocket_api/const.py +++ b/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) diff --git a/homeassistant/components/websocket_api/http.py b/homeassistant/components/websocket_api/http.py index 85051dcae73..80592cc7151 100644 --- a/homeassistant/components/websocket_api/http.py +++ b/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) diff --git a/tests/components/websocket_api/test_connection.py b/tests/components/websocket_api/test_connection.py new file mode 100644 index 00000000000..eeac9af24cd --- /dev/null +++ b/tests/components/websocket_api/test_connection.py @@ -0,0 +1,30 @@ +"""Test WebSocket Connection class.""" +from homeassistant.components import websocket_api +from homeassistant.components.websocket_api import const + + +async def test_send_big_result(hass, websocket_client): + """Test sending big results over the WS.""" + @websocket_api.websocket_command({ + 'type': 'big_result' + }) + @websocket_api.async_response + async def send_big_result(hass, connection, msg): + await connection.send_big_result( + msg['id'], {'big': 'result'} + ) + + hass.components.websocket_api.async_register_command( + send_big_result + ) + + await websocket_client.send_json({ + 'id': 5, + 'type': 'big_result', + }) + + msg = await websocket_client.receive_json() + assert msg['id'] == 5 + assert msg['type'] == const.TYPE_RESULT + assert msg['success'] + assert msg['result'] == {'big': 'result'}