Store a Box<dyn AsyncUdpSocket> instead of Arcing it, make poll_recv take &mut self

quinn/src/endpoint.rs

@@ -129,7 +129,7 @@ impl Endpoint {
     pub fn new_with_abstract_socket(
         config: EndpointConfig,
         server_config: Option<ServerConfig>,
-        socket: Arc<dyn AsyncUdpSocket>,
+        socket: Box<dyn AsyncUdpSocket>,
         runtime: Arc<dyn Runtime>,
     ) -> io::Result<Self> {
         let addr = socket.local_addr()?;
@@ -224,7 +224,7 @@ impl Endpoint {
             .inner
             .connect(self.runtime.now(), config, addr, server_name)?;
 
-        let sender = endpoint.socket.clone().create_sender();
+        let sender = endpoint.socket.create_sender();
         endpoint.stats.outgoing_handshakes += 1;
         Ok(endpoint
             .recv_state
@@ -246,7 +246,7 @@ impl Endpoint {
     /// connections and connections to servers unreachable from the new address will be lost.
     ///
     /// On error, the old UDP socket is retained.
-    pub fn rebind_abstract(&self, socket: Arc<dyn AsyncUdpSocket>) -> io::Result<()> {
+    pub fn rebind_abstract(&self, socket: Box<dyn AsyncUdpSocket>) -> io::Result<()> {
         let addr = socket.local_addr()?;
         let mut inner = self.inner.state.lock().unwrap();
         inner.prev_socket = Some(mem::replace(&mut inner.socket, socket));
@@ -255,9 +255,7 @@ impl Endpoint {
         // Update connection socket references
         for sender in inner.recv_state.connections.senders.values() {
             // Ignoring errors from dropped connections
-            let _ = sender.send(ConnectionEvent::Rebind(
-                inner.socket.clone().create_sender(),
-            ));
+            let _ = sender.send(ConnectionEvent::Rebind(inner.socket.create_sender()));
         }
         if let Some(driver) = inner.driver.take() {
             // Ensure the driver can register for wake-ups from the new socket
@@ -426,7 +424,7 @@ impl EndpointInner {
         {
             Ok((handle, conn)) => {
                 state.stats.accepted_handshakes += 1;
-                let sender = state.socket.clone().create_sender();
+                let sender = state.socket.create_sender();
                 let runtime = state.runtime.clone();
                 Ok(state
                     .recv_state
@@ -467,11 +465,11 @@ impl EndpointInner {
 
 #[derive(Debug)]
 pub(crate) struct State {
-    socket: Arc<dyn AsyncUdpSocket>,
+    socket: Box<dyn AsyncUdpSocket>,
     sender: Pin<Box<dyn UdpSender>>,
     /// During an active migration, abandoned_socket receives traffic
     /// until the first packet arrives on the new socket.
-    prev_socket: Option<Arc<dyn AsyncUdpSocket>>,
+    prev_socket: Option<Box<dyn AsyncUdpSocket>>,
     inner: proto::Endpoint,
     recv_state: RecvState,
     driver: Option<Waker>,
@@ -494,12 +492,12 @@ impl State {
     fn drive_recv(&mut self, cx: &mut Context, now: Instant) -> Result<bool, io::Error> {
         let get_time = || self.runtime.now();
         self.recv_state.recv_limiter.start_cycle(get_time);
-        if let Some(socket) = &self.prev_socket {
+        if let Some(socket) = &mut self.prev_socket {
             // We don't care about the `PollProgress` from old sockets.
             let poll_res = self.recv_state.poll_socket(
                 cx,
                 &mut self.inner,
-                &**socket,
+                &mut **socket,
                 &mut self.sender,
                 &*self.runtime,
                 now,
@@ -511,7 +509,7 @@ impl State {
         let poll_res = self.recv_state.poll_socket(
             cx,
             &mut self.inner,
-            &*self.socket,
+            &mut *self.socket,
             &mut self.sender,
             &*self.runtime,
             now,
@@ -712,14 +710,14 @@ pub(crate) struct EndpointRef(Arc<EndpointInner>);
 
 impl EndpointRef {
     pub(crate) fn new(
-        socket: Arc<dyn AsyncUdpSocket>,
+        socket: Box<dyn AsyncUdpSocket>,
         inner: proto::Endpoint,
         ipv6: bool,
         runtime: Arc<dyn Runtime>,
     ) -> Self {
         let (sender, events) = mpsc::unbounded_channel();
         let recv_state = RecvState::new(sender, socket.max_receive_segments(), &inner);
-        let sender = socket.clone().create_sender();
+        let sender = socket.create_sender();
         Self(Arc::new(EndpointInner {
             shared: Shared {
                 incoming: Notify::new(),
@@ -809,7 +807,7 @@ impl RecvState {
         &mut self,
         cx: &mut Context,
         endpoint: &mut proto::Endpoint,
-        socket: &dyn AsyncUdpSocket,
+        socket: &mut dyn AsyncUdpSocket,
         sender: &mut Pin<Box<dyn UdpSender>>,
         runtime: &dyn Runtime,
         now: Instant,

quinn/src/runtime.rs

@@ -21,7 +21,7 @@ pub trait Runtime: Send + Sync + Debug + 'static {
     fn spawn(&self, future: Pin<Box<dyn Future<Output = ()> + Send>>);
     /// Convert `t` into the socket type used by this runtime
     #[cfg(not(wasm_browser))]
-    fn wrap_udp_socket(&self, t: std::net::UdpSocket) -> io::Result<Arc<dyn AsyncUdpSocket>>;
+    fn wrap_udp_socket(&self, t: std::net::UdpSocket) -> io::Result<Box<dyn AsyncUdpSocket>>;
     /// Look up the current time
     ///
     /// Allows simulating the flow of time for testing.
@@ -50,11 +50,11 @@ pub trait AsyncUdpSocket: Send + Sync + Debug + 'static {
     /// [`Waker`].
     ///
     /// [`Waker`]: std::task::Waker
-    fn create_sender(self: Arc<Self>) -> Pin<Box<dyn UdpSender>>;
+    fn create_sender(&self) -> Pin<Box<dyn UdpSender>>;
 
     /// Receive UDP datagrams, or register to be woken if receiving may succeed in the future
     fn poll_recv(
-        &self,
+        &mut self,
         cx: &mut Context,
         bufs: &mut [IoSliceMut<'_>],
         meta: &mut [RecvMeta],

quinn/src/runtime/async_io.rs

@@ -39,8 +39,8 @@ mod smol {
         fn wrap_udp_socket(
             &self,
             sock: std::net::UdpSocket,
-        ) -> io::Result<Arc<dyn AsyncUdpSocket>> {
-            Ok(Arc::new(UdpSocket::new(sock)?))
+        ) -> io::Result<Box<dyn AsyncUdpSocket>> {
+            Ok(Box::new(UdpSocket::new(sock)?))
         }
     }
 }
@@ -55,25 +55,25 @@ impl AsyncTimer for Timer {
     }
 }
 
-#[cfg(feature = "runtime-smol")]
-#[derive(Debug)]
+#[cfg(any(feature = "runtime-smol"))]
+#[derive(Debug, Clone)]
 struct UdpSocket {
-    io: Async<std::net::UdpSocket>,
-    inner: udp::UdpSocketState,
+    io: Arc<Async<std::net::UdpSocket>>,
+    inner: Arc<udp::UdpSocketState>,
 }
 
 #[cfg(feature = "runtime-smol")]
 impl UdpSocket {
     fn new(sock: std::net::UdpSocket) -> io::Result<Self> {
         Ok(Self {
-            inner: udp::UdpSocketState::new((&sock).into())?,
-            io: Async::new_nonblocking(sock)?,
+            inner: Arc::new(udp::UdpSocketState::new((&sock).into())?),
+            io: Arc::new(Async::new_nonblocking(sock)?),
         })
     }
 }
 
 #[cfg(feature = "runtime-smol")]
-impl UdpSenderHelperSocket for Arc<UdpSocket> {
+impl UdpSenderHelperSocket for UdpSocket {
     fn max_transmit_segments(&self) -> usize {
         self.inner.max_gso_segments()
     }
@@ -85,18 +85,15 @@ impl UdpSenderHelperSocket for Arc<UdpSocket> {
 
 #[cfg(feature = "runtime-smol")]
 impl AsyncUdpSocket for UdpSocket {
-    fn create_sender(self: Arc<Self>) -> Pin<Box<dyn UdpSender>> {
-        Box::pin(UdpSenderHelper::new(
-            Arc::clone(&self),
-            |socket: &Arc<Self>| {
-                let socket = socket.clone();
-                async move { socket.io.writable().await }
-            },
-        ))
+    fn create_sender(&self) -> Pin<Box<dyn UdpSender>> {
+        Box::pin(UdpSenderHelper::new(self.clone(), |socket: &Self| {
+            let socket = socket.clone();
+            async move { socket.io.writable().await }
+        }))
     }
 
     fn poll_recv(
-        &self,
+        &mut self,
         cx: &mut Context,
         bufs: &mut [io::IoSliceMut<'_>],
         meta: &mut [udp::RecvMeta],
@@ -110,7 +107,7 @@ impl AsyncUdpSocket for UdpSocket {
     }
 
     fn local_addr(&self) -> io::Result<std::net::SocketAddr> {
-        self.io.as_ref().local_addr()
+        self.io.as_ref().as_ref().local_addr()
     }
 
     fn may_fragment(&self) -> bool {

quinn/src/runtime/tokio.rs

@@ -27,10 +27,10 @@ impl Runtime for TokioRuntime {
         tokio::spawn(future);
     }
 
-    fn wrap_udp_socket(&self, sock: std::net::UdpSocket) -> io::Result<Arc<dyn AsyncUdpSocket>> {
-        Ok(Arc::new(UdpSocket {
-            inner: udp::UdpSocketState::new((&sock).into())?,
-            io: tokio::net::UdpSocket::from_std(sock)?,
+    fn wrap_udp_socket(&self, sock: std::net::UdpSocket) -> io::Result<Box<dyn AsyncUdpSocket>> {
+        Ok(Box::new(UdpSocket {
+            inner: Arc::new(udp::UdpSocketState::new((&sock).into())?),
+            io: Arc::new(tokio::net::UdpSocket::from_std(sock)?),
         }))
     }
 
@@ -48,13 +48,13 @@ impl AsyncTimer for Sleep {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 struct UdpSocket {
-    io: tokio::net::UdpSocket,
-    inner: udp::UdpSocketState,
+    io: Arc<tokio::net::UdpSocket>,
+    inner: Arc<udp::UdpSocketState>,
 }
 
-impl UdpSenderHelperSocket for Arc<UdpSocket> {
+impl UdpSenderHelperSocket for UdpSocket {
     fn max_transmit_segments(&self) -> usize {
         self.inner.max_gso_segments()
     }
@@ -67,18 +67,15 @@ impl UdpSenderHelperSocket for Arc<UdpSocket> {
 }
 
 impl AsyncUdpSocket for UdpSocket {
-    fn create_sender(self: Arc<Self>) -> Pin<Box<dyn super::UdpSender>> {
-        Box::pin(UdpSenderHelper::new(
-            Arc::clone(&self),
-            |socket: &Arc<Self>| {
-                let socket = socket.clone();
-                async move { socket.io.writable().await }
-            },
-        ))
+    fn create_sender(&self) -> Pin<Box<dyn super::UdpSender>> {
+        Box::pin(UdpSenderHelper::new(self.clone(), |socket: &Self| {
+            let socket = socket.clone();
+            async move { socket.io.writable().await }
+        }))
     }
 
     fn poll_recv(
-        &self,
+        &mut self,
         cx: &mut Context,
         bufs: &mut [std::io::IoSliceMut<'_>],
         meta: &mut [udp::RecvMeta],