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

quinn/src/endpoint.rs

@@ -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
@@ -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,7 +465,7 @@ 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.
@@ -499,7 +497,7 @@ impl State {
             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,
@@ -719,7 +717,7 @@ impl EndpointRef {
     ) -> 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

@@ -49,7 +49,7 @@ 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(

quinn/src/runtime/async_io.rs

@@ -55,7 +55,7 @@ impl AsyncTimer for Timer {
     }
 }
 
-#[cfg(feature = "runtime-smol")]
+#[cfg(any(feature = "runtime-smol"))]
 #[derive(Debug, Clone)]
 struct UdpSocket {
     io: Arc<Async<std::net::UdpSocket>>,
@@ -73,7 +73,7 @@ impl UdpSocket {
 }
 
 #[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,14 +85,11 @@ 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(

quinn/src/runtime/tokio.rs

@@ -55,7 +55,7 @@ struct 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()
     }
@@ -68,14 +68,11 @@ 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(