From 21d530dbd015e5b40e87eeb03ed7cdc83fcbd21e Mon Sep 17 00:00:00 2001 From: Ruediger Klaehn Date: Thu, 29 Jan 2026 20:01:56 +0200 Subject: [PATCH] Experiment: simplify the read side interface --- Cargo.lock | 2 + Cargo.toml | 1 + quinn-udp/Cargo.toml | 2 + quinn-udp/src/fallback.rs | 38 ++++++++++++- quinn-udp/src/lib.rs | 74 +++++++++++++++++++++++++ quinn-udp/src/unix.rs | 61 ++++++++++++++++++++- quinn-udp/src/windows.rs | 46 +++++++++++++++- quinn/src/endpoint.rs | 107 ++++++++++++++----------------------- quinn/src/runtime/mod.rs | 12 ++++- quinn/src/runtime/smol.rs | 88 +++++++++++++++++++++++++----- quinn/src/runtime/tokio.rs | 88 +++++++++++++++++++++++++----- 11 files changed, 423 insertions(+), 96 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ca3ae2ab1..5b1dcf3a3 100644 --- a/Cargo.lock +++ b/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", diff --git a/Cargo.toml b/Cargo.toml index 75733f8dd..fadab6ff7 100644 --- a/Cargo.toml +++ b/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" } diff --git a/quinn-udp/Cargo.toml b/quinn-udp/Cargo.toml index 46eddf838..826815544 100644 --- a/quinn-udp/Cargo.toml +++ b/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] diff --git a/quinn-udp/src/fallback.rs b/quinn-udp/src/fallback.rs index 01a951e51..608db42d1 100644 --- a/quinn-udp/src/fallback.rs +++ b/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 { + 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 diff --git a/quinn-udp/src/lib.rs b/quinn-udp/src/lib.rs index d1b1d6342..6a886d774 100644 --- a/quinn-udp/src/lib.rs +++ b/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, + /// The Explicit Congestion Notification bits + pub ecn: Option, +} + +/// 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 { + 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> { diff --git a/quinn-udp/src/unix.rs b/quinn-udp/src/unix.rs index 825feac12..523fc6f31 100644 --- a/quinn-udp/src/unix.rs +++ b/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 { + // 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). /// diff --git a/quinn-udp/src/windows.rs b/quinn-udp/src/windows.rs index e3519a29d..d8bd5fce3 100644 --- a/quinn-udp/src/windows.rs +++ b/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 { + // 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). /// diff --git a/quinn/src/endpoint.rs b/quinn/src/endpoint.rs index 7b616228d..ffcb77b68 100644 --- a/quinn/src/endpoint.rs +++ b/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, ) -> 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, 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 { 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() } diff --git a/quinn/src/runtime/mod.rs b/quinn/src/runtime/mod.rs index 41bcf803d..08b73814f 100644 --- a/quinn/src/runtime/mod.rs +++ b/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>; + /// 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>; + /// Look up the local IP address and port used by this socket fn local_addr(&self) -> io::Result; diff --git a/quinn/src/runtime/smol.rs b/quinn/src/runtime/smol.rs index 0b6681e8e..d9daaaaee 100644 --- a/quinn/src/runtime/smol.rs +++ b/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>, inner: Arc, } +#[derive(Debug)] +struct UdpSocket { + send: UdpSocketSend, + recv_buf: Vec, +} + impl UdpSocket { fn new(sock: std::net::UdpSocket) -> io::Result { 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::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> { 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> { + // 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 { - 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() } } diff --git a/quinn/src/runtime/tokio.rs b/quinn/src/runtime/tokio.rs index f1cdc5663..6d5d4476d 100644 --- a/quinn/src/runtime/tokio.rs +++ b/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> { 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, inner: Arc, } -impl UdpSenderHelperSocket for UdpSocket { +#[derive(Debug)] +struct UdpSocket { + send: UdpSocketSend, + recv_buf: Vec, +} + +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::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> { 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> { + // 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 { - 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() } }