Skip to content

Commit

Permalink
TrioBackend: Fixed UDP listener crashing down application on `recvfro…
Browse files Browse the repository at this point in the history
…m()` errors (#371)
  • Loading branch information
francis-clairicia authored Nov 3, 2024
1 parent cd2e150 commit fff0a3e
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 7 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
__all__ = ["TrioDatagramListenerSocketAdapter"]

import contextlib
import logging
import socket as _socket
import warnings
from collections.abc import Awaitable, Callable, Coroutine, Mapping
Expand Down Expand Up @@ -102,14 +103,20 @@ async def serve(
MAX_DATAGRAM_BUFSIZE = self.MAX_DATAGRAM_BUFSIZE
listener = self.__listener
wait_readable = self.__wait_readable
logger = logging.getLogger(__name__)

while True:
datagram, client_address = await _retry_socket_method(
wait_readable,
listener,
lambda: listener.recvfrom(MAX_DATAGRAM_BUFSIZE),
always_yield=True,
)
try:
datagram, client_address = await _retry_socket_method(
wait_readable,
listener,
lambda: listener.recvfrom(MAX_DATAGRAM_BUFSIZE),
always_yield=True,
)
except OSError as exc:
message = "Unrelated error occurred on datagram reception: %s: %s"
logger.warning(message, type(exc).__name__, exc, exc_info=exc)
continue

task_group.start_soon(handler, datagram, client_address)
# Always drop references on loop end
Expand Down
13 changes: 12 additions & 1 deletion tests/unit_test/test_async/test_trio_backend/test_datagram.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

import contextlib
import logging
from collections.abc import AsyncIterator, Callable
from typing import TYPE_CHECKING, Any

Expand Down Expand Up @@ -295,8 +296,10 @@ async def test____serve____default(
mock_udp_listener_socket: MagicMock,
mock_trio_lowlevel_wait_readable: AsyncMock,
mocker: MockerFixture,
caplog: pytest.LogCaptureFixture,
) -> None:
# Arrange
caplog.set_level("INFO", listener.__class__.__module__)

import trio

Expand All @@ -305,6 +308,9 @@ async def test____serve____default(
(b"received_datagram", ("127.0.0.1", 12345)),
(b"received_datagram_2", ("127.0.0.1", 54321)),
BlockingIOError,
OSError("Unrelated OS Error"),
(b"received_datagram_3", ("127.0.0.1", 11111)),
BlockingIOError,
(await self._get_cancelled_exc()),
]

Expand All @@ -318,8 +324,13 @@ async def test____serve____default(
assert handler.await_args_list == [
mocker.call(b"received_datagram", ("127.0.0.1", 12345)),
mocker.call(b"received_datagram_2", ("127.0.0.1", 54321)),
mocker.call(b"received_datagram_3", ("127.0.0.1", 11111)),
]
assert mock_trio_lowlevel_wait_readable.await_args_list == [mocker.call(mock_udp_listener_socket) for _ in range(2)]
assert mock_trio_lowlevel_wait_readable.await_args_list == [mocker.call(mock_udp_listener_socket) for _ in range(3)]
assert len(caplog.records) == 1
assert caplog.records[0].levelno == logging.WARNING
assert caplog.records[0].getMessage() == "Unrelated error occurred on datagram reception: OSError: Unrelated OS Error"
assert caplog.records[0].exc_info is not None and isinstance(caplog.records[0].exc_info[1], OSError)

@pytest.mark.parametrize("block_count", [2, 1, 0], ids=lambda count: f"block_count=={count}")
async def test____send_to____write_on_socket(
Expand Down

0 comments on commit fff0a3e

Please sign in to comment.