diff --git a/quinn/src/endpoint.rs b/quinn/src/endpoint.rs index d1bdd23b3..761a2fc5f 100644 --- a/quinn/src/endpoint.rs +++ b/quinn/src/endpoint.rs @@ -129,7 +129,7 @@ impl Endpoint { pub fn new_with_abstract_socket( config: EndpointConfig, server_config: Option, - socket: Arc, + socket: Box, runtime: Arc, ) -> io::Result { 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) -> io::Result<()> { + pub fn rebind_abstract(&self, socket: Box) -> 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, + socket: Box, sender: Pin>, /// During an active migration, abandoned_socket receives traffic /// until the first packet arrives on the new socket. - prev_socket: Option>, + prev_socket: Option>, inner: proto::Endpoint, recv_state: RecvState, driver: Option, @@ -494,12 +492,12 @@ impl State { fn drive_recv(&mut self, cx: &mut Context, now: Instant) -> Result { 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); impl EndpointRef { pub(crate) fn new( - socket: Arc, + socket: Box, inner: proto::Endpoint, ipv6: bool, runtime: Arc, ) -> 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>, runtime: &dyn Runtime, now: Instant, diff --git a/quinn/src/runtime.rs b/quinn/src/runtime.rs index a78fbbd81..64734d2ea 100644 --- a/quinn/src/runtime.rs +++ b/quinn/src/runtime.rs @@ -21,7 +21,7 @@ pub trait Runtime: Send + Sync + Debug + 'static { fn spawn(&self, future: Pin + 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>; + fn wrap_udp_socket(&self, t: std::net::UdpSocket) -> io::Result>; /// 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) -> Pin>; + fn create_sender(&self) -> Pin>; /// 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], diff --git a/quinn/src/runtime/async_io.rs b/quinn/src/runtime/async_io.rs index 546647e3a..eb5380665 100644 --- a/quinn/src/runtime/async_io.rs +++ b/quinn/src/runtime/async_io.rs @@ -39,8 +39,8 @@ mod smol { fn wrap_udp_socket( &self, sock: std::net::UdpSocket, - ) -> io::Result> { - Ok(Arc::new(UdpSocket::new(sock)?)) + ) -> io::Result> { + 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, - inner: udp::UdpSocketState, + io: Arc>, + inner: Arc, } #[cfg(feature = "runtime-smol")] impl UdpSocket { fn new(sock: std::net::UdpSocket) -> io::Result { 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 { +impl UdpSenderHelperSocket for UdpSocket { fn max_transmit_segments(&self) -> usize { self.inner.max_gso_segments() } @@ -85,18 +85,15 @@ impl UdpSenderHelperSocket for Arc { #[cfg(feature = "runtime-smol")] impl AsyncUdpSocket for UdpSocket { - fn create_sender(self: Arc) -> Pin> { - Box::pin(UdpSenderHelper::new( - Arc::clone(&self), - |socket: &Arc| { - let socket = socket.clone(); - async move { socket.io.writable().await } - }, - )) + fn create_sender(&self) -> Pin> { + 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 { - self.io.as_ref().local_addr() + self.io.as_ref().as_ref().local_addr() } fn may_fragment(&self) -> bool { diff --git a/quinn/src/runtime/tokio.rs b/quinn/src/runtime/tokio.rs index 6bbfad4d6..6b54bcb09 100644 --- a/quinn/src/runtime/tokio.rs +++ b/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> { - 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> { + 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, + inner: Arc, } -impl UdpSenderHelperSocket for Arc { +impl UdpSenderHelperSocket for UdpSocket { fn max_transmit_segments(&self) -> usize { self.inner.max_gso_segments() } @@ -67,18 +67,15 @@ impl UdpSenderHelperSocket for Arc { } impl AsyncUdpSocket for UdpSocket { - fn create_sender(self: Arc) -> Pin> { - Box::pin(UdpSenderHelper::new( - Arc::clone(&self), - |socket: &Arc| { - let socket = socket.clone(); - async move { socket.io.writable().await } - }, - )) + fn create_sender(&self) -> Pin> { + 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],