diff --git a/docs/versionhistory.rst b/docs/versionhistory.rst index 3acc15fb3..49229e263 100644 --- a/docs/versionhistory.rst +++ b/docs/versionhistory.rst @@ -46,6 +46,9 @@ This library adheres to `Semantic Versioning 2.0 `_. building a lookup table from the ``if TYPE_CHECKING:`` block. A fallback mode has been provided for installations where the source code is unavailable (e.g. PyInstaller). (`#1169 `_) +- Changed UDP sockets on the asyncio backend to make ``send()`` wait until the + datagram has been passed to the operating system + (`#1294 `_; PR by @graingert) - Fixed free-threading compatibility issues arising from the fact that on Python 3.14 free-threading builds, newly created threads inherit the current context by default, causing AnyIO to behave erroneously in relation to ``start_blocking_portal()`` and diff --git a/src/anyio/_backends/_asyncio.py b/src/anyio/_backends/_asyncio.py index 35b0cf371..cb3c57459 100644 --- a/src/anyio/_backends/_asyncio.py +++ b/src/anyio/_backends/_asyncio.py @@ -1299,6 +1299,7 @@ def connection_made(self, transport: asyncio.BaseTransport) -> None: self.write_event = asyncio.Event() self.closed_event = asyncio.Event() self.write_event.set() + cast(asyncio.WriteTransport, transport).set_write_buffer_limits(0) def connection_lost(self, exc: Exception | None) -> None: self.read_event.set() @@ -1710,14 +1711,26 @@ async def receive(self) -> tuple[bytes, IPSockAddrType]: async def send(self, item: UDPPacketType) -> None: with self._send_guard: - await AsyncIOBackend.checkpoint() - await self._protocol.write_event.wait() + await AsyncIOBackend.checkpoint_if_cancelled() + yielded = False + + # Wait out any datagram the transport had to buffer + if not self._protocol.write_event.is_set(): + yielded = True + await self._protocol.write_event.wait() + if self._closed: raise ClosedResourceError elif self._transport.is_closing(): raise BrokenResourceError - else: - self._transport.sendto(*item) + + self._transport.sendto(*item) + + # The high water mark is 0, so the event is clear if the OS refused it + if not self._protocol.write_event.is_set(): + await self._protocol.write_event.wait() + elif not yielded: + await AsyncIOBackend.cancel_shielded_checkpoint() class ConnectedUDPSocket(abc.ConnectedUDPSocket): @@ -1762,14 +1775,26 @@ async def receive(self) -> bytes: async def send(self, item: bytes) -> None: with self._send_guard: - await AsyncIOBackend.checkpoint() - await self._protocol.write_event.wait() + await AsyncIOBackend.checkpoint_if_cancelled() + yielded = False + + # Wait out any datagram the transport had to buffer + if not self._protocol.write_event.is_set(): + yielded = True + await self._protocol.write_event.wait() + if self._closed: raise ClosedResourceError elif self._transport.is_closing(): raise BrokenResourceError - else: - self._transport.sendto(item) + + self._transport.sendto(item) + + # The high water mark is 0, so the event is clear if the OS refused it + if not self._protocol.write_event.is_set(): + await self._protocol.write_event.wait() + elif not yielded: + await AsyncIOBackend.cancel_shielded_checkpoint() class UNIXDatagramSocket(_RawSocketMixin, abc.UNIXDatagramSocket): diff --git a/tests/test_sockets.py b/tests/test_sockets.py index cc661b79a..421b3cfea 100644 --- a/tests/test_sockets.py +++ b/tests/test_sockets.py @@ -117,6 +117,10 @@ not sys.platform.startswith("linux"), reason="Abstract namespace sockets is a Linux only feature", ) +skip_no_dgram_backpressure_mark = pytest.mark.skipif( + not sys.platform.startswith("linux"), + reason="Datagram sockets are only known to refuse datagrams with EAGAIN on Linux", +) @pytest.fixture @@ -1761,6 +1765,73 @@ async def handle(stream: SocketStream) -> None: assert client_addresses == expected_addresses +DGRAM_BACKPRESSURE_PAYLOAD = b"\x00" * 64 + + +def drain_datagram_socket(sock: socket.socket) -> None: + try: + while True: + sock.recv(65536) + except BlockingIOError: + pass + + +async def receive_datagrams_until(sock: socket.socket, sentinel: bytes) -> list[bytes]: + """Receive datagrams until ``sentinel`` arrives, and return all of them.""" + received: list[bytes] = [] + while not received or received[-1] != sentinel: + try: + received.append(sock.recv(65536)) + except BlockingIOError: + await wait_readable(sock) + + return received + + +def make_backpressured_pair( + tmp_path: Path, *, connect: bool +) -> tuple[socket.socket, socket.socket, str, int]: + """ + Create a datagram socket with a full send buffer, along with the peer it sends to. + + UNIX datagram sockets are used because UDP sockets never exert back-pressure on the + loopback interface, but they take the same code path in both backends. + + :return: the socket, its peer, the peer's path, and the number of datagrams the OS + accepts before it starts refusing them + + """ + peer_path = str(tmp_path / "peer.sock") + peer = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) + peer.setblocking(False) + peer.bind(peer_path) + sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) + sock.bind(str(tmp_path / "local.sock")) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 1024) + if connect: + sock.connect(peer_path) + + sock.setblocking(False) + + # Find out how many datagrams the OS accepts before it refuses any more, and then + # empty both buffers again + capacity = 0 + try: + while True: + if connect: + sock.send(DGRAM_BACKPRESSURE_PAYLOAD) + else: + sock.sendto(DGRAM_BACKPRESSURE_PAYLOAD, peer_path) + + capacity += 1 + except BlockingIOError: + pass + + assert capacity >= 2, f"the send buffer only fits {capacity} datagram(s)" + drain_datagram_socket(peer) + return sock, peer, peer_path, capacity + + @pytest.mark.network @pytest.mark.usefixtures("check_asyncio_bug") class TestUDPSocket: @@ -1832,6 +1903,71 @@ async def test_send_receive(self, family: AnyIPAddressFamily) -> None: assert response == b"halb" assert addr == (host, port) + @skip_no_dgram_backpressure_mark + async def test_send_waits_for_the_os_to_accept_the_datagram( + self, tmp_path: Path + ) -> None: + sock, peer, peer_path, capacity = make_backpressured_pair( + tmp_path, connect=False + ) + addr = cast(IPSockAddrType, peer_path) + with peer: + async with await UDPSocket.from_socket(sock) as udp: + # Fill up the send buffer again + for _ in range(capacity): + await udp.send((DGRAM_BACKPRESSURE_PAYLOAD, addr)) + + send_completed = False + + async def send_one_more() -> None: + nonlocal send_completed + await udp.send((DGRAM_BACKPRESSURE_PAYLOAD, addr)) + send_completed = True + + with fail_after(5): + async with create_task_group() as tg: + tg.start_soon(send_one_more) + await wait_all_tasks_blocked() + blocked = not send_completed + + # Make room in the send buffer so the sender can continue + drain_datagram_socket(peer) + + assert blocked, "send() returned before the OS accepted the datagram" + assert send_completed + + @skip_no_dgram_backpressure_mark + async def test_cancelled_send_does_not_buffer_the_next_datagram( + self, tmp_path: Path + ) -> None: + sock, peer, peer_path, capacity = make_backpressured_pair( + tmp_path, connect=False + ) + addr = cast(IPSockAddrType, peer_path) + with peer: + async with await UDPSocket.from_socket(sock) as udp: + # Fill up the send buffer again + for _ in range(capacity): + await udp.send((DGRAM_BACKPRESSURE_PAYLOAD, addr)) + + async def send_and_cancel(payload: bytes) -> None: + async with create_task_group() as tg: + tg.start_soon(udp.send, (payload, addr)) + await wait_all_tasks_blocked() + tg.cancel_scope.cancel() + + with fail_after(5): + # The first cancelled send may leave its datagram in the asyncio + # transport's buffer, but the second one must never reach it + await send_and_cancel(b"one") + await send_and_cancel(b"two") + + async with create_task_group() as tg: + tg.start_soon(udp.send, (b"three", addr)) + received = await receive_datagrams_until(peer, b"three") + + assert b"two" not in received + async def test_iterate(self, family: AnyIPAddressFamily) -> None: async def serve() -> None: async for packet, addr in server: @@ -2011,6 +2147,69 @@ async def test_send_receive(self, family: AnyIPAddressFamily) -> None: response = await udp2.receive() assert response == b"halb" + @skip_no_dgram_backpressure_mark + async def test_send_waits_for_the_os_to_accept_the_datagram( + self, tmp_path: Path + ) -> None: + sock, peer, _peer_path, capacity = make_backpressured_pair( + tmp_path, connect=True + ) + with peer: + async with await ConnectedUDPSocket.from_socket(sock) as udp: + # Fill up the send buffer again + for _ in range(capacity): + await udp.send(DGRAM_BACKPRESSURE_PAYLOAD) + + send_completed = False + + async def send_one_more() -> None: + nonlocal send_completed + await udp.send(DGRAM_BACKPRESSURE_PAYLOAD) + send_completed = True + + with fail_after(5): + async with create_task_group() as tg: + tg.start_soon(send_one_more) + await wait_all_tasks_blocked() + blocked = not send_completed + + # Make room in the send buffer so the sender can continue + drain_datagram_socket(peer) + + assert blocked, "send() returned before the OS accepted the datagram" + assert send_completed + + @skip_no_dgram_backpressure_mark + async def test_cancelled_send_does_not_buffer_the_next_datagram( + self, tmp_path: Path + ) -> None: + sock, peer, _peer_path, capacity = make_backpressured_pair( + tmp_path, connect=True + ) + with peer: + async with await ConnectedUDPSocket.from_socket(sock) as udp: + # Fill up the send buffer again + for _ in range(capacity): + await udp.send(DGRAM_BACKPRESSURE_PAYLOAD) + + async def send_and_cancel(payload: bytes) -> None: + async with create_task_group() as tg: + tg.start_soon(udp.send, payload) + await wait_all_tasks_blocked() + tg.cancel_scope.cancel() + + with fail_after(5): + # The first cancelled send may leave its datagram in the asyncio + # transport's buffer, but the second one must never reach it + await send_and_cancel(b"one") + await send_and_cancel(b"two") + + async with create_task_group() as tg: + tg.start_soon(udp.send, b"three") + received = await receive_datagrams_until(peer, b"three") + + assert b"two" not in received + async def test_iterate(self, family: AnyIPAddressFamily) -> None: async def serve() -> None: async for packet in udp2: