From 3eb23bc730bc2ead0e560a8b1eaa49f7bebd28d6 Mon Sep 17 00:00:00 2001 From: Benjamin Saunders Date: Fri, 21 Jul 2023 15:38:44 -0700 Subject: [PATCH] Fold UdpState into AsyncUdpSocket Allows knowledge of UdpState to be isolated entirely within AsyncUdpSocket implementations, simplifying quinn::Endpoint and the poll_send API, and exposing more control over UDP feature checks to implementers. --- quinn/src/connection.rs | 15 +++++++-------- quinn/src/endpoint.rs | 28 +++++++++++----------------- quinn/src/runtime.rs | 24 +++++++++++++++--------- quinn/src/runtime/async_std.rs | 24 +++++++++++++++--------- quinn/src/runtime/tokio.rs | 24 +++++++++++++++--------- 5 files changed, 63 insertions(+), 52 deletions(-) diff --git a/quinn/src/connection.rs b/quinn/src/connection.rs index d21313f9e..4be517368 100644 --- a/quinn/src/connection.rs +++ b/quinn/src/connection.rs @@ -9,7 +9,7 @@ use std::{ time::{Duration, Instant}, }; -use crate::runtime::{AsyncTimer, Runtime}; +use crate::runtime::{AsyncTimer, AsyncUdpSocket, Runtime}; use bytes::Bytes; use pin_project_lite::pin_project; use proto::{ConnectionError, ConnectionHandle, ConnectionStats, Dir, StreamEvent, StreamId}; @@ -17,7 +17,6 @@ use rustc_hash::FxHashMap; use thiserror::Error; use tokio::sync::{futures::Notified, mpsc, oneshot, Notify}; use tracing::debug_span; -use udp::UdpState; use crate::{ mutex::Mutex, @@ -42,7 +41,7 @@ impl Connecting { conn: proto::Connection, endpoint_events: mpsc::UnboundedSender<(ConnectionHandle, EndpointEvent)>, conn_events: mpsc::UnboundedReceiver, - udp_state: Arc, + socket: Arc, runtime: Arc, ) -> Self { let (on_handshake_data_send, on_handshake_data_recv) = oneshot::channel(); @@ -54,7 +53,7 @@ impl Connecting { conn_events, on_handshake_data_send, on_connected_send, - udp_state, + socket, runtime.clone(), ); @@ -747,7 +746,7 @@ impl ConnectionRef { conn_events: mpsc::UnboundedReceiver, on_handshake_data: oneshot::Sender<()>, on_connected: oneshot::Sender, - udp_state: Arc, + socket: Arc, runtime: Arc, ) -> Self { Self(Arc::new(ConnectionInner { @@ -768,7 +767,7 @@ impl ConnectionRef { stopped: FxHashMap::default(), error: None, ref_count: 0, - udp_state, + socket, runtime, }), shared: Shared::default(), @@ -846,7 +845,7 @@ pub(crate) struct State { pub(crate) error: Option, /// Number of live handles that can be used to initiate or handle I/O; excludes the driver ref_count: usize, - udp_state: Arc, + socket: Arc, runtime: Arc, } @@ -855,7 +854,7 @@ impl State { let now = Instant::now(); let mut transmits = 0; - let max_datagrams = self.udp_state.max_gso_segments(); + let max_datagrams = self.socket.max_transmit_segments(); while let Some(t) = self.inner.poll_transmit(now, max_datagrams) { transmits += match t.segment_size { diff --git a/quinn/src/endpoint.rs b/quinn/src/endpoint.rs index d5e1e2406..1aa26f77d 100644 --- a/quinn/src/endpoint.rs +++ b/quinn/src/endpoint.rs @@ -20,7 +20,7 @@ use proto::{ }; use rustc_hash::FxHashMap; use tokio::sync::{futures::Notified, mpsc, Notify}; -use udp::{RecvMeta, UdpState, BATCH_SIZE}; +use udp::{RecvMeta, BATCH_SIZE}; use crate::{ connection::Connecting, work_limiter::WorkLimiter, ConnectionEvent, EndpointConfig, @@ -102,7 +102,7 @@ impl Endpoint { pub fn new_with_abstract_socket( config: EndpointConfig, server_config: Option, - socket: Box, + socket: Arc, runtime: Arc, ) -> io::Result { let addr = socket.local_addr()?; @@ -183,10 +183,10 @@ impl Endpoint { addr }; let (ch, conn) = endpoint.inner.connect(config, addr, server_name)?; - let udp_state = endpoint.udp_state.clone(); + let socket = endpoint.socket.clone(); Ok(endpoint .connections - .insert(ch, conn, udp_state, self.runtime.clone())) + .insert(ch, conn, socket, self.runtime.clone())) } /// Switch to a new UDP socket @@ -355,8 +355,7 @@ pub(crate) struct EndpointInner { #[derive(Debug)] pub(crate) struct State { - socket: Box, - udp_state: Arc, + socket: Arc, inner: proto::Endpoint, outgoing: VecDeque, incoming: VecDeque, @@ -415,7 +414,7 @@ impl State { let conn = self.connections.insert( handle, conn, - self.udp_state.clone(), + self.socket.clone(), self.runtime.clone(), ); self.incoming.push_back(conn); @@ -483,10 +482,7 @@ impl State { break Ok(true); } - match self - .socket - .poll_send(&self.udp_state, cx, self.outgoing.as_slices().0) - { + match self.socket.poll_send(cx, self.outgoing.as_slices().0) { Poll::Ready(Ok(n)) => { let contents_len: usize = self.outgoing.drain(..n).map(|t| t.contents.len()).sum(); @@ -596,7 +592,7 @@ impl ConnectionSet { &mut self, handle: ConnectionHandle, conn: proto::Connection, - udp_state: Arc, + socket: Arc, runtime: Arc, ) -> Connecting { let (send, recv) = mpsc::unbounded_channel(); @@ -608,7 +604,7 @@ impl ConnectionSet { .unwrap(); } self.senders.insert(handle, send); - Connecting::new(handle, conn, self.sender.clone(), recv, udp_state, runtime) + Connecting::new(handle, conn, self.sender.clone(), recv, socket, runtime) } fn is_empty(&self) -> bool { @@ -664,16 +660,15 @@ pub(crate) struct EndpointRef(Arc); impl EndpointRef { pub(crate) fn new( - socket: Box, + socket: Arc, inner: proto::Endpoint, ipv6: bool, runtime: Arc, ) -> Self { - let udp_state = Arc::new(UdpState::new()); let recv_buf = vec![ 0; inner.config().get_max_udp_payload_size().min(64 * 1024) as usize - * udp_state.gro_segments() + * socket.max_receive_segments() * BATCH_SIZE ]; let (sender, events) = mpsc::unbounded_channel(); @@ -684,7 +679,6 @@ impl EndpointRef { }, state: Mutex::new(State { socket, - udp_state, inner, ipv6, events, diff --git a/quinn/src/runtime.rs b/quinn/src/runtime.rs index 9bcb73600..d4d5bd3f8 100644 --- a/quinn/src/runtime.rs +++ b/quinn/src/runtime.rs @@ -9,7 +9,7 @@ use std::{ time::Instant, }; -use udp::{RecvMeta, Transmit, UdpState}; +use udp::{RecvMeta, Transmit}; /// Abstracts I/O and timer operations for runtime independence pub trait Runtime: Send + Sync + Debug + 'static { @@ -18,7 +18,7 @@ pub trait Runtime: Send + Sync + Debug + 'static { /// Drive `future` to completion in the background fn spawn(&self, future: Pin + Send>>); /// Convert `t` into the socket type used by this runtime - fn wrap_udp_socket(&self, t: std::net::UdpSocket) -> io::Result>; + fn wrap_udp_socket(&self, t: std::net::UdpSocket) -> io::Result>; } /// Abstract implementation of an async timer for runtime independence @@ -30,15 +30,11 @@ pub trait AsyncTimer: Send + Debug + 'static { } /// Abstract implementation of a UDP socket for runtime independence -pub trait AsyncUdpSocket: Send + Debug + 'static { +pub trait AsyncUdpSocket: Send + Sync + Debug + 'static { /// Send UDP datagrams from `transmits`, or register to be woken if sending may succeed in the /// future - fn poll_send( - &self, - state: &UdpState, - cx: &mut Context, - transmits: &[Transmit], - ) -> Poll>; + fn poll_send(&self, cx: &mut Context, transmits: &[Transmit]) + -> Poll>; /// Receive UDP datagrams, or register to be woken if receiving may succeed in the future fn poll_recv( @@ -51,6 +47,16 @@ pub trait AsyncUdpSocket: Send + Debug + 'static { /// Look up the local IP address and port used by this socket fn local_addr(&self) -> io::Result; + /// Maximum number of datagrams that a [`Transmit`] may encode + fn max_transmit_segments(&self) -> usize { + 1 + } + + /// Maximum number of datagrams that might be described by a single [`RecvMeta`] + fn max_receive_segments(&self) -> usize { + 1 + } + /// Whether datagrams might get fragmented into multiple parts /// /// Sockets should prevent this for best performance. See e.g. the `IPV6_DONTFRAG` socket diff --git a/quinn/src/runtime/async_std.rs b/quinn/src/runtime/async_std.rs index da01201b7..98d8821db 100644 --- a/quinn/src/runtime/async_std.rs +++ b/quinn/src/runtime/async_std.rs @@ -2,6 +2,7 @@ use std::{ future::Future, io, pin::Pin, + sync::Arc, task::{Context, Poll}, time::Instant, }; @@ -23,10 +24,11 @@ impl Runtime for AsyncStdRuntime { async_std::task::spawn(future); } - fn wrap_udp_socket(&self, sock: std::net::UdpSocket) -> io::Result> { + fn wrap_udp_socket(&self, sock: std::net::UdpSocket) -> io::Result> { udp::UdpSocketState::configure((&sock).into())?; - Ok(Box::new(UdpSocket { + Ok(Arc::new(UdpSocket { io: Async::new(sock)?, + state: udp::UdpState::new(), inner: udp::UdpSocketState::new(), })) } @@ -45,19 +47,15 @@ impl AsyncTimer for Timer { #[derive(Debug)] struct UdpSocket { io: Async, + state: udp::UdpState, inner: udp::UdpSocketState, } impl AsyncUdpSocket for UdpSocket { - fn poll_send( - &self, - state: &udp::UdpState, - cx: &mut Context, - transmits: &[udp::Transmit], - ) -> Poll> { + fn poll_send(&self, cx: &mut Context, transmits: &[udp::Transmit]) -> Poll> { loop { ready!(self.io.poll_writable(cx))?; - if let Ok(res) = self.inner.send((&self.io).into(), state, transmits) { + if let Ok(res) = self.inner.send((&self.io).into(), &self.state, transmits) { return Poll::Ready(Ok(res)); } } @@ -84,4 +82,12 @@ impl AsyncUdpSocket for UdpSocket { fn may_fragment(&self) -> bool { udp::may_fragment() } + + fn max_transmit_segments(&self) -> usize { + self.state.max_gso_segments() + } + + fn max_receive_segments(&self) -> usize { + self.state.gro_segments() + } } diff --git a/quinn/src/runtime/tokio.rs b/quinn/src/runtime/tokio.rs index d8ac4ac10..5c71fdea3 100644 --- a/quinn/src/runtime/tokio.rs +++ b/quinn/src/runtime/tokio.rs @@ -2,6 +2,7 @@ use std::{ future::Future, io, pin::Pin, + sync::Arc, task::{Context, Poll}, time::Instant, }; @@ -26,10 +27,11 @@ impl Runtime for TokioRuntime { tokio::spawn(future); } - fn wrap_udp_socket(&self, sock: std::net::UdpSocket) -> io::Result> { + fn wrap_udp_socket(&self, sock: std::net::UdpSocket) -> io::Result> { udp::UdpSocketState::configure((&sock).into())?; - Ok(Box::new(UdpSocket { + Ok(Arc::new(UdpSocket { io: tokio::net::UdpSocket::from_std(sock)?, + state: udp::UdpState::new(), inner: udp::UdpSocketState::new(), })) } @@ -47,22 +49,18 @@ impl AsyncTimer for Sleep { #[derive(Debug)] struct UdpSocket { io: tokio::net::UdpSocket, + state: udp::UdpState, inner: udp::UdpSocketState, } impl AsyncUdpSocket for UdpSocket { - fn poll_send( - &self, - state: &udp::UdpState, - cx: &mut Context, - transmits: &[udp::Transmit], - ) -> Poll> { + fn poll_send(&self, cx: &mut Context, transmits: &[udp::Transmit]) -> Poll> { let inner = &self.inner; let io = &self.io; loop { ready!(io.poll_send_ready(cx))?; if let Ok(res) = io.try_io(Interest::WRITABLE, || { - inner.send(io.into(), state, transmits) + inner.send(io.into(), &self.state, transmits) }) { return Poll::Ready(Ok(res)); } @@ -92,4 +90,12 @@ impl AsyncUdpSocket for UdpSocket { fn may_fragment(&self) -> bool { udp::may_fragment() } + + fn max_transmit_segments(&self) -> usize { + self.state.max_gso_segments() + } + + fn max_receive_segments(&self) -> usize { + self.state.gro_segments() + } }