call set_write_buffer_limits(0) for UDPSockets so that send returns once data hits the kernel

docs/versionhistory.rst

@@ -46,6 +46,9 @@ This library adheres to `Semantic Versioning 2.0 <http://semver.org/>`_.
   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 <https://github.com/agronholm/anyio/pull/1169>`_)
+- Changed UDP sockets on the asyncio backend to make ``send()`` wait until the
+  datagram has been passed to the operating system
+  (`#1294 <https://github.com/agronholm/anyio/pull/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

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):

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: