Experiment: simplify the read side interface

Cargo.lock

@@ -1510,10 +1510,12 @@ dependencies = [
 name = "iroh-quinn-udp"
 version = "0.8.0"
 dependencies = [
+ "bytes",
  "cfg_aliases",
  "criterion",
  "libc",
  "log",
+ "smallvec",
  "socket2",
  "tokio",
  "tracing",

Cargo.toml

@@ -44,6 +44,7 @@ rustls-pki-types = "1.7"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1"
 slab = "0.4.9"
+smallvec = "1"
 smol = "2"
 socket2 = ">=0.5, <0.7"
 sorted-index-buffer = { version = "0.2.0" }

quinn-udp/Cargo.toml

@@ -21,8 +21,10 @@ log = ["dep:log"]
 fast-apple-datapath = []
 
 [dependencies]
+bytes = { workspace = true }
 libc = "0.2.175"
 log = { workspace = true, optional = true }
+smallvec = { workspace = true }
 tracing = { workspace = true, optional = true }
 
 [target.'cfg(not(all(target_family = "wasm", target_os = "unknown")))'.dependencies]

quinn-udp/src/fallback.rs

@@ -1,10 +1,16 @@
 use std::{
     io::{self, IoSliceMut},
+    num::NonZeroUsize,
     sync::Mutex,
     time::Instant,
 };
 
-use super::{IO_ERROR_LOG_INTERVAL, RecvMeta, Transmit, UdpSockRef, log_sendmsg_error};
+use bytes::BytesMut;
+
+use super::{
+    IO_ERROR_LOG_INTERVAL, RecvMeta, ReceivedDatagram, ReceivedDatagrams, Transmit, UdpSockRef,
+    log_sendmsg_error,
+};
 
 /// Fallback UDP socket interface that stubs out all special functionality
 ///
@@ -77,6 +83,36 @@ impl UdpSocketState {
         Ok(1)
     }
 
+    /// Receives datagrams from the socket, returning owned data.
+    ///
+    /// This is a higher-level API that handles buffer management internally.
+    /// Each datagram in the returned collection contains its own `BytesMut`
+    /// buffer suitable for in-place decryption.
+    pub fn recv_datagrams(
+        &self,
+        socket: UdpSockRef<'_>,
+        max_payload_size: usize,
+    ) -> io::Result<ReceivedDatagrams> {
+        let mut recv_buf = vec![0u8; max_payload_size];
+        let mut bufs = [IoSliceMut::new(&mut recv_buf)];
+        let mut metas = [RecvMeta::default()];
+
+        let msg_count = self.recv(socket, &mut bufs, &mut metas)?;
+
+        let mut result = ReceivedDatagrams::new();
+        for meta in metas.iter().take(msg_count) {
+            let data = BytesMut::from(&recv_buf[..meta.len]);
+            result.push(ReceivedDatagram {
+                data,
+                remote: meta.addr,
+                local_ip: meta.dst_ip,
+                ecn: meta.ecn,
+            });
+        }
+
+        Ok(result)
+    }
+
     #[inline]
     pub fn max_gso_segments(&self) -> usize {
         1

quinn-udp/src/lib.rs

@@ -28,6 +28,9 @@
 #![warn(clippy::use_self)]
 
 use std::net::{IpAddr, Ipv6Addr, SocketAddr};
+
+use bytes::BytesMut;
+use smallvec::SmallVec;
 #[cfg(unix)]
 use std::os::unix::io::AsFd;
 #[cfg(windows)]
@@ -134,6 +137,77 @@ impl Default for RecvMeta {
     }
 }
 
+/// A single received UDP datagram
+#[derive(Debug)]
+pub struct ReceivedDatagram {
+    /// The payload of the datagram
+    pub data: BytesMut,
+    /// The source address of the datagram
+    pub remote: SocketAddr,
+    /// The destination IP address the datagram was sent to
+    pub local_ip: Option<IpAddr>,
+    /// The Explicit Congestion Notification bits
+    pub ecn: Option<EcnCodepoint>,
+}
+
+/// Maximum number of datagrams to store inline without heap allocation
+///
+/// This is set to accommodate a single GRO batch (up to 64 segments on Linux).
+const DATAGRAM_VEC_INLINE_CAP: usize = 64;
+
+/// A collection of received datagrams
+///
+/// This type uses inline storage for small batches to avoid heap allocation
+/// in the common case. It implements [`IntoIterator`] for convenient consumption.
+#[derive(Debug)]
+pub struct ReceivedDatagrams {
+    inner: SmallVec<[ReceivedDatagram; DATAGRAM_VEC_INLINE_CAP]>,
+}
+
+impl ReceivedDatagrams {
+    /// Creates an empty collection
+    pub fn new() -> Self {
+        Self {
+            inner: SmallVec::new(),
+        }
+    }
+
+    /// Adds a datagram to the collection
+    pub fn push(&mut self, datagram: ReceivedDatagram) {
+        self.inner.push(datagram);
+    }
+
+    /// Returns the number of datagrams in the collection
+    pub fn len(&self) -> usize {
+        self.inner.len()
+    }
+
+    /// Returns `true` if the collection contains no datagrams
+    pub fn is_empty(&self) -> bool {
+        self.inner.is_empty()
+    }
+
+    /// Returns an iterator over the datagrams
+    pub fn iter(&self) -> impl Iterator<Item = &ReceivedDatagram> {
+        self.inner.iter()
+    }
+}
+
+impl Default for ReceivedDatagrams {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl IntoIterator for ReceivedDatagrams {
+    type Item = ReceivedDatagram;
+    type IntoIter = smallvec::IntoIter<[ReceivedDatagram; DATAGRAM_VEC_INLINE_CAP]>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.inner.into_iter()
+    }
+}
+
 /// An outgoing packet
 #[derive(Debug, Clone)]
 pub struct Transmit<'a> {

quinn-udp/src/unix.rs

@@ -15,8 +15,11 @@ use std::{
 
 use socket2::SockRef;
 
+use bytes::BytesMut;
+
 use super::{
-    EcnCodepoint, IO_ERROR_LOG_INTERVAL, RecvMeta, Transmit, UdpSockRef, cmsg, log_sendmsg_error,
+    EcnCodepoint, IO_ERROR_LOG_INTERVAL, RecvMeta, ReceivedDatagram, ReceivedDatagrams, Transmit,
+    UdpSockRef, cmsg, log_sendmsg_error,
 };
 
 // Adapted from https://github.com/apple-oss-distributions/xnu/blob/8d741a5de7ff4191bf97d57b9f54c2f6d4a15585/bsd/sys/socket_private.h
@@ -234,6 +237,62 @@ impl UdpSocketState {
         recv(socket.0, bufs, meta)
     }
 
+    /// Receives datagrams from the socket, returning owned data.
+    ///
+    /// This is a higher-level API that handles buffer management and GRO splitting
+    /// internally. Each datagram in the returned collection contains its own `BytesMut`
+    /// buffer suitable for in-place decryption.
+    ///
+    /// # Arguments
+    ///
+    /// * `socket` - The UDP socket to receive from
+    /// * `max_payload_size` - Maximum expected UDP payload size (typically 65535 or less)
+    ///
+    /// # Returns
+    ///
+    /// A collection of received datagrams, or an error if the receive failed.
+    pub fn recv_datagrams(
+        &self,
+        socket: UdpSockRef<'_>,
+        max_payload_size: usize,
+    ) -> io::Result<ReceivedDatagrams> {
+        // Allocate buffer sized for GRO coalescing
+        let gro_segments = self.gro_segments.get();
+        let buf_size = max_payload_size * gro_segments;
+        let mut recv_buf = vec![0u8; buf_size * BATCH_SIZE];
+
+        // Prepare IoSliceMut array for recv
+        let mut bufs: [IoSliceMut<'_>; BATCH_SIZE] =
+            std::array::from_fn(|_| IoSliceMut::new(&mut []));
+        for (i, chunk) in recv_buf.chunks_mut(buf_size).enumerate().take(BATCH_SIZE) {
+            bufs[i] = IoSliceMut::new(chunk);
+        }
+
+        let mut metas = [RecvMeta::default(); BATCH_SIZE];
+
+        // Call the underlying recv
+        let msg_count = recv(socket.0, &mut bufs, &mut metas)?;
+
+        // Convert to ReceivedDatagrams, splitting by stride
+        let mut result = ReceivedDatagrams::new();
+        for (meta, buf) in metas.iter().zip(bufs.iter()).take(msg_count) {
+            let mut offset = 0;
+            while offset < meta.len {
+                let stride = meta.stride.min(meta.len - offset);
+                let data = BytesMut::from(&buf[offset..offset + stride]);
+                result.push(ReceivedDatagram {
+                    data,
+                    remote: meta.addr,
+                    local_ip: meta.dst_ip,
+                    ecn: meta.ecn,
+                });
+                offset += stride;
+            }
+        }
+
+        Ok(result)
+    }
+
     /// The maximum amount of segments which can be transmitted if a platform
     /// supports Generic Send Offload (GSO).
     ///

quinn-udp/src/windows.rs

@@ -15,8 +15,11 @@ use std::{
 use libc::{c_int, c_uint};
 use windows_sys::Win32::Networking::WinSock;
 
+use bytes::BytesMut;
+
 use crate::{
-    EcnCodepoint, IO_ERROR_LOG_INTERVAL, RecvMeta, Transmit, UdpSockRef,
+    EcnCodepoint, IO_ERROR_LOG_INTERVAL, RecvMeta, ReceivedDatagram, ReceivedDatagrams, Transmit,
+    UdpSockRef,
     cmsg::{self, CMsgHdr},
     log::debug,
     log_sendmsg_error,
@@ -281,6 +284,47 @@ impl UdpSocketState {
         Ok(1)
     }
 
+    /// Receives datagrams from the socket, returning owned data.
+    ///
+    /// This is a higher-level API that handles buffer management and GRO splitting
+    /// internally. Each datagram in the returned collection contains its own `BytesMut`
+    /// buffer suitable for in-place decryption.
+    pub fn recv_datagrams(
+        &self,
+        socket: UdpSockRef<'_>,
+        max_payload_size: usize,
+    ) -> io::Result<ReceivedDatagrams> {
+        // Allocate buffer sized for URO coalescing
+        let gro_segments = self.gro_segments().get();
+        let buf_size = max_payload_size * gro_segments;
+        let mut recv_buf = vec![0u8; buf_size];
+
+        let mut bufs = [IoSliceMut::new(&mut recv_buf)];
+        let mut metas = [RecvMeta::default()];
+
+        // Call the underlying recv
+        let msg_count = self.recv(socket, &mut bufs, &mut metas)?;
+
+        // Convert to ReceivedDatagrams, splitting by stride
+        let mut result = ReceivedDatagrams::new();
+        for meta in metas.iter().take(msg_count) {
+            let mut offset = 0;
+            while offset < meta.len {
+                let stride = meta.stride.min(meta.len - offset);
+                let data = BytesMut::from(&recv_buf[offset..offset + stride]);
+                result.push(ReceivedDatagram {
+                    data,
+                    remote: meta.addr,
+                    local_ip: meta.dst_ip,
+                    ecn: meta.ecn,
+                });
+                offset += stride;
+            }
+        }
+
+        Ok(result)
+    }
+
     /// The maximum amount of segments which can be transmitted if a platform
     /// supports Generic Send Offload (GSO).
     ///

quinn/src/endpoint.rs

@@ -2,10 +2,9 @@ use std::{
     collections::VecDeque,
     fmt,
     future::Future,
-    io::{self, IoSliceMut},
+    io,
     mem,
     net::{SocketAddr, SocketAddrV6},
-    num::NonZeroUsize,
     pin::Pin,
     str,
     sync::{Arc, Mutex},
@@ -23,7 +22,7 @@ use crate::{
     runtime::{AsyncUdpSocket, Runtime, UdpSender},
     udp_transmit,
 };
-use bytes::{Bytes, BytesMut};
+use bytes::Bytes;
 use pin_project_lite::pin_project;
 use proto::{
     self as proto, ClientConfig, ConnectError, ConnectionError, ConnectionHandle, DatagramEvent,
@@ -38,7 +37,6 @@ use rustc_hash::FxHashMap;
 use socket2::{Domain, Protocol, Socket, Type};
 use tokio::sync::{Notify, futures::Notified, mpsc};
 use tracing::{Instrument, Span};
-use udp::{BATCH_SIZE, RecvMeta};
 
 use crate::{
     ConnectionEvent, EndpointConfig, IO_LOOP_BOUND, RECV_TIME_BOUND, VarInt,
@@ -731,7 +729,7 @@ impl EndpointRef {
         runtime: Arc<dyn Runtime>,
     ) -> Self {
         let (sender, events) = mpsc::unbounded_channel();
-        let recv_state = RecvState::new(sender, socket.max_receive_segments(), &inner);
+        let recv_state = RecvState::new(sender);
         let sender = socket.create_sender();
         Self(Arc::new(EndpointInner {
             shared: Shared {
@@ -791,22 +789,13 @@ impl std::ops::Deref for EndpointRef {
 struct RecvState {
     incoming: VecDeque<proto::Incoming>,
     connections: ConnectionSet,
-    recv_buf: Box<[u8]>,
     recv_limiter: WorkLimiter,
 }
 
 impl RecvState {
     fn new(
         sender: mpsc::UnboundedSender<(ConnectionHandle, EndpointEvent)>,
-        max_receive_segments: NonZeroUsize,
-        endpoint: &proto::Endpoint,
     ) -> Self {
-        let recv_buf = vec![
-            0;
-            endpoint.config().get_max_udp_payload_size().min(64 * 1024) as usize
-                * max_receive_segments.get()
-                * BATCH_SIZE
-        ];
         Self {
             connections: ConnectionSet {
                 senders: FxHashMap::default(),
@@ -814,7 +803,6 @@ impl RecvState {
                 close: None,
             },
             incoming: VecDeque::new(),
-            recv_buf: recv_buf.into(),
             recv_limiter: WorkLimiter::new(RECV_TIME_BOUND),
         }
     }
@@ -829,62 +817,46 @@ impl RecvState {
         now: Instant,
     ) -> Result<PollProgress, io::Error> {
         let mut received_connection_packet = false;
-        let mut metas = [RecvMeta::default(); BATCH_SIZE];
-        let mut iovs: [IoSliceMut<'_>; BATCH_SIZE] = {
-            let mut bufs = self
-                .recv_buf
-                .chunks_mut(self.recv_buf.len() / BATCH_SIZE)
-                .map(IoSliceMut::new);
-
-            // expect() safe as self.recv_buf is chunked into BATCH_SIZE items
-            // and iovs will be of size BATCH_SIZE, thus from_fn is called
-            // exactly BATCH_SIZE times.
-            std::array::from_fn(|_| bufs.next().expect("BATCH_SIZE elements"))
-        };
         loop {
-            match socket.poll_recv(cx, &mut iovs, &mut metas) {
-                Poll::Ready(Ok(msgs)) => {
-                    self.recv_limiter.record_work(msgs);
-                    for (meta, buf) in metas.iter().zip(iovs.iter()).take(msgs) {
-                        let mut data: BytesMut = buf[0..meta.len].into();
-                        while !data.is_empty() {
-                            let buf = data.split_to(meta.stride.min(data.len()));
-                            let mut response_buffer = Vec::new();
-                            let addresses = FourTuple {
-                                remote: meta.addr,
-                                local_ip: meta.dst_ip,
-                            };
-                            match endpoint.handle(
-                                now,
-                                addresses,
-                                meta.ecn.map(proto_ecn),
-                                buf,
-                                &mut response_buffer,
-                            ) {
-                                Some(DatagramEvent::NewConnection(incoming)) => {
-                                    if self.connections.close.is_none() {
-                                        self.incoming.push_back(incoming);
-                                    } else {
-                                        let transmit =
-                                            endpoint.refuse(incoming, &mut response_buffer);
-                                        respond(transmit, &response_buffer, sender);
-                                    }
-                                }
-                                Some(DatagramEvent::ConnectionEvent(handle, event)) => {
-                                    // Ignoring errors from dropped connections that haven't yet been cleaned up
-                                    received_connection_packet = true;
-                                    let _ = self
-                                        .connections
-                                        .senders
-                                        .get_mut(&handle)
-                                        .unwrap()
-                                        .send(ConnectionEvent::Proto(event));
-                                }
-                                Some(DatagramEvent::Response(transmit)) => {
+            match socket.poll_recv_datagrams(cx) {
+                Poll::Ready(Ok(datagrams)) => {
+                    self.recv_limiter.record_work(datagrams.len());
+                    for datagram in datagrams {
+                        let mut response_buffer = Vec::new();
+                        let addresses = FourTuple {
+                            remote: datagram.remote,
+                            local_ip: datagram.local_ip,
+                        };
+                        match endpoint.handle(
+                            now,
+                            addresses,
+                            datagram.ecn.map(proto_ecn),
+                            datagram.data,
+                            &mut response_buffer,
+                        ) {
+                            Some(DatagramEvent::NewConnection(incoming)) => {
+                                if self.connections.close.is_none() {
+                                    self.incoming.push_back(incoming);
+                                } else {
+                                    let transmit =
+                                        endpoint.refuse(incoming, &mut response_buffer);
                                     respond(transmit, &response_buffer, sender);
                                 }
-                                None => {}
                             }
+                            Some(DatagramEvent::ConnectionEvent(handle, event)) => {
+                                // Ignoring errors from dropped connections that haven't yet been cleaned up
+                                received_connection_packet = true;
+                                let _ = self
+                                    .connections
+                                    .senders
+                                    .get_mut(&handle)
+                                    .unwrap()
+                                    .send(ConnectionEvent::Proto(event));
+                            }
+                            Some(DatagramEvent::Response(transmit)) => {
+                                respond(transmit, &response_buffer, sender);
+                            }
+                            None => {}
                         }
                     }
                 }
@@ -918,7 +890,6 @@ impl fmt::Debug for RecvState {
         f.debug_struct("RecvState")
             .field("incoming", &self.incoming)
             .field("connections", &self.connections)
-            // recv_buf too large
             .field("recv_limiter", &self.recv_limiter)
             .finish_non_exhaustive()
     }

quinn/src/runtime/mod.rs

@@ -10,7 +10,7 @@ use std::{
     task::{Context, Poll},
 };
 
-use udp::{RecvMeta, Transmit};
+use udp::{RecvMeta, ReceivedDatagrams, Transmit};
 
 use crate::Instant;
 
@@ -62,6 +62,16 @@ pub trait AsyncUdpSocket: Send + Sync + Debug + 'static {
         meta: &mut [RecvMeta],
     ) -> Poll<io::Result<usize>>;
 
+    /// Receive UDP datagrams as owned data, or register to be woken if receiving may succeed
+    ///
+    /// This is a higher-level API that handles buffer management and GRO splitting internally.
+    /// Each datagram in the returned collection contains its own `BytesMut` buffer suitable
+    /// for in-place decryption.
+    fn poll_recv_datagrams(
+        &mut self,
+        cx: &mut Context<'_>,
+    ) -> Poll<io::Result<ReceivedDatagrams>>;
+
     /// Look up the local IP address and port used by this socket
     fn local_addr(&self) -> io::Result<SocketAddr>;
 

quinn/src/runtime/smol.rs

@@ -1,11 +1,14 @@
 use std::{
     future::Future,
+    io::{self, IoSliceMut},
     num::NonZeroUsize,
     pin::Pin,
-    task::{Context, Poll},
+    sync::Arc,
+    task::{Context, Poll, ready},
     time::Instant,
 };
-use std::{io, sync::Arc, task::ready};
+
+use bytes::BytesMut;
 
 use async_io::Async;
 use async_io::Timer;
@@ -41,22 +44,35 @@ impl AsyncTimer for Timer {
     }
 }
 
+/// The parts of a UDP socket needed for sending
+///
+/// This is separated from UdpSocket so that senders can clone just what they need
+/// without carrying the receive buffer.
 #[derive(Debug, Clone)]
-struct UdpSocket {
+struct UdpSocketSend {
     io: Arc<Async<std::net::UdpSocket>>,
     inner: Arc<udp::UdpSocketState>,
 }
 
+#[derive(Debug)]
+struct UdpSocket {
+    send: UdpSocketSend,
+    recv_buf: Vec<u8>,
+}
+
 impl UdpSocket {
     fn new(sock: std::net::UdpSocket) -> io::Result<Self> {
         Ok(Self {
-            inner: Arc::new(udp::UdpSocketState::new((&sock).into())?),
-            io: Arc::new(Async::new_nonblocking(sock)?),
+            send: UdpSocketSend {
+                inner: Arc::new(udp::UdpSocketState::new((&sock).into())?),
+                io: Arc::new(Async::new_nonblocking(sock)?),
+            },
+            recv_buf: Vec::new(),
         })
     }
 }
 
-impl UdpSenderHelperSocket for UdpSocket {
+impl UdpSenderHelperSocket for UdpSocketSend {
     fn max_transmit_segments(&self) -> NonZeroUsize {
         self.inner.max_gso_segments()
     }
@@ -68,7 +84,8 @@ impl UdpSenderHelperSocket for UdpSocket {
 
 impl AsyncUdpSocket for UdpSocket {
     fn create_sender(&self) -> Pin<Box<dyn UdpSender>> {
-        Box::pin(UdpSenderHelper::new(self.clone(), |socket: &Self| {
+        let core = self.send.clone();
+        Box::pin(UdpSenderHelper::new(core, |socket: &UdpSocketSend| {
             let socket = socket.clone();
             async move { socket.io.writable().await }
         }))
@@ -81,22 +98,69 @@ impl AsyncUdpSocket for UdpSocket {
         meta: &mut [udp::RecvMeta],
     ) -> Poll<io::Result<usize>> {
         loop {
-            ready!(self.io.poll_readable(cx))?;
-            if let Ok(res) = self.inner.recv((&self.io).into(), bufs, meta) {
+            ready!(self.send.io.poll_readable(cx))?;
+            if let Ok(res) = self.send.inner.recv((&self.send.io).into(), bufs, meta) {
                 return Poll::Ready(Ok(res));
             }
         }
     }
 
+    fn poll_recv_datagrams(
+        &mut self,
+        cx: &mut Context<'_>,
+    ) -> Poll<io::Result<udp::ReceivedDatagrams>> {
+        // Ensure buffer is sized for GRO coalescing
+        // Use 1500 (typical Ethernet MTU) as max payload size
+        const MAX_PAYLOAD_SIZE: usize = 1500;
+        let gro_segments = self.send.inner.gro_segments().get();
+        let buf_size = MAX_PAYLOAD_SIZE * gro_segments;
+        let total_size = buf_size * udp::BATCH_SIZE;
+        if self.recv_buf.len() < total_size {
+            self.recv_buf.resize(total_size, 0);
+        }
+
+        loop {
+            ready!(self.send.io.poll_readable(cx))?;
+
+            // Prepare IoSliceMut array
+            let mut bufs: [IoSliceMut<'_>; udp::BATCH_SIZE] =
+                std::array::from_fn(|_| IoSliceMut::new(&mut []));
+            for (i, chunk) in self.recv_buf.chunks_mut(buf_size).enumerate().take(udp::BATCH_SIZE) {
+                bufs[i] = IoSliceMut::new(chunk);
+            }
+            let mut metas = [udp::RecvMeta::default(); udp::BATCH_SIZE];
+
+            if let Ok(msg_count) = self.send.inner.recv((&self.send.io).into(), &mut bufs, &mut metas) {
+                // Convert to ReceivedDatagrams, splitting by stride
+                let mut result = udp::ReceivedDatagrams::new();
+                for (meta, buf) in metas.iter().zip(bufs.iter()).take(msg_count) {
+                    let mut offset = 0;
+                    while offset < meta.len {
+                        let stride = meta.stride.min(meta.len - offset);
+                        let data = BytesMut::from(&buf[offset..offset + stride]);
+                        result.push(udp::ReceivedDatagram {
+                            data,
+                            remote: meta.addr,
+                            local_ip: meta.dst_ip,
+                            ecn: meta.ecn,
+                        });
+                        offset += stride;
+                    }
+                }
+                return Poll::Ready(Ok(result));
+            }
+        }
+    }
+
     fn local_addr(&self) -> io::Result<std::net::SocketAddr> {
-        self.io.as_ref().as_ref().local_addr()
+        self.send.io.as_ref().as_ref().local_addr()
     }
 
     fn may_fragment(&self) -> bool {
-        self.inner.may_fragment()
+        self.send.inner.may_fragment()
     }
 
     fn max_receive_segments(&self) -> NonZeroUsize {
-        self.inner.gro_segments()
+        self.send.inner.gro_segments()
     }
 }

quinn/src/runtime/tokio.rs

@@ -1,7 +1,7 @@
 use std::{
     fmt::Debug,
     future::Future,
-    io,
+    io::{self, IoSliceMut},
     num::NonZeroUsize,
     pin::Pin,
     sync::Arc,
@@ -9,6 +9,7 @@ use std::{
     time::Instant,
 };
 
+use bytes::BytesMut;
 use tokio::{
     io::Interest,
     time::{Sleep, sleep_until},
@@ -31,8 +32,11 @@ impl Runtime for TokioRuntime {
 
     fn wrap_udp_socket(&self, sock: std::net::UdpSocket) -> io::Result<Box<dyn AsyncUdpSocket>> {
         Ok(Box::new(UdpSocket {
-            inner: Arc::new(udp::UdpSocketState::new((&sock).into())?),
-            io: Arc::new(tokio::net::UdpSocket::from_std(sock)?),
+            send: UdpSocketSend {
+                inner: Arc::new(udp::UdpSocketState::new((&sock).into())?),
+                io: Arc::new(tokio::net::UdpSocket::from_std(sock)?),
+            },
+            recv_buf: Vec::new(),
         }))
     }
 
@@ -50,13 +54,23 @@ impl AsyncTimer for Sleep {
     }
 }
 
+/// The parts of a UDP socket needed for sending
+///
+/// This is separated from UdpSocket so that senders can clone just what they need
+/// without carrying the receive buffer.
 #[derive(Debug, Clone)]
-struct UdpSocket {
+struct UdpSocketSend {
     io: Arc<tokio::net::UdpSocket>,
     inner: Arc<udp::UdpSocketState>,
 }
 
-impl UdpSenderHelperSocket for UdpSocket {
+#[derive(Debug)]
+struct UdpSocket {
+    send: UdpSocketSend,
+    recv_buf: Vec<u8>,
+}
+
+impl UdpSenderHelperSocket for UdpSocketSend {
     fn max_transmit_segments(&self) -> NonZeroUsize {
         self.inner.max_gso_segments()
     }
@@ -70,7 +84,8 @@ impl UdpSenderHelperSocket for UdpSocket {
 
 impl AsyncUdpSocket for UdpSocket {
     fn create_sender(&self) -> Pin<Box<dyn super::UdpSender>> {
-        Box::pin(UdpSenderHelper::new(self.clone(), |socket: &Self| {
+        let core = self.send.clone();
+        Box::pin(UdpSenderHelper::new(core, |socket: &UdpSocketSend| {
             let socket = socket.clone();
             async move { socket.io.writable().await }
         }))
@@ -83,24 +98,73 @@ impl AsyncUdpSocket for UdpSocket {
         meta: &mut [udp::RecvMeta],
     ) -> Poll<io::Result<usize>> {
         loop {
-            ready!(self.io.poll_recv_ready(cx))?;
-            if let Ok(res) = self.io.try_io(Interest::READABLE, || {
-                self.inner.recv((&self.io).into(), bufs, meta)
+            ready!(self.send.io.poll_recv_ready(cx))?;
+            if let Ok(res) = self.send.io.try_io(Interest::READABLE, || {
+                self.send.inner.recv((&self.send.io).into(), bufs, meta)
             }) {
                 return Poll::Ready(Ok(res));
             }
         }
     }
 
+    fn poll_recv_datagrams(
+        &mut self,
+        cx: &mut Context<'_>,
+    ) -> Poll<io::Result<udp::ReceivedDatagrams>> {
+        // Ensure buffer is sized for GRO coalescing
+        // Use 1500 (typical Ethernet MTU) as max payload size
+        const MAX_PAYLOAD_SIZE: usize = 1500;
+        let gro_segments = self.send.inner.gro_segments().get();
+        let buf_size = MAX_PAYLOAD_SIZE * gro_segments;
+        let total_size = buf_size * udp::BATCH_SIZE;
+        if self.recv_buf.len() < total_size {
+            self.recv_buf.resize(total_size, 0);
+        }
+
+        loop {
+            ready!(self.send.io.poll_recv_ready(cx))?;
+
+            // Prepare IoSliceMut array
+            let mut bufs: [IoSliceMut<'_>; udp::BATCH_SIZE] =
+                std::array::from_fn(|_| IoSliceMut::new(&mut []));
+            for (i, chunk) in self.recv_buf.chunks_mut(buf_size).enumerate().take(udp::BATCH_SIZE) {
+                bufs[i] = IoSliceMut::new(chunk);
+            }
+            let mut metas = [udp::RecvMeta::default(); udp::BATCH_SIZE];
+
+            if let Ok(msg_count) = self.send.io.try_io(Interest::READABLE, || {
+                self.send.inner.recv((&self.send.io).into(), &mut bufs, &mut metas)
+            }) {
+                // Convert to ReceivedDatagrams, splitting by stride
+                let mut result = udp::ReceivedDatagrams::new();
+                for (meta, buf) in metas.iter().zip(bufs.iter()).take(msg_count) {
+                    let mut offset = 0;
+                    while offset < meta.len {
+                        let stride = meta.stride.min(meta.len - offset);
+                        let data = BytesMut::from(&buf[offset..offset + stride]);
+                        result.push(udp::ReceivedDatagram {
+                            data,
+                            remote: meta.addr,
+                            local_ip: meta.dst_ip,
+                            ecn: meta.ecn,
+                        });
+                        offset += stride;
+                    }
+                }
+                return Poll::Ready(Ok(result));
+            }
+        }
+    }
+
     fn local_addr(&self) -> io::Result<std::net::SocketAddr> {
-        self.io.local_addr()
+        self.send.io.local_addr()
     }
 
     fn may_fragment(&self) -> bool {
-        self.inner.may_fragment()
+        self.send.inner.may_fragment()
     }
 
     fn max_receive_segments(&self) -> NonZeroUsize {
-        self.inner.gro_segments()
+        self.send.inner.gro_segments()
     }
 }