Make failure to increase the recv buffer size non-fatal

We now try to set the buffer to
2097152, 1048576, 524288, 262144, 131072

before giving up and logging a warning
This commit is contained in:
J. Nick Koston 2023-12-19 09:48:07 -10:00
parent 6565d67c73
commit 8871ec0668
No known key found for this signature in database
3 changed files with 124 additions and 5 deletions

View File

@ -152,3 +152,5 @@ cdef class APIConnection:
cdef void _set_fatal_exception_if_unset(self, Exception err)
cdef void _register_internal_message_handlers(self)
cdef void _increase_recv_buffer_size(self)

View File

@ -62,7 +62,8 @@ else:
_LOGGER = logging.getLogger(__name__)
BUFFER_SIZE = 1024 * 1024 * 2 # Set buffer limit to 2MB
PREFERRED_BUFFER_SIZE = 2097152 # Set buffer limit to 2MB
MIN_BUFFER_SIZE = 131072 # Minimum buffer size to use
DISCONNECT_REQUEST_MESSAGE = DisconnectRequest()
DISCONNECT_RESPONSE_MESSAGES = (DisconnectResponse(),)
@ -384,9 +385,7 @@ class APIConnection:
self._socket = sock
sock.setblocking(False)
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
# Try to reduce the pressure on esphome device as it measures
# ram in bytes and we measure ram in megabytes.
sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, BUFFER_SIZE)
self._increase_recv_buffer_size()
self.connected_address = sock.getpeername()[0]
if self._debug_enabled:
@ -397,6 +396,32 @@ class APIConnection:
self._params.port,
)
def _increase_recv_buffer_size(self) -> None:
"""Increase the recv buffer size."""
if TYPE_CHECKING:
assert self._socket is not None
new_buffer_size = PREFERRED_BUFFER_SIZE
while True:
# Try to reduce the pressure on ESPHome device as it measures
# ram in bytes and we measure ram in megabytes.
try:
self._socket.setsockopt(
socket.SOL_SOCKET, socket.SO_RCVBUF, new_buffer_size
)
return
except OSError as err:
if new_buffer_size <= MIN_BUFFER_SIZE:
_LOGGER.warning(
"%s: Unable to increase the socket receive buffer size to %s; "
"The connection may not be unstable if the ESPHome device "
"sends at volume (ex. a Bluetooth proxy or camera): %s",
self.log_name,
new_buffer_size,
err,
)
return
new_buffer_size //= 2
async def _connect_init_frame_helper(self) -> None:
"""Step 3 in connect process: initialize the frame helper and init read loop."""
fh: APIPlaintextFrameHelper | APINoiseFrameHelper

View File

@ -2,10 +2,11 @@ from __future__ import annotations
import asyncio
import logging
import socket
from datetime import timedelta
from functools import partial
from typing import Callable, cast
from unittest.mock import AsyncMock, MagicMock, call, patch
from unittest.mock import AsyncMock, MagicMock, call, create_autospec, patch
import pytest
from google.protobuf import message
@ -236,6 +237,97 @@ async def test_start_connection_socket_error(
await asyncio.sleep(0)
@pytest.mark.asyncio
async def test_start_connection_cannot_increase_recv_buffer(
conn: APIConnection,
resolve_host,
aiohappyeyeballs_start_connection: MagicMock,
caplog: pytest.LogCaptureFixture,
):
"""Test failing to increase the recv buffer."""
loop = asyncio.get_event_loop()
transport = MagicMock()
connected = asyncio.Event()
tried_sizes = []
def _setsockopt(*args, **kwargs):
if args[0] == socket.SOL_SOCKET and args[1] == socket.SO_RCVBUF:
size = args[2]
tried_sizes.append(size)
raise OSError("Socket error")
mock_socket: socket.socket = create_autospec(
socket.socket, spec_set=True, instance=True, name="bad_buffer_socket"
)
mock_socket.type = socket.SOCK_STREAM
mock_socket.fileno.return_value = 1
mock_socket.getpeername.return_value = ("10.0.0.512", 323)
mock_socket.setsockopt = _setsockopt
mock_socket.sendmsg.side_effect = OSError("Socket error")
aiohappyeyeballs_start_connection.return_value = mock_socket
with patch.object(
loop,
"create_connection",
side_effect=partial(_create_mock_transport_protocol, transport, connected),
):
connect_task = asyncio.create_task(connect(conn, login=False))
await asyncio.sleep(0)
await connected.wait()
protocol = conn._frame_helper
send_plaintext_hello(protocol)
await connect_task
assert "Unable to increase the socket receive buffer size to 131072" in caplog.text
assert tried_sizes == [2097152, 1048576, 524288, 262144, 131072]
@pytest.mark.asyncio
async def test_start_connection_can_only_increase_buffer_size_to_262144(
conn: APIConnection,
resolve_host,
aiohappyeyeballs_start_connection: MagicMock,
caplog: pytest.LogCaptureFixture,
):
"""Test the receive buffer can only be increased to 262144."""
loop = asyncio.get_event_loop()
transport = MagicMock()
connected = asyncio.Event()
tried_sizes = []
def _setsockopt(*args, **kwargs):
if args[0] == socket.SOL_SOCKET and args[1] == socket.SO_RCVBUF:
size = args[2]
tried_sizes.append(size)
if size != 262144:
raise OSError("Socket error")
mock_socket: socket.socket = create_autospec(
socket.socket, spec_set=True, instance=True, name="bad_buffer_socket"
)
mock_socket.type = socket.SOCK_STREAM
mock_socket.fileno.return_value = 1
mock_socket.getpeername.return_value = ("10.0.0.512", 323)
mock_socket.setsockopt = _setsockopt
mock_socket.sendmsg.side_effect = OSError("Socket error")
aiohappyeyeballs_start_connection.return_value = mock_socket
with patch.object(
loop,
"create_connection",
side_effect=partial(_create_mock_transport_protocol, transport, connected),
):
connect_task = asyncio.create_task(connect(conn, login=False))
await asyncio.sleep(0)
await connected.wait()
protocol = conn._frame_helper
send_plaintext_hello(protocol)
await connect_task
assert "Unable to increase the socket receive buffer size" not in caplog.text
assert tried_sizes == [2097152, 1048576, 524288, 262144]
@pytest.mark.asyncio
async def test_start_connection_times_out(
conn: APIConnection, resolve_host, aiohappyeyeballs_start_connection