refactor!: use NonZeroUsize for segments sizes

Start introducing more type safety for implicit assumptions, step by step

quinn-proto/src/connection/mod.rs

@@ -4,7 +4,7 @@ use std::{
     convert::TryFrom,
     fmt, io, mem,
     net::{IpAddr, SocketAddr},
-    num::NonZeroU32,
+    num::{NonZeroU32, NonZeroUsize},
     ops::Not,
     sync::Arc,
 };
@@ -888,7 +888,7 @@ impl Connection {
     pub fn poll_transmit(
         &mut self,
         now: Instant,
-        max_datagrams: usize,
+        max_datagrams: NonZeroUsize,
         buf: &mut Vec<u8>,
     ) -> Option<Transmit> {
         if let Some(probing) = self
@@ -911,9 +911,8 @@ impl Connection {
             });
         }
 
-        assert!(max_datagrams != 0);
         let max_datagrams = match self.config.enable_segmentation_offload {
-            false => 1,
+            false => NonZeroUsize::new(1).expect("known"),
             true => max_datagrams,
         };
 
@@ -1154,7 +1153,7 @@ impl Connection {
 
             // If the datagram is full, we need to start a new one.
             if transmit.datagram_remaining_mut() == 0 {
-                if transmit.num_datagrams() >= transmit.max_datagrams() {
+                if transmit.num_datagrams() >= transmit.max_datagrams().get() {
                     // No more datagrams allowed
                     break;
                 }

quinn-proto/src/connection/transmit_buf.rs

@@ -1,3 +1,5 @@
+use std::num::NonZeroUsize;
+
 use bytes::BufMut;
 use tracing::trace;
 
@@ -40,7 +42,7 @@ pub(super) struct TransmitBuf<'a> {
     /// size. All datagrams in between need to be exactly this size.
     buf_capacity: usize,
     /// The maximum number of datagrams allowed to write into [`TransmitBuf::buf`]
-    max_datagrams: usize,
+    max_datagrams: NonZeroUsize,
     /// The number of datagrams already (partially) written into the buffer
     ///
     /// Incremented by a call to [`TransmitBuf::start_new_datagram`].
@@ -57,7 +59,7 @@ pub(super) struct TransmitBuf<'a> {
 }
 
 impl<'a> TransmitBuf<'a> {
-    pub(super) fn new(buf: &'a mut Vec<u8>, max_datagrams: usize, mtu: usize) -> Self {
+    pub(super) fn new(buf: &'a mut Vec<u8>, max_datagrams: NonZeroUsize, mtu: usize) -> Self {
         Self {
             buf,
             datagram_start: 0,
@@ -106,12 +108,12 @@ impl<'a> TransmitBuf<'a> {
         // (e.g. purely containing ACKs), modern memory allocators (e.g. mimalloc and
         // jemalloc) will pool certain allocation sizes and therefore this is still rather
         // efficient.
-        let max_capacity_hint = self.max_datagrams * self.segment_size;
+        let max_capacity_hint = self.max_datagrams.get() * self.segment_size;
         self.new_datagram_inner(self.segment_size, max_capacity_hint)
     }
 
     fn new_datagram_inner(&mut self, datagram_size: usize, max_capacity_hint: usize) {
-        debug_assert!(self.num_datagrams < self.max_datagrams);
+        debug_assert!(self.num_datagrams < self.max_datagrams.into());
         if self.num_datagrams == 1 {
             // Set the segment size to the size of the first datagram.
             self.segment_size = self.buf.len();
@@ -121,7 +123,7 @@ impl<'a> TransmitBuf<'a> {
             if datagram_size < self.segment_size {
                 // If this is a GSO batch and this datagram is smaller than the segment
                 // size, this must be the last datagram in the batch.
-                self.max_datagrams = self.num_datagrams + 1;
+                self.max_datagrams = NonZeroUsize::new(self.num_datagrams + 1).expect("known");
             }
         }
         self.datagram_start = self.buf.len();
@@ -181,7 +183,7 @@ impl<'a> TransmitBuf<'a> {
     }
 
     /// Returns the maximum number of datagrams allowed to be written into the buffer
-    pub(super) fn max_datagrams(&self) -> usize {
+    pub(super) fn max_datagrams(&self) -> NonZeroUsize {
         self.max_datagrams
     }
 

quinn-proto/src/tests/util.rs

@@ -5,6 +5,7 @@ use std::{
     io::{self, Write},
     mem,
     net::{Ipv6Addr, SocketAddr, UdpSocket},
+    num::NonZeroUsize,
     ops::RangeFrom,
     str,
     sync::{Arc, LazyLock, Mutex},
@@ -791,7 +792,7 @@ pub(super) fn min_opt<T: Ord>(x: Option<T>, y: Option<T>) -> Option<T> {
 }
 
 /// The maximum of datagrams TestEndpoint will produce via `poll_transmit`
-const MAX_DATAGRAMS: usize = 10;
+const MAX_DATAGRAMS: NonZeroUsize = NonZeroUsize::new(10).expect("known");
 
 fn split_transmit(transmit: Transmit, buffer: &[u8]) -> Vec<(Transmit, Bytes)> {
     let mut buffer = Bytes::copy_from_slice(buffer);

quinn-udp/benches/throughput.rs

@@ -2,6 +2,7 @@ use std::{
     cmp::min,
     io::{ErrorKind, IoSliceMut},
     net::{Ipv4Addr, Ipv6Addr, UdpSocket},
+    num::NonZeroUsize,
 };
 
 use criterion::{Criterion, criterion_group, criterion_main};
@@ -54,9 +55,9 @@ pub fn criterion_benchmark(c: &mut Criterion) {
         let gso_segments = if gso_enabled {
             send_state.max_gso_segments()
         } else {
-            1
+            NonZeroUsize::new(1).expect("known")
         };
-        let msg = vec![0xAB; min(MAX_DATAGRAM_SIZE, SEGMENT_SIZE * gso_segments)];
+        let msg = vec![0xAB; min(MAX_DATAGRAM_SIZE, SEGMENT_SIZE * gso_segments.get())];
         let transmit = Transmit {
             destination: dst_addr,
             ecn: None,
@@ -67,13 +68,14 @@ pub fn criterion_benchmark(c: &mut Criterion) {
         let gro_segments = if gro_enabled {
             recv_state.gro_segments()
         } else {
-            1
+            NonZeroUsize::new(1).expect("known")
         };
         let batch_size = if recvmmsg_enabled { BATCH_SIZE } else { 1 };
 
         group.bench_function("throughput", |b| {
             b.to_async(&rt).iter(|| async {
-                let mut receive_buffers = vec![vec![0; SEGMENT_SIZE * gro_segments]; batch_size];
+                let mut receive_buffers =
+                    vec![vec![0; SEGMENT_SIZE * gro_segments.get()]; batch_size];
                 let mut receive_slices = receive_buffers
                     .iter_mut()
                     .map(|buf| IoSliceMut::new(buf))

quinn-udp/src/fallback.rs

@@ -83,8 +83,8 @@ impl UdpSocketState {
     }
 
     #[inline]
-    pub fn gro_segments(&self) -> usize {
-        1
+    pub fn gro_segments(&self) -> NonZeroUsize {
+        NonZeroUsize::new(1).expect("known")
     }
 
     /// Resize the send buffer of `socket` to `bytes`

quinn-udp/src/unix.rs

@@ -4,6 +4,7 @@ use std::{
     io::{self, IoSliceMut},
     mem::{self, MaybeUninit},
     net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6},
+    num::NonZeroUsize,
     os::unix::io::AsRawFd,
     sync::{
         Mutex,
@@ -71,7 +72,7 @@ type IpTosTy = libc::c_int;
 pub struct UdpSocketState {
     last_send_error: Mutex<Instant>,
     max_gso_segments: AtomicUsize,
-    gro_segments: usize,
+    gro_segments: NonZeroUsize,
     may_fragment: bool,
 
     /// True if we have received EINVAL error from `sendmsg` system call at least once.
@@ -244,8 +245,11 @@ impl UdpSocketState {
     /// This is 1 if the platform doesn't support GSO. Subject to change if errors are detected
     /// while using GSO.
     #[inline]
-    pub fn max_gso_segments(&self) -> usize {
-        self.max_gso_segments.load(Ordering::Relaxed)
+    pub fn max_gso_segments(&self) -> NonZeroUsize {
+        self.max_gso_segments
+            .load(Ordering::Relaxed)
+            .try_into()
+            .expect("must have non zero GSO segments")
     }
 
     /// The number of segments to read when GRO is enabled. Used as a factor to
@@ -253,7 +257,7 @@ impl UdpSocketState {
     ///
     /// Returns 1 if the platform doesn't support GRO.
     #[inline]
-    pub fn gro_segments(&self) -> usize {
+    pub fn gro_segments(&self) -> NonZeroUsize {
         self.gro_segments
     }
 
@@ -1023,12 +1027,12 @@ mod gro {
     // TODO: Add this to libc
     pub(crate) const UDP_GRO: libc::c_int = 104;
 
-    pub(crate) fn gro_segments() -> usize {
+    pub(crate) fn gro_segments() -> NonZeroUsize {
         let socket = match std::net::UdpSocket::bind("[::]:0")
             .or_else(|_| std::net::UdpSocket::bind((Ipv4Addr::LOCALHOST, 0)))
         {
             Ok(socket) => socket,
-            Err(_) => return 1,
+            Err(_) => return NonZeroUsize::new(1).expect("known"),
         };
 
         // As defined in net/ipv4/udp_offload.c
@@ -1039,8 +1043,8 @@ mod gro {
         // list the kernel might potentially produce. See
         // https://github.com/quinn-rs/quinn/pull/1354.
         match set_socket_option(&socket, libc::SOL_UDP, UDP_GRO, OPTION_ON) {
-            Ok(()) => 64,
-            Err(_) => 1,
+            Ok(()) => NonZeroUsize::new(64).expect("known"),
+            Err(_) => NonZeroUsize::new(1).expect("known"),
         }
     }
 }
@@ -1089,7 +1093,9 @@ const OPTION_ON: libc::c_int = 1;
 
 #[cfg(not(any(target_os = "linux", target_os = "android")))]
 mod gro {
-    pub(super) fn gro_segments() -> usize {
-        1
+    use std::num::NonZeroUsize;
+
+    pub(super) fn gro_segments() -> NonZeroUsize {
+        NonZeroUsize::new(1).expect("known")
     }
 }

quinn-udp/src/windows.rs

@@ -295,9 +295,9 @@ impl UdpSocketState {
     ///
     /// Returns 1 if the platform doesn't support GRO.
     #[inline]
-    pub fn gro_segments(&self) -> usize {
+    pub fn gro_segments(&self) -> NonZeroUsize {
         // Arbitrary reasonable value inspired by Linux and msquic
-        64
+        NonZeroUsize::new(64).expect("known")
     }
 
     /// Resize the send buffer of `socket` to `bytes`

quinn-udp/tests/tests.rs

@@ -195,7 +195,7 @@ fn gso() {
         .max_gso_segments();
     let dst_addr = recv.local_addr().unwrap();
     const SEGMENT_SIZE: usize = 128;
-    let msg = vec![0xAB; SEGMENT_SIZE * max_segments];
+    let msg = vec![0xAB; SEGMENT_SIZE * max_segments.get()];
     test_send_recv(
         &send.into(),
         &recv.into(),
@@ -293,7 +293,7 @@ fn test_send_recv(send: &Socket, recv: &Socket, transmit: Transmit) {
         Ok(_) => (),
         Err(err) if err.kind() == io::ErrorKind::InvalidInput && cfg!(target_os = "windows") => {
             // GSO can fail on windows. It should have disabled GSO now.
-            assert_eq!(send_state.max_gso_segments(), 1);
+            assert_eq!(send_state.max_gso_segments().get(), 1);
             return;
         }
         Err(err) => panic!("send failed: {err:?}"),

quinn/src/connection.rs

@@ -4,6 +4,7 @@ use std::{
     future::Future,
     io,
     net::{IpAddr, SocketAddr},
+    num::NonZeroUsize,
     pin::Pin,
     sync::{Arc, Weak},
     task::{Context, Poll, Waker, ready},
@@ -1728,4 +1729,4 @@ const MAX_TRANSMIT_DATAGRAMS: usize = 20;
 /// This can be lower than the maximum platform capabilities, to avoid excessive
 /// memory allocations when calling `poll_transmit()`. Benchmarks have shown
 /// that numbers around 10 are a good compromise.
-const MAX_TRANSMIT_SEGMENTS: usize = 10;
+const MAX_TRANSMIT_SEGMENTS: NonZeroUsize = NonZeroUsize::new(10).expect("known");

quinn/src/endpoint.rs

@@ -5,6 +5,7 @@ use std::{
     io::{self, IoSliceMut},
     mem,
     net::{SocketAddr, SocketAddrV6},
+    num::NonZeroUsize,
     pin::Pin,
     str,
     sync::{Arc, Mutex},
@@ -798,13 +799,13 @@ struct RecvState {
 impl RecvState {
     fn new(
         sender: mpsc::UnboundedSender<(ConnectionHandle, EndpointEvent)>,
-        max_receive_segments: usize,
+        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
+                * max_receive_segments.get()
                 * BATCH_SIZE
         ];
         Self {

quinn/src/runtime/mod.rs

@@ -5,6 +5,7 @@ use std::{
     future::Future,
     io::{self, IoSliceMut},
     net::SocketAddr,
+    num::NonZeroUsize,
     pin::Pin,
     task::{Context, Poll},
 };
@@ -65,8 +66,8 @@ pub trait AsyncUdpSocket: Send + Sync + Debug + 'static {
     fn local_addr(&self) -> io::Result<SocketAddr>;
 
     /// Maximum number of datagrams that might be described by a single [`RecvMeta`]
-    fn max_receive_segments(&self) -> usize {
-        1
+    fn max_receive_segments(&self) -> NonZeroUsize {
+        NonZeroUsize::new(1).expect("known")
     }
 
     /// Whether datagrams might get fragmented into multiple parts
@@ -100,8 +101,8 @@ pub trait UdpSender: Send + Sync + Debug + 'static {
     ) -> Poll<io::Result<()>>;
 
     /// Maximum number of datagrams that a [`Transmit`] may encode.
-    fn max_transmit_segments(&self) -> usize {
-        1
+    fn max_transmit_segments(&self) -> NonZeroUsize {
+        NonZeroUsize::new(1).expect("known")
     }
 }
 
@@ -195,7 +196,7 @@ where
         }
     }
 
-    fn max_transmit_segments(&self) -> usize {
+    fn max_transmit_segments(&self) -> NonZeroUsize {
         self.socket.max_transmit_segments()
     }
 }
@@ -212,7 +213,7 @@ trait UdpSenderHelperSocket: Send + Sync + 'static {
     fn try_send(&self, transmit: &udp::Transmit) -> io::Result<()>;
 
     /// See [`UdpSender::max_transmit_segments`].
-    fn max_transmit_segments(&self) -> usize;
+    fn max_transmit_segments(&self) -> NonZeroUsize;
 }
 
 /// Automatically select an appropriate runtime from those enabled at compile time

quinn/src/runtime/smol.rs

@@ -1,5 +1,6 @@
 use std::{
     future::Future,
+    num::NonZeroUsize,
     pin::Pin,
     task::{Context, Poll},
     time::Instant,
@@ -56,7 +57,7 @@ impl UdpSocket {
 }
 
 impl UdpSenderHelperSocket for UdpSocket {
-    fn max_transmit_segments(&self) -> usize {
+    fn max_transmit_segments(&self) -> NonZeroUsize {
         self.inner.max_gso_segments()
     }
 
@@ -95,7 +96,7 @@ impl AsyncUdpSocket for UdpSocket {
         self.inner.may_fragment()
     }
 
-    fn max_receive_segments(&self) -> usize {
+    fn max_receive_segments(&self) -> NonZeroUsize {
         self.inner.gro_segments()
     }
 }

quinn/src/runtime/tokio.rs

@@ -2,6 +2,7 @@ use std::{
     fmt::Debug,
     future::Future,
     io,
+    num::NonZeroUsize,
     pin::Pin,
     sync::Arc,
     task::{Context, Poll, ready},
@@ -56,7 +57,7 @@ struct UdpSocket {
 }
 
 impl UdpSenderHelperSocket for UdpSocket {
-    fn max_transmit_segments(&self) -> usize {
+    fn max_transmit_segments(&self) -> NonZeroUsize {
         self.inner.max_gso_segments()
     }
 
@@ -99,7 +100,7 @@ impl AsyncUdpSocket for UdpSocket {
         self.inner.may_fragment()
     }
 
-    fn max_receive_segments(&self) -> usize {
+    fn max_receive_segments(&self) -> NonZeroUsize {
         self.inner.gro_segments()
     }
 }