feat(udp): handle partial sendmsg_x progress

quinn-udp/benches/throughput.rs

@@ -57,7 +57,7 @@ pub fn criterion_benchmark(c: &mut Criterion) {
             1
         };
         let msg = vec![0xAB; min(MAX_DATAGRAM_SIZE, SEGMENT_SIZE * gso_segments)];
-        let transmit = Transmit {
+        let make_transmit = || Transmit {
             destination: dst_addr,
             ecn: None,
             contents: &msg,
@@ -82,14 +82,25 @@ pub fn criterion_benchmark(c: &mut Criterion) {
 
                 let mut sent: usize = 0;
                 let mut received: usize = 0;
+                let mut pending = make_transmit();
                 while sent < TOTAL_BYTES {
                     send_socket.writable().await.unwrap();
-                    send_socket
+                    let sent_datagrams = send_socket
                         .try_io(Interest::WRITABLE, || {
-                            send_state.send((&send_socket).into(), &transmit)
+                            send_state.send((&send_socket).into(), &pending)
                         })
                         .unwrap();
-                    sent += transmit.contents.len();
+                    let previous_len = pending.contents.len();
+                    match pending.advance(sent_datagrams) {
+                        Some(remainder) => {
+                            sent += previous_len - remainder.contents.len();
+                            pending = remainder;
+                        }
+                        None => {
+                            sent += previous_len;
+                            pending = make_transmit();
+                        }
+                    }
 
                     while received < sent {
                         recv_socket.readable().await.unwrap();

quinn-udp/src/apple_fast.rs

@@ -16,42 +16,36 @@ use crate::{
 pub(crate) fn send(
     state: &UdpSocketState,
     io: SockRef<'_>,
-    transmit: &Transmit<'_>,
-) -> io::Result<()> {
+    transmit: Transmit<'_>,
+) -> io::Result<usize> {
     if state.is_apple_fast_path_enabled() {
         send_via_sendmsg_x(state, io, transmit)
     } else {
-        send_single(state, io, transmit)
+        send_single(state, io, &transmit.limit(1))
     }
 }
 
 /// Send using the fast `sendmsg_x` API
 fn send_via_sendmsg_x(
     state: &UdpSocketState,
     io: SockRef<'_>,
-    transmit: &Transmit<'_>,
-) -> io::Result<()> {
+    transmit: Transmit<'_>,
+) -> io::Result<usize> {
     let mut hdrs = unsafe { mem::zeroed::<[msghdr_x; BATCH_SIZE]>() };
     let mut iovs = unsafe { mem::zeroed::<[libc::iovec; BATCH_SIZE]>() };
     let mut ctrls = [cmsg::Aligned([0u8; cmsg::LEN]); BATCH_SIZE];
     let addr = socket2::SockAddr::from(transmit.destination);
-    let segment_size = transmit.segment_size.unwrap_or(transmit.contents.len());
-    let mut cnt = 0;
-    debug_assert!(transmit.contents.len().div_ceil(segment_size) <= BATCH_SIZE);
-    for (i, chunk) in transmit
-        .contents
-        .chunks(segment_size)
-        .enumerate()
-        .take(BATCH_SIZE)
-    {
+
+    let mut prepare = |i: usize, chunk: &[u8]| -> io::Result<()> {
+        let segment = Transmit {
+            destination: transmit.destination,
+            ecn: transmit.ecn,
+            contents: chunk,
+            segment_size: None,
+            src_ip: transmit.src_ip,
+        };
         prepare_msg_x(
-            &Transmit {
-                destination: transmit.destination,
-                ecn: transmit.ecn,
-                contents: chunk,
-                segment_size: Some(chunk.len()),
-                src_ip: transmit.src_ip,
-            },
+            &segment,
             &addr,
             &mut hdrs[i],
             &mut iovs[i],
@@ -61,13 +55,45 @@ fn send_via_sendmsg_x(
         );
         hdrs[i].msg_datalen = chunk.len();
         state.check_send_buffer_limit(chunk.len(), &hdrs[i])?;
-        cnt += 1;
-    }
+
+        Ok(())
+    };
+
+    debug_assert!(transmit.datagram_count() <= BATCH_SIZE);
+    let cnt = if transmit.contents.is_empty() {
+        prepare(0, &[])?;
+        1
+    } else {
+        let segment_size = transmit.segment_size.unwrap_or(transmit.contents.len());
+        let mut cnt = 0;
+
+        for (i, chunk) in transmit
+            .contents
+            .chunks(segment_size)
+            .enumerate()
+            .take(BATCH_SIZE)
+        {
+            prepare(i, chunk)?;
+            cnt += 1;
+        }
+
+        cnt
+    };
+
     let Some(sendmsg_x) = state.resolve_apple_fast_fn(sendmsg_x_fn) else {
-        return send_single(state, io, transmit);
+        return send_single(state, io, &transmit.limit(1));
     };
-    retry_if_interrupted(|| unsafe { sendmsg_x(io.as_raw_fd(), hdrs.as_ptr(), cnt as u32, 0) })?;
-    Ok(())
+
+    let sent = retry_if_interrupted(|| unsafe {
+        sendmsg_x(io.as_raw_fd(), hdrs.as_ptr(), cnt as u32, 0)
+    })? as usize;
+
+    if sent == 0 {
+        return Err(io::ErrorKind::WriteZero.into());
+    }
+    debug_assert!(sent <= cnt);
+
+    Ok(sent)
 }
 
 /// Prepares an `msghdr_x` for use with `sendmsg_x`

quinn-udp/src/fallback.rs

@@ -4,7 +4,7 @@ use std::{
     time::Instant,
 };
 
-use super::{IO_ERROR_LOG_INTERVAL, RecvMeta, Transmit, UdpSockRef, log_sendmsg_error};
+use super::{IO_ERROR_LOG_INTERVAL, RecvMeta, SendCount, Transmit, UdpSockRef, log_sendmsg_error};
 
 /// Fallback UDP socket interface that stubs out all special functionality
 ///
@@ -35,21 +35,30 @@ impl UdpSocketState {
     ///
     /// If you would like to handle these errors yourself, use [`UdpSocketState::try_send`]
     /// instead.
-    pub fn send(&self, socket: UdpSockRef<'_>, transmit: &Transmit<'_>) -> io::Result<()> {
-        match send(socket, transmit) {
-            Ok(()) => Ok(()),
+    pub fn send(&self, socket: UdpSockRef<'_>, transmit: &Transmit<'_>) -> io::Result<SendCount> {
+        let transmit = transmit.limit(1);
+
+        match send(socket, &transmit) {
+            Ok(sent) => Ok(SendCount::from_datagram_count(sent)),
             Err(e) if e.kind() == io::ErrorKind::WouldBlock => Err(e),
             Err(e) => {
-                log_sendmsg_error(&self.last_send_error, e, transmit);
+                log_sendmsg_error(&self.last_send_error, e, &transmit);
 
-                Ok(())
+                Ok(SendCount::from_datagram_count(1))
             }
         }
     }
 
-    /// Sends a [`Transmit`] on the given socket without any additional error handling.
-    pub fn try_send(&self, socket: UdpSockRef<'_>, transmit: &Transmit<'_>) -> io::Result<()> {
-        send(socket, transmit)
+    /// Sends the first datagram of a [`Transmit`] without any additional error handling.
+    pub fn try_send(
+        &self,
+        socket: UdpSockRef<'_>,
+        transmit: &Transmit<'_>,
+    ) -> io::Result<SendCount> {
+        let transmit = transmit.limit(1);
+        let sent = send(socket, &transmit)?;
+
+        Ok(SendCount::from_datagram_count(sent))
     }
 
     pub fn recv(
@@ -117,11 +126,13 @@ impl UdpSocketState {
     }
 }
 
-fn send(socket: UdpSockRef<'_>, transmit: &Transmit<'_>) -> io::Result<()> {
+fn send(socket: UdpSockRef<'_>, transmit: &Transmit<'_>) -> io::Result<usize> {
     socket.0.send_to(
         transmit.contents,
         &socket2::SockAddr::from(transmit.destination),
-    )
+    )?;
+
+    Ok(transmit.datagram_count())
 }
 
 pub(crate) const BATCH_SIZE: usize = 1;