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

@@ -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<ConnectionEvent>,
-        udp_state: Arc<UdpState>,
+        socket: Arc<dyn AsyncUdpSocket>,
         runtime: Arc<dyn Runtime>,
     ) -> 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<ConnectionEvent>,
         on_handshake_data: oneshot::Sender<()>,
         on_connected: oneshot::Sender<bool>,
-        udp_state: Arc<UdpState>,
+        socket: Arc<dyn AsyncUdpSocket>,
         runtime: Arc<dyn Runtime>,
     ) -> 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<ConnectionError>,
     /// Number of live handles that can be used to initiate or handle I/O; excludes the driver
     ref_count: usize,
-    udp_state: Arc<UdpState>,
+    socket: Arc<dyn AsyncUdpSocket>,
     runtime: Arc<dyn Runtime>,
 }
 
@@ -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 {

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<ServerConfig>,
-        socket: Box<dyn AsyncUdpSocket>,
+        socket: Arc<dyn AsyncUdpSocket>,
         runtime: Arc<dyn Runtime>,
     ) -> io::Result<Self> {
         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<dyn AsyncUdpSocket>,
-    udp_state: Arc<UdpState>,
+    socket: Arc<dyn AsyncUdpSocket>,
     inner: proto::Endpoint,
     outgoing: VecDeque<udp::Transmit>,
     incoming: VecDeque<Connecting>,
@@ -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<UdpState>,
+        socket: Arc<dyn AsyncUdpSocket>,
         runtime: Arc<dyn Runtime>,
     ) -> 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<EndpointInner>);
 
 impl EndpointRef {
     pub(crate) fn new(
-        socket: Box<dyn AsyncUdpSocket>,
+        socket: Arc<dyn AsyncUdpSocket>,
         inner: proto::Endpoint,
         ipv6: bool,
         runtime: Arc<dyn Runtime>,
     ) -> 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,

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<Box<dyn Future<Output = ()> + Send>>);
     /// Convert `t` into the socket type used by this runtime
-    fn wrap_udp_socket(&self, t: std::net::UdpSocket) -> io::Result<Box<dyn AsyncUdpSocket>>;
+    fn wrap_udp_socket(&self, t: std::net::UdpSocket) -> io::Result<Arc<dyn AsyncUdpSocket>>;
 }
 
 /// 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<Result<usize, io::Error>>;
+    fn poll_send(&self, cx: &mut Context, transmits: &[Transmit])
+        -> Poll<Result<usize, io::Error>>;
 
     /// 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<SocketAddr>;
 
+    /// 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

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<Box<dyn AsyncUdpSocket>> {
+    fn wrap_udp_socket(&self, sock: std::net::UdpSocket) -> io::Result<Arc<dyn AsyncUdpSocket>> {
         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<std::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<io::Result<usize>> {
+    fn poll_send(&self, cx: &mut Context, transmits: &[udp::Transmit]) -> Poll<io::Result<usize>> {
         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()
+    }
 }

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<Box<dyn AsyncUdpSocket>> {
+    fn wrap_udp_socket(&self, sock: std::net::UdpSocket) -> io::Result<Arc<dyn AsyncUdpSocket>> {
         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<io::Result<usize>> {
+    fn poll_send(&self, cx: &mut Context, transmits: &[udp::Transmit]) -> Poll<io::Result<usize>> {
         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()
+    }
 }