From 40fffcb74bb149fc640d1a96371bee73e3dddf74 Mon Sep 17 00:00:00 2001 From: dignifiedquire Date: Fri, 12 Dec 2025 16:43:54 +0100 Subject: [PATCH] refactor!: use NonZeroUsize for segments sizes Start introducing more type safety for implicit assumptions, step by step --- quinn-proto/src/connection/mod.rs | 9 ++++---- quinn-proto/src/connection/transmit_buf.rs | 14 +++++++----- quinn-proto/src/tests/util.rs | 3 ++- quinn-udp/benches/throughput.rs | 10 +++++---- quinn-udp/src/fallback.rs | 4 ++-- quinn-udp/src/unix.rs | 26 +++++++++++++--------- quinn-udp/src/windows.rs | 4 ++-- quinn-udp/tests/tests.rs | 4 ++-- quinn/src/connection.rs | 3 ++- quinn/src/endpoint.rs | 5 +++-- quinn/src/runtime/mod.rs | 13 ++++++----- quinn/src/runtime/smol.rs | 5 +++-- quinn/src/runtime/tokio.rs | 5 +++-- 13 files changed, 60 insertions(+), 45 deletions(-) diff --git a/quinn-proto/src/connection/mod.rs b/quinn-proto/src/connection/mod.rs index e996be829..16e391624 100644 --- a/quinn-proto/src/connection/mod.rs +++ b/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, ) -> Option { 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; } diff --git a/quinn-proto/src/connection/transmit_buf.rs b/quinn-proto/src/connection/transmit_buf.rs index d12d1c8cc..42930ab1a 100644 --- a/quinn-proto/src/connection/transmit_buf.rs +++ b/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, max_datagrams: usize, mtu: usize) -> Self { + pub(super) fn new(buf: &'a mut Vec, 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 } diff --git a/quinn-proto/src/tests/util.rs b/quinn-proto/src/tests/util.rs index 10940739a..1a3ac0ade 100644 --- a/quinn-proto/src/tests/util.rs +++ b/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(x: Option, y: Option) -> Option { } /// 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); diff --git a/quinn-udp/benches/throughput.rs b/quinn-udp/benches/throughput.rs index 6c429817b..9c4ddf88f 100644 --- a/quinn-udp/benches/throughput.rs +++ b/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)) diff --git a/quinn-udp/src/fallback.rs b/quinn-udp/src/fallback.rs index 1f50d7034..f89a77d25 100644 --- a/quinn-udp/src/fallback.rs +++ b/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` diff --git a/quinn-udp/src/unix.rs b/quinn-udp/src/unix.rs index 393a361b2..c1d711357 100644 --- a/quinn-udp/src/unix.rs +++ b/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, 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") } } diff --git a/quinn-udp/src/windows.rs b/quinn-udp/src/windows.rs index d6eee6f1f..55018c698 100644 --- a/quinn-udp/src/windows.rs +++ b/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` diff --git a/quinn-udp/tests/tests.rs b/quinn-udp/tests/tests.rs index de97499cb..00067e2d7 100644 --- a/quinn-udp/tests/tests.rs +++ b/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:?}"), diff --git a/quinn/src/connection.rs b/quinn/src/connection.rs index 646cd44f5..04a48e561 100644 --- a/quinn/src/connection.rs +++ b/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"); diff --git a/quinn/src/endpoint.rs b/quinn/src/endpoint.rs index 496a75849..b62aeb98c 100644 --- a/quinn/src/endpoint.rs +++ b/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 { diff --git a/quinn/src/runtime/mod.rs b/quinn/src/runtime/mod.rs index 66c35ca15..6100df149 100644 --- a/quinn/src/runtime/mod.rs +++ b/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; /// 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>; /// 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 diff --git a/quinn/src/runtime/smol.rs b/quinn/src/runtime/smol.rs index f9bff2194..20a7768bf 100644 --- a/quinn/src/runtime/smol.rs +++ b/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() } } diff --git a/quinn/src/runtime/tokio.rs b/quinn/src/runtime/tokio.rs index 3fa7a27b1..e2da37b8a 100644 --- a/quinn/src/runtime/tokio.rs +++ b/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() } }