laantungir/fips
1
0
Fork 0
mirror of https://github.com/jmcorgan/fips.git synced 2026-08-09 08:14:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

udp: batched recvmmsg receive on Linux (32-pkt bursts)

Browse Source 
The UDP recv loop drained the kernel queue one packet per recvmsg(2).
Each call paid full per-syscall + per-task-wakeup overhead (~50us avg
including a futex-based scheduler hop), so under sustained load the
loop ran at one rx event per scheduler quantum — the dominant cap on
inbound packet rate.

On Linux, switch the steady-state path to recvmmsg(2) with a 32-packet
batch. A single readable() wakeup drains up to 32 datagrams in one
syscall before yielding back to the reactor. Stack-allocated mmsghdr
arrays sized to a module-level `BATCH_SIZE` constant.

`SO_RXQ_OVFL` is sampled once per batch off the cmsg chain of `msgs[0]`
and plumbed through `AsyncUdpSocket::recv_batch` as `(count, drops)`.
The counter is socket-wide and monotonic, so a single sample per batch
gives the 1Hz `sample_transport_congestion()` detector ample fresh
values under load (one batch = up to 32 datagrams). Cost is one
stack-allocated CMSG_SPACE(4) buffer + one CMSG_FIRSTHDR walk per
batch syscall.

macOS / Windows fall through to the per-packet recv_from loop —
recvmmsg is Linux-specific and the per-packet API is fast enough on
those platforms for now (recvmsg_x for Darwin can be added later).

The slice-array build also drops the `MaybeUninit::uninit().assume_init()`
+ `transmute` pair for `std::array::from_fn` over a single shared
`backing.iter_mut()` — same disjoint mutable borrows, no `unsafe`.
This commit is contained in:
Martti Malmi
2026-05-10 00:28:45 +03:00
committed by Dev
parent cd56fee7cf
commit 253dddabe3
2 changed files with 257 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/transport/udp/mod.rs
+117 -49
View File
@@ -412,6 +412,11 @@ impl Drop for UdpTransport {
}
/// UDP receive loop - runs as a spawned task.
///
/// On Linux, drains the kernel UDP queue in 32-packet bursts via `recvmmsg`
/// to amortise the per-syscall + per-task-wakeup overhead. macOS / Windows
/// fall through to single-packet `recv_from`. Either way every datagram
/// is forwarded to `packet_tx` in arrival order.
async fn udp_receive_loop(
socket: AsyncUdpSocket,
transport_id: TransportId,
@@ -419,66 +424,129 @@ async fn udp_receive_loop(
mtu: u16,
stats: Arc<UdpStats>,
) {
// Buffer with headroom for slightly oversized packets
let mut buf = vec![0u8; mtu as usize + 100];
debug!(transport_id = %transport_id, "UDP receive loop starting");
loop {
match socket.recv_from(&mut buf).await {
Ok((len, remote_addr, kernel_drops)) => {
stats.record_recv(len);
stats.set_kernel_drops(kernel_drops as u64);
#[cfg(target_os = "linux")]
{
const BATCH: usize = 32;
let buf_size = mtu as usize + 100;
// One contiguous backing alloc; slice it for recvmmsg.
let mut backing: Vec<Vec<u8>> = (0..BATCH).map(|_| vec![0u8; buf_size]).collect();
let mut addrs: [Option<std::net::SocketAddr>; BATCH] = std::array::from_fn(|_| None);
let mut lens: [usize; BATCH] = [0; BATCH];
// Drop stray punch probes / acks. After bootstrap-handoff
// adopts a socket, the remote side may keep retrying its
// own punch attempt for several seconds; those probes
// arrive here and would otherwise be parsed as FMP frames
// (their first byte 0x4E has high-nibble 0x4 → bogus
// "FMP version 4" warnings).
if is_punch_packet(&buf[..len]) {
trace!(
transport_id = %transport_id,
remote_addr = %remote_addr,
bytes = len,
"Dropping stray punch probe/ack on UDP transport"
);
continue;
loop {
// Build mutable slice references for the syscall layer.
// Drawing from a single `iter_mut()` keeps the borrows disjoint
// without `MaybeUninit`/`transmute`.
let mut bufs: [&mut [u8]; BATCH] = {
let mut iter = backing.iter_mut();
std::array::from_fn(|_| iter.next().unwrap().as_mut_slice())
};
match socket.recv_batch(&mut bufs, &mut addrs, &mut lens).await {
Ok((count, kernel_drops)) => {
stats.set_kernel_drops(kernel_drops as u64);
for i in 0..count {
let len = lens[i];
let Some(remote_addr) = addrs[i] else {
continue;
};
stats.record_recv(len);
let buf = &backing[i][..len];
if is_punch_packet(buf) {
trace!(
transport_id = %transport_id,
remote_addr = %remote_addr,
bytes = len,
"Dropping stray punch probe/ack on UDP transport"
);
continue;
}
let data = buf.to_vec();
let addr = TransportAddr::from_string(&remote_addr.to_string());
let packet = ReceivedPacket::new(transport_id, addr, data);
trace!(
transport_id = %transport_id,
remote_addr = %remote_addr,
bytes = len,
"UDP packet received"
);
if packet_tx.send(packet).await.is_err() {
debug!(
transport_id = %transport_id,
"Packet channel closed, stopping receive loop"
);
return;
}
}
}
let data = buf[..len].to_vec();
let addr = TransportAddr::from_string(&remote_addr.to_string());
let packet = ReceivedPacket::new(transport_id, addr, data);
trace!(
transport_id = %transport_id,
remote_addr = %remote_addr,
bytes = len,
"UDP packet received"
);
if packet_tx.send(packet).await.is_err() {
// Receiver dropped, exit loop
debug!(
Err(e) => {
stats.record_recv_error();
warn!(
transport_id = %transport_id,
"Packet channel closed, stopping receive loop"
error = %e,
"UDP receive error"
);
break;
}
}
Err(e) => {
stats.record_recv_error();
// Log error but continue - transient errors are expected
warn!(
transport_id = %transport_id,
error = %e,
"UDP receive error"
);
}
}
}
debug!(transport_id = %transport_id, "UDP receive loop stopped");
#[cfg(not(target_os = "linux"))]
{
let mut buf = vec![0u8; mtu as usize + 100];
loop {
match socket.recv_from(&mut buf).await {
Ok((len, remote_addr, kernel_drops)) => {
stats.record_recv(len);
stats.set_kernel_drops(kernel_drops as u64);
if is_punch_packet(&buf[..len]) {
trace!(
transport_id = %transport_id,
remote_addr = %remote_addr,
bytes = len,
"Dropping stray punch probe/ack on UDP transport"
);
continue;
}
let data = buf[..len].to_vec();
let addr = TransportAddr::from_string(&remote_addr.to_string());
let packet = ReceivedPacket::new(transport_id, addr, data);
trace!(
transport_id = %transport_id,
remote_addr = %remote_addr,
bytes = len,
"UDP packet received"
);
if packet_tx.send(packet).await.is_err() {
debug!(
transport_id = %transport_id,
"Packet channel closed, stopping receive loop"
);
break;
}
}
Err(e) => {
stats.record_recv_error();
warn!(
transport_id = %transport_id,
error = %e,
"UDP receive error"
);
}
}
}
}
}
// ============================================================================
src/transport/udp/socket.rs
+140 -1
View File
@@ -29,6 +29,13 @@ mod platform {
use std::os::unix::io::{AsRawFd, RawFd};
use tokio::io::unix::AsyncFd;
/// Maximum number of datagrams a single recvmmsg / sendmmsg syscall
/// will pull from / push to the kernel. Tuned to amortise syscall +
/// per-task-wakeup overhead across a useful burst without blowing
/// the stack (each slot owns an mmsghdr + sockaddr_storage + iovec).
#[cfg(target_os = "linux")]
const BATCH_SIZE: usize = 32;
/// Wrapper around a `socket2::Socket` providing sync send/recv with
/// `SO_RXQ_OVFL` ancillary data parsing.
pub struct UdpRawSocket {
@@ -232,6 +239,11 @@ mod platform {
/// Returns `(bytes_read, source_addr, kernel_drops)`. The `kernel_drops`
/// value is a cumulative counter since socket creation; it is 0 if
/// `SO_RXQ_OVFL` is not supported.
///
/// On Linux the production receive path uses `recv_batch` (recvmmsg);
/// this single-packet variant remains for non-Linux unix targets and
/// for the local `tests` module.
#[cfg_attr(target_os = "linux", allow(dead_code))]
pub fn recv_from(&self, buf: &mut [u8]) -> std::io::Result<(usize, SocketAddr, u32)> {
let fd = self.inner.as_raw_fd();
@@ -287,6 +299,98 @@ mod platform {
Ok((n as usize, addr, drops))
}
/// Receive up to `BATCH_SIZE` datagrams in a single recvmmsg syscall
/// (Linux only — macOS falls through to per-packet recvmsg).
///
/// Returns `(count, kernel_drops)`. Caller pre-sizes `bufs` (each
/// must be at least the configured MTU) and the matching `addrs` /
/// `lens` slices; on return, slots `[0..count)` are valid.
///
/// `kernel_drops` is the `SO_RXQ_OVFL` cumulative counter sampled
/// from the cmsg chain of the FIRST datagram in the batch. The
/// counter is monotonic per-socket since `SO_RXQ_OVFL` was enabled,
/// so a single sample per batch is sufficient to feed the 1Hz
/// congestion detector in `sample_transport_congestion()`. Returns
/// `(0, 0)` on a spurious wakeup with no datagrams ready.
#[cfg(target_os = "linux")]
pub fn recv_batch(
&self,
bufs: &mut [&mut [u8]],
addrs: &mut [Option<SocketAddr>],
lens: &mut [usize],
) -> std::io::Result<(usize, u32)> {
let n = bufs.len().min(addrs.len()).min(lens.len()).min(BATCH_SIZE);
if n == 0 {
return Ok((0, 0));
}
let fd = self.inner.as_raw_fd();
// CMSG buffer wired to msgs[0] only. SO_RXQ_OVFL delivers a
// monotonic u32 drop counter; sampling once per batch gives
// the 1Hz congestion detector ample fresh values under load
// (one batch = up to 32 datagrams).
const CMSG_BUF_SIZE: usize = unsafe { libc::CMSG_SPACE(4) } as usize;
let mut cmsg_buf = [0u8; CMSG_BUF_SIZE];
// Stack-allocated parallel arrays; lifetime tied to this call.
let mut iovs: [libc::iovec; BATCH_SIZE] = unsafe { std::mem::zeroed() };
let mut storages: [libc::sockaddr_storage; BATCH_SIZE] = unsafe { std::mem::zeroed() };
let mut msgs: [libc::mmsghdr; BATCH_SIZE] = unsafe { std::mem::zeroed() };
for i in 0..n {
iovs[i].iov_base = bufs[i].as_mut_ptr() as *mut libc::c_void;
iovs[i].iov_len = bufs[i].len();
msgs[i].msg_hdr.msg_name = &mut storages[i] as *mut _ as *mut libc::c_void;
msgs[i].msg_hdr.msg_namelen =
std::mem::size_of::<libc::sockaddr_storage>() as libc::socklen_t;
msgs[i].msg_hdr.msg_iov = &mut iovs[i];
msgs[i].msg_hdr.msg_iovlen = 1;
msgs[i].msg_len = 0;
}
// Only msgs[0] carries a cmsg buffer — sampling the OVFL counter
// there is enough since it is socket-wide and monotonic.
msgs[0].msg_hdr.msg_control = cmsg_buf.as_mut_ptr() as *mut libc::c_void;
msgs[0].msg_hdr.msg_controllen = cmsg_buf.len() as _;
let r = unsafe {
libc::recvmmsg(
fd,
msgs.as_mut_ptr(),
n as libc::c_uint,
0,
std::ptr::null_mut(),
)
};
if r < 0 {
return Err(std::io::Error::last_os_error());
}
let count = r as usize;
for i in 0..count {
lens[i] = msgs[i].msg_len as usize;
addrs[i] = sockaddr_to_socket_addr(&storages[i]).ok();
}
// Walk msgs[0] cmsg chain for SO_RXQ_OVFL. Skip when no
// datagram landed (cmsg buffer is undefined in that case).
let mut drops: u32 = 0;
if count > 0 {
unsafe {
let mut cmsg = libc::CMSG_FIRSTHDR(&msgs[0].msg_hdr);
while !cmsg.is_null() {
if (*cmsg).cmsg_level == libc::SOL_SOCKET
&& (*cmsg).cmsg_type == libc::SO_RXQ_OVFL
{
let data = libc::CMSG_DATA(cmsg);
drops = std::ptr::read_unaligned(data as *const u32);
}
cmsg = libc::CMSG_NXTHDR(&msgs[0].msg_hdr, cmsg);
}
}
}
Ok((count, drops))
}
/// Wrap this socket in a tokio `AsyncFd` for async I/O.
pub fn into_async(self) -> Result<AsyncUdpSocket, TransportError> {
let async_fd = AsyncFd::new(self)
@@ -336,7 +440,11 @@ mod platform {
/// Receive a payload, source address, and kernel drop counter.
///
/// Returns `(bytes_read, source_addr, kernel_drops)`.
/// Returns `(bytes_read, source_addr, kernel_drops)`. On Linux the
/// production receive path uses `recv_batch`; this single-packet
/// variant remains for non-Linux unix targets and for the local
/// `tests` module.
#[cfg_attr(target_os = "linux", allow(dead_code))]
pub async fn recv_from(
&self,
buf: &mut [u8],
@@ -355,6 +463,37 @@ mod platform {
}
}
}
/// Drain up to `BATCH_SIZE` datagrams from the kernel via
/// `recvmmsg` (Linux). Returns `(count, kernel_drops)`; same
/// buffer / addr / len contract as `UdpRawSocket::recv_batch`.
#[cfg(target_os = "linux")]
pub async fn recv_batch(
&self,
bufs: &mut [&mut [u8]],
addrs: &mut [Option<SocketAddr>],
lens: &mut [usize],
) -> Result<(usize, u32), TransportError> {
loop {
let mut guard = self
.inner
.readable()
.await
.map_err(|e| TransportError::RecvFailed(format!("readable wait: {}", e)))?;
match guard.try_io(|inner| inner.get_ref().recv_batch(bufs, addrs, lens)) {
Ok(Ok((0, _))) => {
// Spurious wakeup or no datagrams ready — yield
// back to the reactor instead of busy-looping.
guard.clear_ready();
continue;
}
Ok(Ok(result)) => return Ok(result),
Ok(Err(e)) => return Err(TransportError::RecvFailed(format!("{}", e))),
Err(_would_block) => continue,
}
}
}
}
/// Convert a `libc::sockaddr_storage` to `std::net::SocketAddr`.