From aa2d7736ded46ff09d14fd84eaedc485439efb10 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Philipp=20Kr=C3=BCger?= Date: Thu, 5 Jun 2025 14:35:42 +0200 Subject: [PATCH] Refactor `UdpPoller` into `UdpSender` and use it in favor of `AsyncUdpSocket` in connections --- quinn/src/connection.rs | 51 ++++----- quinn/src/endpoint.rs | 89 +++++++++++----- quinn/src/lib.rs | 6 +- quinn/src/runtime.rs | 189 ++++++++++++++++++++++------------ quinn/src/runtime/async_io.rs | 28 +++-- quinn/src/runtime/tokio.rs | 27 +++-- 6 files changed, 244 insertions(+), 146 deletions(-) diff --git a/quinn/src/connection.rs b/quinn/src/connection.rs index ca2014b6d..ab3777d71 100644 --- a/quinn/src/connection.rs +++ b/quinn/src/connection.rs @@ -20,7 +20,7 @@ use crate::{ ConnectionEvent, Duration, Instant, VarInt, mutex::Mutex, recv_stream::RecvStream, - runtime::{AsyncTimer, AsyncUdpSocket, Runtime, UdpPoller}, + runtime::{AsyncTimer, Runtime, UdpSender}, send_stream::SendStream, udp_transmit, }; @@ -43,7 +43,7 @@ impl Connecting { conn: proto::Connection, endpoint_events: mpsc::UnboundedSender<(ConnectionHandle, EndpointEvent)>, conn_events: mpsc::UnboundedReceiver, - socket: Arc, + sender: Pin>, runtime: Arc, ) -> Self { let (on_handshake_data_send, on_handshake_data_recv) = oneshot::channel(); @@ -55,7 +55,7 @@ impl Connecting { conn_events, on_handshake_data_send, on_connected_send, - socket, + sender, runtime.clone(), ); @@ -882,7 +882,7 @@ impl ConnectionRef { conn_events: mpsc::UnboundedReceiver, on_handshake_data: oneshot::Sender<()>, on_connected: oneshot::Sender, - socket: Arc, + sender: Pin>, runtime: Arc, ) -> Self { Self(Arc::new(ConnectionInner { @@ -902,8 +902,7 @@ impl ConnectionRef { stopped: FxHashMap::default(), error: None, ref_count: 0, - io_poller: socket.clone().create_io_poller(), - socket, + sender, runtime, send_buffer: Vec::new(), buffered_transmit: None, @@ -983,8 +982,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, - socket: Arc, - io_poller: Pin>, + sender: Pin>, runtime: Arc, send_buffer: Vec, /// We buffer a transmit when the underlying I/O would block @@ -997,7 +995,7 @@ impl State { let mut transmits = 0; let max_datagrams = self - .socket + .sender .max_transmit_segments() .min(MAX_TRANSMIT_SEGMENTS); @@ -1024,28 +1022,18 @@ impl State { } }; - if self.io_poller.as_mut().poll_writable(cx)?.is_pending() { - // Retry after a future wakeup - self.buffered_transmit = Some(t); - return Ok(false); - } - let len = t.size; - let retry = match self - .socket - .try_send(&udp_transmit(&t, &self.send_buffer[..len])) + match self + .sender + .as_mut() + .poll_send(&udp_transmit(&t, &self.send_buffer[..len]), cx) { - Ok(()) => false, - Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => true, - Err(e) => return Err(e), - }; - if retry { - // We thought the socket was writable, but it wasn't. Retry so that either another - // `poll_writable` call determines that the socket is indeed not writable and - // registers us for a wakeup, or the send succeeds if this really was just a - // transient failure. - self.buffered_transmit = Some(t); - continue; + Poll::Pending => { + self.buffered_transmit = Some(t); + return Ok(false); + } + Poll::Ready(Err(e)) => return Err(e), + Poll::Ready(Ok(())) => {} } if transmits >= MAX_TRANSMIT_DATAGRAMS { @@ -1075,9 +1063,8 @@ impl State { ) -> Result<(), ConnectionError> { loop { match self.conn_events.poll_recv(cx) { - Poll::Ready(Some(ConnectionEvent::Rebind(socket))) => { - self.socket = socket; - self.io_poller = self.socket.clone().create_io_poller(); + Poll::Ready(Some(ConnectionEvent::Rebind(sender))) => { + self.sender = sender; self.inner.local_address_changed(); } Poll::Ready(Some(ConnectionEvent::Proto(event))) => { diff --git a/quinn/src/endpoint.rs b/quinn/src/endpoint.rs index d6879d22f..d1bdd23b3 100644 --- a/quinn/src/endpoint.rs +++ b/quinn/src/endpoint.rs @@ -2,21 +2,20 @@ use std::{ collections::VecDeque, fmt, future::Future, - io, - io::IoSliceMut, + io::{self, IoSliceMut}, mem, net::{SocketAddr, SocketAddrV6}, pin::Pin, str, sync::{Arc, Mutex}, - task::{Context, Poll, Waker}, + task::{Context, Poll, RawWaker, RawWakerVTable, Waker}, }; #[cfg(all(not(wasm_browser), any(feature = "aws-lc-rs", feature = "ring")))] use crate::runtime::default_runtime; use crate::{ Instant, - runtime::{AsyncUdpSocket, Runtime}, + runtime::{AsyncUdpSocket, Runtime, UdpSender}, udp_transmit, }; use bytes::{Bytes, BytesMut}; @@ -225,12 +224,12 @@ impl Endpoint { .inner .connect(self.runtime.now(), config, addr, server_name)?; - let socket = endpoint.socket.clone(); + let sender = endpoint.socket.clone().create_sender(); endpoint.stats.outgoing_handshakes += 1; Ok(endpoint .recv_state .connections - .insert(ch, conn, socket, self.runtime.clone())) + .insert(ch, conn, sender, self.runtime.clone())) } /// Switch to a new UDP socket @@ -256,7 +255,9 @@ 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())); + let _ = sender.send(ConnectionEvent::Rebind( + inner.socket.clone().create_sender(), + )); } if let Some(driver) = inner.driver.take() { // Ensure the driver can register for wake-ups from the new socket @@ -425,16 +426,16 @@ impl EndpointInner { { Ok((handle, conn)) => { state.stats.accepted_handshakes += 1; - let socket = state.socket.clone(); + let sender = state.socket.clone().create_sender(); let runtime = state.runtime.clone(); Ok(state .recv_state .connections - .insert(handle, conn, socket, runtime)) + .insert(handle, conn, sender, runtime)) } Err(error) => { if let Some(transmit) = error.response { - respond(transmit, &response_buffer, &*state.socket); + respond(transmit, &response_buffer, &mut state.sender); } Err(error.cause) } @@ -446,14 +447,14 @@ impl EndpointInner { state.stats.refused_handshakes += 1; let mut response_buffer = Vec::new(); let transmit = state.inner.refuse(incoming, &mut response_buffer); - respond(transmit, &response_buffer, &*state.socket); + respond(transmit, &response_buffer, &mut state.sender); } pub(crate) fn retry(&self, incoming: proto::Incoming) -> Result<(), proto::RetryError> { let mut state = self.state.lock().unwrap(); let mut response_buffer = Vec::new(); let transmit = state.inner.retry(incoming, &mut response_buffer)?; - respond(transmit, &response_buffer, &*state.socket); + respond(transmit, &response_buffer, &mut state.sender); Ok(()) } @@ -467,6 +468,7 @@ impl EndpointInner { #[derive(Debug)] pub(crate) struct State { socket: Arc, + sender: Pin>, /// During an active migration, abandoned_socket receives traffic /// until the first packet arrives on the new socket. prev_socket: Option>, @@ -494,16 +496,26 @@ impl State { self.recv_state.recv_limiter.start_cycle(get_time); if let Some(socket) = &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, &*self.runtime, now); + let poll_res = self.recv_state.poll_socket( + cx, + &mut self.inner, + &**socket, + &mut self.sender, + &*self.runtime, + now, + ); if poll_res.is_err() { self.prev_socket = None; } }; - let poll_res = - self.recv_state - .poll_socket(cx, &mut self.inner, &*self.socket, &*self.runtime, now); + let poll_res = self.recv_state.poll_socket( + cx, + &mut self.inner, + &*self.socket, + &mut self.sender, + &*self.runtime, + now, + ); self.recv_state.recv_limiter.finish_cycle(get_time); let poll_res = poll_res?; if poll_res.received_connection_packet { @@ -555,7 +567,11 @@ impl Drop for State { } } -fn respond(transmit: proto::Transmit, response_buffer: &[u8], socket: &dyn AsyncUdpSocket) { +fn respond( + transmit: proto::Transmit, + response_buffer: &[u8], + sender: &mut Pin>, +) { // Send if there's kernel buffer space; otherwise, drop it // // As an endpoint-generated packet, we know this is an @@ -576,7 +592,29 @@ fn respond(transmit: proto::Transmit, response_buffer: &[u8], socket: &dyn Async // to transmit. This is morally equivalent to the packet getting // lost due to congestion further along the link, which // similarly relies on peer retries for recovery. - _ = socket.try_send(&udp_transmit(&transmit, &response_buffer[..transmit.size])); + + // Copied from rust 1.85's std::task::Waker::noop() implementation for backwards compatibility + const NOOP: RawWaker = { + const VTABLE: RawWakerVTable = RawWakerVTable::new( + // Cloning just returns a new no-op raw waker + |_| NOOP, + // `wake` does nothing + |_| {}, + // `wake_by_ref` does nothing + |_| {}, + // Dropping does nothing as we don't allocate anything + |_| {}, + ); + RawWaker::new(std::ptr::null(), &VTABLE) + }; + // SAFETY: Copied from rust stdlib, the NOOP waker is thread-safe and doesn't violate the RawWakerVTable contract, + // it doesn't access the data pointer at all. + let waker = unsafe { Waker::from_raw(NOOP) }; + let mut cx = Context::from_waker(&waker); + _ = sender.as_mut().poll_send( + &udp_transmit(&transmit, &response_buffer[..transmit.size]), + &mut cx, + ); } #[inline] @@ -603,7 +641,7 @@ impl ConnectionSet { &mut self, handle: ConnectionHandle, conn: proto::Connection, - socket: Arc, + sender: Pin>, runtime: Arc, ) -> Connecting { let (send, recv) = mpsc::unbounded_channel(); @@ -615,7 +653,7 @@ impl ConnectionSet { .unwrap(); } self.senders.insert(handle, send); - Connecting::new(handle, conn, self.sender.clone(), recv, socket, runtime) + Connecting::new(handle, conn, self.sender.clone(), recv, sender, runtime) } fn is_empty(&self) -> bool { @@ -681,6 +719,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(); Self(Arc::new(EndpointInner { shared: Shared { incoming: Notify::new(), @@ -688,6 +727,7 @@ impl EndpointRef { }, state: Mutex::new(State { socket, + sender, prev_socket: None, inner, ipv6, @@ -770,6 +810,7 @@ impl RecvState { cx: &mut Context, endpoint: &mut proto::Endpoint, socket: &dyn AsyncUdpSocket, + sender: &mut Pin>, runtime: &dyn Runtime, now: Instant, ) -> Result { @@ -809,7 +850,7 @@ impl RecvState { } else { let transmit = endpoint.refuse(incoming, &mut response_buffer); - respond(transmit, &response_buffer, socket); + respond(transmit, &response_buffer, sender); } } Some(DatagramEvent::ConnectionEvent(handle, event)) => { @@ -823,7 +864,7 @@ impl RecvState { .send(ConnectionEvent::Proto(event)); } Some(DatagramEvent::Response(transmit)) => { - respond(transmit, &response_buffer, socket); + respond(transmit, &response_buffer, sender); } None => {} } diff --git a/quinn/src/lib.rs b/quinn/src/lib.rs index 0aff67922..64f4800bb 100644 --- a/quinn/src/lib.rs +++ b/quinn/src/lib.rs @@ -41,7 +41,7 @@ #![warn(unreachable_pub)] #![warn(clippy::use_self)] -use std::sync::Arc; +use std::pin::Pin; mod connection; mod endpoint; @@ -85,7 +85,7 @@ pub use crate::recv_stream::{ReadError, ReadExactError, ReadToEndError, RecvStre pub use crate::runtime::SmolRuntime; #[cfg(feature = "runtime-tokio")] pub use crate::runtime::TokioRuntime; -pub use crate::runtime::{AsyncTimer, AsyncUdpSocket, Runtime, UdpPoller, default_runtime}; +pub use crate::runtime::{AsyncTimer, AsyncUdpSocket, Runtime, UdpSender, default_runtime}; pub use crate::send_stream::{SendStream, StoppedError, WriteError}; #[cfg(test)] @@ -98,7 +98,7 @@ enum ConnectionEvent { reason: bytes::Bytes, }, Proto(proto::ConnectionEvent), - Rebind(Arc), + Rebind(Pin>), } fn udp_transmit<'a>(t: &proto::Transmit, buffer: &'a [u8]) -> udp::Transmit<'a> { diff --git a/quinn/src/runtime.rs b/quinn/src/runtime.rs index e2ba7c60a..a78fbbd81 100644 --- a/quinn/src/runtime.rs +++ b/quinn/src/runtime.rs @@ -1,5 +1,5 @@ use std::{ - fmt::Debug, + fmt::{self, Debug}, future::Future, io::{self, IoSliceMut}, net::SocketAddr, @@ -40,23 +40,17 @@ pub trait AsyncTimer: Send + Debug + 'static { /// Abstract implementation of a UDP socket for runtime independence pub trait AsyncUdpSocket: Send + Sync + Debug + 'static { - /// Create a [`UdpPoller`] that can register a single task for write-readiness notifications + /// Create a [`UdpSender`] that can register a single task for write-readiness notifications + /// and send a transmit, if ready. /// /// A `poll_send` method on a single object can usually store only one [`Waker`] at a time, /// i.e. allow at most one caller to wait for an event. This method allows any number of - /// interested tasks to construct their own [`UdpPoller`] object. They can all then wait for the - /// same event and be notified concurrently, because each [`UdpPoller`] can store a separate + /// interested tasks to construct their own [`UdpSender`] object. They can all then wait for the + /// same event and be notified concurrently, because each [`UdpSender`] can store a separate /// [`Waker`]. /// /// [`Waker`]: std::task::Waker - fn create_io_poller(self: Arc) -> Pin>; - - /// Send UDP datagrams from `transmits`, or return `WouldBlock` and clear the underlying - /// socket's readiness, or return an I/O error - /// - /// If this returns [`io::ErrorKind::WouldBlock`], [`UdpPoller::poll_writable`] must be called - /// to register the calling task to be woken when a send should be attempted again. - fn try_send(&self, transmit: &Transmit) -> io::Result<()>; + fn create_sender(self: Arc) -> Pin>; /// Receive UDP datagrams, or register to be woken if receiving may succeed in the future fn poll_recv( @@ -69,11 +63,6 @@ pub trait AsyncUdpSocket: Send + Sync + 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 @@ -88,71 +77,141 @@ pub trait AsyncUdpSocket: Send + Sync + Debug + 'static { } } -/// An object polled to detect when an associated [`AsyncUdpSocket`] is writable +/// An object for asynchronously writing to an associated [`AsyncUdpSocket`]. /// -/// Any number of `UdpPoller`s may exist for a single [`AsyncUdpSocket`]. Each `UdpPoller` is -/// responsible for notifying at most one task when that socket becomes writable. -pub trait UdpPoller: Send + Sync + Debug + 'static { - /// Check whether the associated socket is likely to be writable +/// Any number of [`UdpSender`]s may exist for a single [`AsyncUdpSocket`]. Each [`UdpSender`] is +/// responsible for notifying at most one task for send readiness. +pub trait UdpSender: Send + Sync + Debug + 'static { + /// Send a UDP datagram, or register to be woken if sending may succeed in the future. /// - /// Must be called after [`AsyncUdpSocket::try_send`] returns [`io::ErrorKind::WouldBlock`] to - /// register the task associated with `cx` to be woken when a send should be attempted - /// again. Unlike in [`Future::poll`], a [`UdpPoller`] may be reused indefinitely no matter how - /// many times `poll_writable` returns [`Poll::Ready`]. - fn poll_writable(self: Pin<&mut Self>, cx: &mut Context) -> Poll>; + /// Usually implementations of this will poll the socket for writability before trying to + /// write to them, and retry both if writing fails. + /// + /// Quinn will create multiple [`UdpSender`]s, one for each task it's using it from. Thus it's + /// important to poll the underlying socket in a way that doesn't overwrite wakers. + /// + /// A single [`UdpSender`] will be re-used, even if `poll_send` returns `Poll::Ready` once, + /// unlike [`Future::poll`], so calling it again after readiness should not panic. + fn poll_send( + self: Pin<&mut Self>, + transmit: &Transmit, + cx: &mut Context, + ) -> Poll>; + + /// Maximum number of datagrams that a [`Transmit`] may encode. + fn max_transmit_segments(&self) -> usize { + 1 + } } pin_project_lite::pin_project! { - /// Helper adapting a function `MakeFut` that constructs a single-use future `Fut` into a - /// [`UdpPoller`] that may be reused indefinitely - struct UdpPollHelper { - make_fut: MakeFut, + /// A helper for constructing [`UdpSender`]s from an underlying `Socket` type. + /// + /// This struct implements [`UdpSender`] if `MakeWritableFn` produces a `WritableFut`. + /// + /// Also serves as a trick, since `WritableFut` doesn't need to be a named future, + /// it can be an anonymous async block, as long as `MakeWritableFn` produces that + /// anonymous async block type. + /// + /// The `UdpSenderHelper` generic type parameters don't need to named, as it will be + /// used in its dyn-compatible form as a `Pin>`. + struct UdpSenderHelper { + socket: Socket, + make_writable_fut_fn: MakeWritableFutFn, #[pin] - fut: Option, + writable_fut: Option, } } -impl UdpPollHelper { - /// Construct a [`UdpPoller`] that calls `make_fut` to get the future to poll, storing it until - /// it yields [`Poll::Ready`], then creating a new one on the next - /// [`poll_writable`](UdpPoller::poll_writable) - #[cfg(any(feature = "runtime-smol", feature = "runtime-tokio",))] - fn new(make_fut: MakeFut) -> Self { - Self { - make_fut, - fut: None, - } - } -} - -impl UdpPoller for UdpPollHelper -where - MakeFut: Fn() -> Fut + Send + Sync + 'static, - Fut: Future> + Send + Sync + 'static, +impl Debug + for UdpSenderHelper { - fn poll_writable(self: Pin<&mut Self>, cx: &mut Context) -> Poll> { - let mut this = self.project(); - if this.fut.is_none() { - this.fut.set(Some((this.make_fut)())); + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("UdpSender") + } +} + +impl + UdpSenderHelper +{ + /// Create helper that implements [`UdpSender`] from a socket. + /// + /// Additionally you need to provide what is essentially an async function + /// that resolves once the socket is write-ready. + /// + /// See also the bounds on this struct's [`UdpSender`] implementation. + #[cfg(any(feature = "runtime-smol", feature = "runtime-tokio",))] + fn new(inner: Socket, make_fut: MakeWritableFutFn) -> Self { + Self { + socket: inner, + make_writable_fut_fn: make_fut, + writable_fut: None, } - // We're forced to `unwrap` here because `Fut` may be `!Unpin`, which means we can't safely - // obtain an `&mut Fut` after storing it in `self.fut` when `self` is already behind `Pin`, - // and if we didn't store it then we wouldn't be able to keep it alive between - // `poll_writable` calls. - let result = this.fut.as_mut().as_pin_mut().unwrap().poll(cx); - if result.is_ready() { + } +} + +impl super::UdpSender + for UdpSenderHelper +where + Socket: UdpSenderHelperSocket, + MakeWritableFutFn: Fn(&Socket) -> WritableFut + Send + Sync + 'static, + WritableFut: Future> + Send + Sync + 'static, +{ + fn poll_send( + self: Pin<&mut Self>, + transmit: &udp::Transmit, + cx: &mut Context, + ) -> Poll> { + let mut this = self.project(); + loop { + if this.writable_fut.is_none() { + this.writable_fut + .set(Some((this.make_writable_fut_fn)(this.socket))); + } + // We're forced to `unwrap` here because `Fut` may be `!Unpin`, which means we can't safely + // obtain an `&mut WritableFut` after storing it in `self.writable_fut` when `self` is already behind `Pin`, + // and if we didn't store it then we wouldn't be able to keep it alive between + // `poll_send` calls. + let result = + std::task::ready!(this.writable_fut.as_mut().as_pin_mut().unwrap().poll(cx)); + // Polling an arbitrary `Future` after it becomes ready is a logic error, so arrange for // a new `Future` to be created on the next call. - this.fut.set(None); + this.writable_fut.set(None); + + // If .writable() fails, propagate the error + result?; + + match this.socket.try_send(transmit) { + // We thought the socket was writable, but it wasn't, then retry so that either another + // `writable().await` call determines that the socket is indeed not writable and + // registers us for a wakeup, or the send succeeds if this really was just a + // transient failure. + Err(e) if e.kind() == io::ErrorKind::WouldBlock => continue, + // In all other cases, either propagate the error or we're Ok + result => return Poll::Ready(result), + } } - result + } + + fn max_transmit_segments(&self) -> usize { + self.socket.max_transmit_segments() } } -impl Debug for UdpPollHelper { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("UdpPollHelper").finish_non_exhaustive() - } +/// Parts of the [`UdpSender`] trait that aren't asynchronous or require storing wakers. +/// +/// This trait is used by [`UdpSenderHelper`] to help construct [`UdpSender`]s. +trait UdpSenderHelperSocket: Send + Sync + 'static { + /// Try to send a transmit, if the socket happens to be write-ready. + /// + /// If not write-ready, this is allowed to return [`std::io::ErrorKind::WouldBlock`]. + /// + /// The [`UdpSenderHelper`] will use this to implement [`UdpSender::poll_send`]. + fn try_send(&self, transmit: &udp::Transmit) -> io::Result<()>; + + /// See [`UdpSender::max_transmit_segments`]. + fn max_transmit_segments(&self) -> usize; } /// Automatically select an appropriate runtime from those enabled at compile time diff --git a/quinn/src/runtime/async_io.rs b/quinn/src/runtime/async_io.rs index 5ebcd3230..546647e3a 100644 --- a/quinn/src/runtime/async_io.rs +++ b/quinn/src/runtime/async_io.rs @@ -13,7 +13,7 @@ use async_io::Timer; use super::AsyncTimer; #[cfg(feature = "runtime-smol")] -use super::{AsyncUdpSocket, Runtime, UdpPollHelper}; +use super::{AsyncUdpSocket, Runtime, UdpSender, UdpSenderHelper, UdpSenderHelperSocket}; #[cfg(feature = "runtime-smol")] // Due to MSRV, we must specify `self::` where there's crate/module ambiguity @@ -73,17 +73,27 @@ impl UdpSocket { } #[cfg(feature = "runtime-smol")] -impl AsyncUdpSocket for UdpSocket { - fn create_io_poller(self: Arc) -> Pin> { - Box::pin(UdpPollHelper::new(move || { - let socket = self.clone(); - async move { socket.io.writable().await } - })) +impl UdpSenderHelperSocket for Arc { + fn max_transmit_segments(&self) -> usize { + self.inner.max_gso_segments() } fn try_send(&self, transmit: &udp::Transmit) -> io::Result<()> { self.inner.send((&self.io).into(), transmit) } +} + +#[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 poll_recv( &self, @@ -107,10 +117,6 @@ impl AsyncUdpSocket for UdpSocket { self.inner.may_fragment() } - fn max_transmit_segments(&self) -> usize { - self.inner.max_gso_segments() - } - fn max_receive_segments(&self) -> usize { self.inner.gro_segments() } diff --git a/quinn/src/runtime/tokio.rs b/quinn/src/runtime/tokio.rs index 0e423660d..6bbfad4d6 100644 --- a/quinn/src/runtime/tokio.rs +++ b/quinn/src/runtime/tokio.rs @@ -12,7 +12,7 @@ use tokio::{ time::{Sleep, sleep_until}, }; -use super::{AsyncTimer, AsyncUdpSocket, Runtime, UdpPollHelper}; +use super::{AsyncTimer, AsyncUdpSocket, Runtime, UdpSenderHelper, UdpSenderHelperSocket}; /// A Quinn runtime for Tokio #[derive(Debug)] @@ -54,12 +54,9 @@ struct UdpSocket { inner: udp::UdpSocketState, } -impl AsyncUdpSocket for UdpSocket { - fn create_io_poller(self: Arc) -> Pin> { - Box::pin(UdpPollHelper::new(move || { - let socket = self.clone(); - async move { socket.io.writable().await } - })) +impl UdpSenderHelperSocket for Arc { + fn max_transmit_segments(&self) -> usize { + self.inner.max_gso_segments() } fn try_send(&self, transmit: &udp::Transmit) -> io::Result<()> { @@ -67,6 +64,18 @@ impl AsyncUdpSocket for UdpSocket { self.inner.send((&self.io).into(), transmit) }) } +} + +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 poll_recv( &self, @@ -92,10 +101,6 @@ impl AsyncUdpSocket for UdpSocket { self.inner.may_fragment() } - fn max_transmit_segments(&self) -> usize { - self.inner.max_gso_segments() - } - fn max_receive_segments(&self) -> usize { self.inner.gro_segments() }