From a24898ada6c3e1d7908387270e3c72c8e78be744 Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Tue, 4 Aug 2026 19:17:33 +0700 Subject: [PATCH] refactor(client): share pollable async transactions --- crates/irtt-client/src/async_client.rs | 755 +++++++++++++++---------- 1 file changed, 469 insertions(+), 286 deletions(-) diff --git a/crates/irtt-client/src/async_client.rs b/crates/irtt-client/src/async_client.rs index 59fc34c..d3616d8 100644 --- a/crates/irtt-client/src/async_client.rs +++ b/crates/irtt-client/src/async_client.rs @@ -1,6 +1,9 @@ use std::{ + future::{poll_fn, Future}, io, net::SocketAddr, + pin::Pin, + task::{Context, Poll}, time::{Duration, Instant}, }; @@ -8,7 +11,6 @@ use std::{ use std::{ cell::{Cell, RefCell}, collections::VecDeque, - future, }; use tokio::time; @@ -49,6 +51,86 @@ struct PreparedAsyncOpenFailure { machine: PreparedOpenAcceptance, } +#[derive(Debug)] +enum AsyncOpenCleanup { + Acceptance { + primary: ClientError, + packet: Option>, + }, + Adapter(Box), +} + +impl AsyncOpenCleanup { + fn packet(&self) -> Option<&[u8]> { + match self { + Self::Acceptance { packet, .. } => packet.as_deref(), + Self::Adapter(failure) => failure.machine.cleanup_close_packet(), + } + } + + fn into_primary(self) -> ClientError { + match self { + Self::Acceptance { primary, .. } => primary, + Self::Adapter(failure) => failure.primary, + } + } +} + +/// Owned state for one asynchronous open transaction. +/// +/// This state deliberately contains no reference to [`AsyncClient`], allowing +/// a future managed target to store the client and opening state side by side. +#[derive(Debug)] +pub(crate) struct AsyncOpenState { + attempt: usize, + deadline: Option, + deadline_timer: Option>>, + request_submitted: bool, + buffer: [u8; MAX_OPEN_PACKET_SIZE], + cleanup: Option, +} + +impl AsyncOpenState { + pub(crate) fn new() -> Self { + Self { + attempt: 0, + deadline: None, + deadline_timer: None, + request_submitted: false, + buffer: [0_u8; MAX_OPEN_PACKET_SIZE], + cleanup: None, + } + } + + fn start_attempt(&mut self, timeout: Duration) -> Result<(), ClientError> { + let deadline = Instant::now() + .checked_add(timeout) + .ok_or(ClientError::DurationOverflow)?; + self.deadline = Some(deadline); + self.deadline_timer = Some(Box::pin(time::sleep_until(deadline.into()))); + self.request_submitted = false; + Ok(()) + } + + fn finish_attempt(&mut self) { + self.attempt += 1; + self.deadline = None; + self.deadline_timer = None; + self.request_submitted = false; + } + + fn deadline(&self) -> Instant { + self.deadline + .expect("an open attempt has an active deadline") + } + + fn deadline_timer(&mut self) -> &mut Pin> { + self.deadline_timer + .as_mut() + .expect("an open attempt has an active deadline timer") + } +} + #[cfg(test)] #[derive(Debug)] enum InjectedSend { @@ -155,69 +237,8 @@ impl AsyncClient { /// is dropped during best-effort post-token cleanup, cleanup may remain /// incomplete because this low-level client never detaches cleanup work. pub async fn open(&mut self) -> Result { - if self.prepared_open.is_none() { - let _ = self.machine.prepare_open_request()?; - unreachable!("connected async clients retain an open request"); - } - - let attempt_count = self.machine.config().open_timeouts.len(); - let mut open_buffer = [0_u8; MAX_OPEN_PACKET_SIZE]; - - for attempt in 0..attempt_count { - let timeout = self.machine.config().open_timeouts[attempt]; - let deadline = Instant::now() - .checked_add(timeout) - .ok_or(ClientError::DurationOverflow)?; - if !self.send_open_request(deadline).await? { - continue; - } - - loop { - let Some(datagram) = self.recv_open_datagram(&mut open_buffer, deadline).await? - else { - break; - }; - if datagram.received_at.mono > deadline { - break; - } - - let reply = match self - .machine - .inspect_open_datagram(&open_buffer[..datagram.len])? - { - OpenDatagramDisposition::Ignore => continue, - OpenDatagramDisposition::Trusted(reply) => reply, - }; - let machine = match self - .machine - .prepare_open_acceptance(reply, datagram.received_at) - { - Ok(machine) => machine, - Err(failure) => { - self.send_cleanup_close_best_effort( - failure.cleanup_close.as_deref(), - deadline, - ) - .await; - return Err(failure.primary); - } - }; - let prepared = match self.prepare_async_open(machine, datagram.received_at) { - Ok(prepared) => prepared, - Err(failure) => { - self.send_cleanup_close_best_effort( - failure.machine.cleanup_close_packet(), - deadline, - ) - .await; - return Err(failure.primary); - } - }; - return Ok(self.commit_async_open(prepared)); - } - } - - Err(ClientError::OpenTimeout) + let mut state = AsyncOpenState::new(); + poll_fn(|cx| self.poll_open(&mut state, cx)).await } /// Send one caller-paced probe with transactional protocol and schedule @@ -227,102 +248,7 @@ impl AsyncClient { /// cancellation, and socket errors so a later call retries the same logical /// probe without advancing state before kernel acceptance. pub async fn send_probe(&mut self) -> Result, ClientError> { - self.machine.ensure_open()?; - if self.prepared_probe.is_none() { - self.prepared_probe = self.machine.prepare_probe()?; - } - if self.prepared_probe.is_none() { - return Ok(Vec::new()); - } - - let mut events = Vec::new(); - loop { - #[cfg(test)] - if self.test_hooks.pause_probe_before_writable.replace(false) { - future::pending::<()>().await; - } - self.socket.writable().await?; - - #[cfg(test)] - let test_timestamps = self.test_hooks.take_probe_timestamps(); - #[cfg(not(test))] - let permission_at = Instant::now(); - #[cfg(test)] - let permission_at = test_timestamps - .map(|timestamps| timestamps.permission_at) - .unwrap_or_else(Instant::now); - - let schedule = self - .schedule - .as_mut() - .expect("open sessions always have a probe schedule"); - if !schedule.permit_probe_at(permission_at) { - self.prepared_probe = None; - return Ok(events); - } - let prepared = self - .prepared_probe - .as_ref() - .expect("prepared probe was retained across readiness"); - let machine_preflight = self.machine.preflight_probe_commit(prepared)?; - let schedule_commit = - schedule.preflight_caller_commit(machine_preflight.next_packets_sent)?; - events - .try_reserve(1) - .map_err(|source| ClientError::AllocationFailed { - operation: "probe event result", - source, - })?; - - let expected_bytes = prepared.bytes.len(); - let scheduled_at = schedule_commit.scheduled_at; - #[cfg(not(test))] - let sent_at = ClientTimestamp::now(); - #[cfg(test)] - let sent_at = test_timestamps - .map(|timestamps| timestamps.sent_at) - .unwrap_or_else(ClientTimestamp::now); - let machine_commit = self - .machine - .finalize_probe_commit(machine_preflight, sent_at)?; - #[cfg(not(test))] - let send_call_start = Instant::now(); - #[cfg(test)] - let send_call_start = test_timestamps - .map(|timestamps| timestamps.send_call_start) - .unwrap_or_else(Instant::now); - #[cfg(test)] - let send_result = self.test_hooks.try_send(&self.socket, &prepared.bytes); - #[cfg(not(test))] - let send_result = self.socket.try_send(&prepared.bytes); - let bytes = match send_result { - Ok(bytes) => bytes, - Err(error) if error.kind() == io::ErrorKind::WouldBlock => continue, - Err(error) => return Err(ClientError::Socket(error)), - }; - - #[cfg(not(test))] - let send_finished_at = Instant::now(); - #[cfg(test)] - let send_finished_at = test_timestamps - .map(|timestamps| timestamps.send_finished_at) - .unwrap_or_else(Instant::now); - let sent = self.machine.commit_probe_sent(machine_commit, bytes); - schedule.commit(schedule_commit); - self.prepared_probe = None; - - let send_call = send_finished_at.saturating_duration_since(send_call_start); - let timer_error = instant_abs_diff(sent_at.mono, scheduled_at); - validate_datagram_length(expected_bytes, bytes)?; - events.push(echo_sent_event( - self.remote, - sent, - send_call, - scheduled_at, - timer_error, - )); - return Ok(events); - } + poll_fn(|cx| self.poll_send_probe(cx)).await } /// Await and classify one complete UDP datagram. @@ -331,37 +257,7 @@ impl AsyncClient { /// retry without changing protocol state. Authenticated peer-close events /// remain authoritative even if best-effort DSCP cleanup fails. pub async fn recv(&mut self) -> Result, ClientError> { - self.machine.ensure_open()?; - - let datagram = loop { - self.socket.readable().await?; - #[cfg(test)] - if self.test_hooks.receive_would_block.get() > 0 { - self.test_hooks - .receive_would_block - .set(self.test_hooks.receive_would_block.get() - 1); - continue; - } - match try_recv_tokio_datagram(&self.socket, &mut self.recv_buffer) { - Ok(datagram) => break datagram, - Err(error) if error.kind() == io::ErrorKind::WouldBlock => continue, - Err(error) => return Err(ClientError::Socket(error)), - } - }; - - let events = self.machine.process_received_echo_packet( - &self.recv_buffer[..datagram.len], - datagram.received_at, - datagram.meta, - )?; - if self.machine.is_peer_closed() { - self.schedule = None; - self.prepared_probe = None; - if self.clear_peer_close_dscp().is_ok() { - self.applied_dscp = None; - } - } - Ok(events) + poll_fn(|cx| self.poll_recv(cx)).await } /// Poll protocol timeouts using a newly captured monotonic timestamp. @@ -381,24 +277,334 @@ impl AsyncClient { /// next await, so dropping this future at a suspension point cannot leave /// an otherwise-open session with cleared DSCP. pub async fn close(&mut self) -> Result, ClientError> { - let prepared = self.machine.prepare_close()?; - let mut events = Vec::new(); - events - .try_reserve(1) - .map_err(|source| ClientError::AllocationFailed { - operation: "close event result", - source, - })?; - let previous_dscp = self.applied_dscp; - let expected_bytes = prepared.bytes.len(); + poll_fn(|cx| self.poll_close(cx)).await + } + + pub(crate) fn poll_open( + &mut self, + state: &mut AsyncOpenState, + cx: &mut Context<'_>, + ) -> Poll> { + if self.prepared_open.is_none() { + if let Err(error) = self.machine.prepare_open_request() { + return Poll::Ready(Err(error)); + } + unreachable!("connected async clients retain an open request"); + } + + loop { + if state.cleanup.is_some() { + return self.poll_open_cleanup(state, cx); + } + + let attempt_count = self.machine.config().open_timeouts.len(); + if state.attempt >= attempt_count { + return Poll::Ready(Err(ClientError::OpenTimeout)); + } + if state.deadline.is_none() { + let timeout = self.machine.config().open_timeouts[state.attempt]; + if let Err(error) = state.start_attempt(timeout) { + return Poll::Ready(Err(error)); + } + } + + if !state.request_submitted { + #[cfg(test)] + if self.test_hooks.pause_open_before_writable.replace(false) { + return Poll::Pending; + } + match poll_writable_until(&self.socket, state.deadline_timer(), cx) { + Poll::Pending => return Poll::Pending, + Poll::Ready(Err(error)) => return Poll::Ready(Err(error)), + Poll::Ready(Ok(false)) => { + state.finish_attempt(); + continue; + } + Poll::Ready(Ok(true)) => {} + } + + let request = self + .prepared_open + .as_ref() + .expect("connected clients retain their prepared open request"); + #[cfg(test)] + let send_result = self.test_hooks.try_send(&self.socket, &request.bytes); + #[cfg(not(test))] + let send_result = self.socket.try_send(&request.bytes); + match send_result { + Ok(bytes) => { + if let Err(error) = validate_datagram_length(request.bytes.len(), bytes) { + return Poll::Ready(Err(error)); + } + state.request_submitted = true; + } + Err(error) if error.kind() == io::ErrorKind::WouldBlock => continue, + Err(error) => return Poll::Ready(Err(ClientError::Socket(error))), + } + } + + #[cfg(test)] + if self.test_hooks.pause_open_before_readable.replace(false) { + return Poll::Pending; + } + match poll_readable_until(&self.socket, state.deadline_timer(), cx) { + Poll::Pending => return Poll::Pending, + Poll::Ready(Err(error)) => return Poll::Ready(Err(error)), + Poll::Ready(Ok(false)) => { + state.finish_attempt(); + continue; + } + Poll::Ready(Ok(true)) => {} + } + + let datagram = match try_recv_tokio_datagram(&self.socket, &mut state.buffer) { + Ok(datagram) => datagram, + Err(error) if error.kind() == io::ErrorKind::WouldBlock => continue, + Err(error) => return Poll::Ready(Err(ClientError::Socket(error))), + }; + if datagram.received_at.mono > state.deadline() { + state.finish_attempt(); + continue; + } + + let reply = match self + .machine + .inspect_open_datagram(&state.buffer[..datagram.len]) + { + Ok(OpenDatagramDisposition::Ignore) => continue, + Ok(OpenDatagramDisposition::Trusted(reply)) => reply, + Err(error) => return Poll::Ready(Err(error)), + }; + let machine = match self + .machine + .prepare_open_acceptance(reply, datagram.received_at) + { + Ok(machine) => machine, + Err(failure) => { + state.cleanup = Some(AsyncOpenCleanup::Acceptance { + primary: failure.primary, + packet: failure.cleanup_close, + }); + continue; + } + }; + let prepared = match self.prepare_async_open(machine, datagram.received_at) { + Ok(prepared) => prepared, + Err(failure) => { + state.cleanup = Some(AsyncOpenCleanup::Adapter(failure)); + continue; + } + }; + return Poll::Ready(Ok(self.commit_async_open(prepared))); + } + } + + pub(crate) fn poll_send_probe( + &mut self, + cx: &mut Context<'_>, + ) -> Poll, ClientError>> { + if let Err(error) = self.machine.ensure_open() { + return Poll::Ready(Err(error)); + } + if self.prepared_probe.is_none() { + match self.machine.prepare_probe() { + Ok(prepared) => self.prepared_probe = prepared, + Err(error) => return Poll::Ready(Err(error)), + } + } + if self.prepared_probe.is_none() { + return Poll::Ready(Ok(Vec::new())); + } loop { #[cfg(test)] + if self.test_hooks.pause_probe_before_writable.replace(false) { + return Poll::Pending; + } + match self.socket.poll_send_ready(cx) { + Poll::Pending => return Poll::Pending, + Poll::Ready(Ok(())) => {} + Poll::Ready(Err(error)) => { + return Poll::Ready(Err(ClientError::Socket(error))); + } + } + + #[cfg(test)] + let test_timestamps = self.test_hooks.take_probe_timestamps(); + #[cfg(not(test))] + let permission_at = Instant::now(); + #[cfg(test)] + let permission_at = test_timestamps + .map(|timestamps| timestamps.permission_at) + .unwrap_or_else(Instant::now); + + let schedule = self + .schedule + .as_mut() + .expect("open sessions always have a probe schedule"); + if !schedule.permit_probe_at(permission_at) { + self.prepared_probe = None; + return Poll::Ready(Ok(Vec::new())); + } + let prepared = self + .prepared_probe + .as_ref() + .expect("prepared probe was retained across readiness"); + let machine_preflight = match self.machine.preflight_probe_commit(prepared) { + Ok(preflight) => preflight, + Err(error) => return Poll::Ready(Err(error)), + }; + let schedule_commit = + match schedule.preflight_caller_commit(machine_preflight.next_packets_sent) { + Ok(commit) => commit, + Err(error) => return Poll::Ready(Err(error)), + }; + let mut events = Vec::new(); + if let Err(source) = events.try_reserve(1) { + return Poll::Ready(Err(ClientError::AllocationFailed { + operation: "probe event result", + source, + })); + } + + let expected_bytes = prepared.bytes.len(); + let scheduled_at = schedule_commit.scheduled_at; + #[cfg(not(test))] + let sent_at = ClientTimestamp::now(); + #[cfg(test)] + let sent_at = test_timestamps + .map(|timestamps| timestamps.sent_at) + .unwrap_or_else(ClientTimestamp::now); + let machine_commit = match self + .machine + .finalize_probe_commit(machine_preflight, sent_at) + { + Ok(commit) => commit, + Err(error) => return Poll::Ready(Err(error)), + }; + #[cfg(not(test))] + let send_call_start = Instant::now(); + #[cfg(test)] + let send_call_start = test_timestamps + .map(|timestamps| timestamps.send_call_start) + .unwrap_or_else(Instant::now); + #[cfg(test)] + let send_result = self.test_hooks.try_send(&self.socket, &prepared.bytes); + #[cfg(not(test))] + let send_result = self.socket.try_send(&prepared.bytes); + let bytes = match send_result { + Ok(bytes) => bytes, + Err(error) if error.kind() == io::ErrorKind::WouldBlock => continue, + Err(error) => return Poll::Ready(Err(ClientError::Socket(error))), + }; + + #[cfg(not(test))] + let send_finished_at = Instant::now(); + #[cfg(test)] + let send_finished_at = test_timestamps + .map(|timestamps| timestamps.send_finished_at) + .unwrap_or_else(Instant::now); + let sent = self.machine.commit_probe_sent(machine_commit, bytes); + schedule.commit(schedule_commit); + self.prepared_probe = None; + + let send_call = send_finished_at.saturating_duration_since(send_call_start); + let timer_error = instant_abs_diff(sent_at.mono, scheduled_at); + if let Err(error) = validate_datagram_length(expected_bytes, bytes) { + return Poll::Ready(Err(error)); + } + events.push(echo_sent_event( + self.remote, + sent, + send_call, + scheduled_at, + timer_error, + )); + return Poll::Ready(Ok(events)); + } + } + + pub(crate) fn poll_recv( + &mut self, + cx: &mut Context<'_>, + ) -> Poll, ClientError>> { + if let Err(error) = self.machine.ensure_open() { + return Poll::Ready(Err(error)); + } + + loop { + match self.socket.poll_recv_ready(cx) { + Poll::Pending => return Poll::Pending, + Poll::Ready(Ok(())) => {} + Poll::Ready(Err(error)) => { + return Poll::Ready(Err(ClientError::Socket(error))); + } + } + #[cfg(test)] + if self.test_hooks.receive_would_block.get() > 0 { + self.test_hooks + .receive_would_block + .set(self.test_hooks.receive_would_block.get() - 1); + continue; + } + let datagram = match try_recv_tokio_datagram(&self.socket, &mut self.recv_buffer) { + Ok(datagram) => datagram, + Err(error) if error.kind() == io::ErrorKind::WouldBlock => continue, + Err(error) => return Poll::Ready(Err(ClientError::Socket(error))), + }; + + let events = match self.machine.process_received_echo_packet( + &self.recv_buffer[..datagram.len], + datagram.received_at, + datagram.meta, + ) { + Ok(events) => events, + Err(error) => return Poll::Ready(Err(error)), + }; + if self.machine.is_peer_closed() { + self.schedule = None; + self.prepared_probe = None; + if self.clear_peer_close_dscp().is_ok() { + self.applied_dscp = None; + } + } + return Poll::Ready(Ok(events)); + } + } + + pub(crate) fn poll_close( + &mut self, + cx: &mut Context<'_>, + ) -> Poll, ClientError>> { + loop { + let prepared = match self.machine.prepare_close() { + Ok(prepared) => prepared, + Err(error) => return Poll::Ready(Err(error)), + }; + let mut events = Vec::new(); + if let Err(source) = events.try_reserve(1) { + return Poll::Ready(Err(ClientError::AllocationFailed { + operation: "close event result", + source, + })); + } + let previous_dscp = self.applied_dscp; + let expected_bytes = prepared.bytes.len(); + + #[cfg(test)] if self.test_hooks.pause_close_before_writable.replace(false) { - future::pending::<()>().await; + return Poll::Pending; + } + match self.socket.poll_send_ready(cx) { + Poll::Pending => return Poll::Pending, + Poll::Ready(Ok(())) => {} + Poll::Ready(Err(error)) => { + return Poll::Ready(Err(ClientError::Socket(error))); + } + } + if let Err(error) = self.clear_close_dscp() { + return Poll::Ready(Err(error)); } - self.socket.writable().await?; - self.clear_close_dscp()?; let mut rollback = DscpRollback::armed( &self.socket, self.remote, @@ -413,14 +619,16 @@ impl AsyncClient { let bytes = match send_result { Ok(bytes) => bytes, Err(error) if error.kind() == io::ErrorKind::WouldBlock => { - rollback.restore()?; + if let Err(error) = rollback.restore() { + return Poll::Ready(Err(error)); + } #[cfg(test)] if self.test_hooks.pause_close_after_would_block.replace(false) { - future::pending::<()>().await; + return Poll::Pending; } continue; } - Err(error) => return Err(ClientError::Socket(error)), + Err(error) => return Poll::Ready(Err(ClientError::Socket(error))), }; #[cfg(not(test))] @@ -439,9 +647,11 @@ impl AsyncClient { self.prepared_probe = None; self.applied_dscp = None; - validate_datagram_length(expected_bytes, bytes)?; + if let Err(error) = validate_datagram_length(expected_bytes, bytes) { + return Poll::Ready(Err(error)); + } events.push(event); - return Ok(events); + return Poll::Ready(Ok(events)); } } @@ -472,55 +682,6 @@ impl AsyncClient { self.machine.is_peer_closed() } - async fn send_open_request(&self, deadline: Instant) -> Result { - let request = self - .prepared_open - .as_ref() - .expect("connected clients retain their prepared open request"); - loop { - #[cfg(test)] - if self.test_hooks.pause_open_before_writable.replace(false) { - future::pending::<()>().await; - } - if !wait_for_writable(&self.socket, deadline).await? { - return Ok(false); - } - #[cfg(test)] - let send_result = self.test_hooks.try_send(&self.socket, &request.bytes); - #[cfg(not(test))] - let send_result = self.socket.try_send(&request.bytes); - match send_result { - Ok(bytes) => { - validate_datagram_length(request.bytes.len(), bytes)?; - return Ok(true); - } - Err(error) if error.kind() == io::ErrorKind::WouldBlock => {} - Err(error) => return Err(ClientError::Socket(error)), - } - } - } - - async fn recv_open_datagram( - &self, - buffer: &mut [u8], - deadline: Instant, - ) -> Result, ClientError> { - loop { - #[cfg(test)] - if self.test_hooks.pause_open_before_readable.replace(false) { - future::pending::<()>().await; - } - if !wait_for_readable(&self.socket, deadline).await? { - return Ok(None); - } - match try_recv_tokio_datagram(&self.socket, buffer) { - Ok(datagram) => return Ok(Some(datagram)), - Err(error) if error.kind() == io::ErrorKind::WouldBlock => {} - Err(error) => return Err(ClientError::Socket(error)), - } - } - } - fn prepare_async_open( &mut self, machine: PreparedOpenAcceptance, @@ -603,27 +764,45 @@ impl AsyncClient { outcome } - async fn send_cleanup_close_best_effort(&self, packet: Option<&[u8]>, deadline: Instant) { - let Some(packet) = packet else { - return; - }; + fn poll_open_cleanup( + &self, + state: &mut AsyncOpenState, + cx: &mut Context<'_>, + ) -> Poll> { loop { + let Some(packet) = state + .cleanup + .as_ref() + .expect("cleanup polling requires a retained primary error") + .packet() + else { + return Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())); + }; + #[cfg(test)] if self.test_hooks.fail_cleanup_send.replace(false) { - return; + return Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())); } #[cfg(test)] let send_result = self.test_hooks.try_send(&self.socket, packet); #[cfg(not(test))] let send_result = self.socket.try_send(packet); match send_result { - Ok(_) => return, + Ok(_) => { + return Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())); + } Err(error) if error.kind() == io::ErrorKind::WouldBlock => {} - Err(_) => return, + Err(_) => { + return Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())); + } } - match wait_for_writable(&self.socket, deadline).await { - Ok(true) => {} - Ok(false) | Err(_) => return, + + match poll_writable_until(&self.socket, state.deadline_timer(), cx) { + Poll::Pending => return Poll::Pending, + Poll::Ready(Ok(true)) => {} + Poll::Ready(Ok(false) | Err(_)) => { + return Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())); + } } } } @@ -711,29 +890,33 @@ impl Drop for DscpRollback<'_> { } } -async fn wait_for_writable( +fn poll_writable_until( socket: &tokio::net::UdpSocket, - deadline: Instant, -) -> Result { - match time::timeout_at(deadline.into(), socket.writable()).await { - Ok(result) => { - result?; - Ok(true) - } - Err(_) => Ok(false), + deadline: &mut Pin>, + cx: &mut Context<'_>, +) -> Poll> { + match socket.poll_send_ready(cx) { + Poll::Ready(Ok(())) => Poll::Ready(Ok(true)), + Poll::Ready(Err(error)) => Poll::Ready(Err(ClientError::Socket(error))), + Poll::Pending => match deadline.as_mut().poll(cx) { + Poll::Ready(()) => Poll::Ready(Ok(false)), + Poll::Pending => Poll::Pending, + }, } } -async fn wait_for_readable( +fn poll_readable_until( socket: &tokio::net::UdpSocket, - deadline: Instant, -) -> Result { - match time::timeout_at(deadline.into(), socket.readable()).await { - Ok(result) => { - result?; - Ok(true) - } - Err(_) => Ok(false), + deadline: &mut Pin>, + cx: &mut Context<'_>, +) -> Poll> { + match socket.poll_recv_ready(cx) { + Poll::Ready(Ok(())) => Poll::Ready(Ok(true)), + Poll::Ready(Err(error)) => Poll::Ready(Err(ClientError::Socket(error))), + Poll::Pending => match deadline.as_mut().poll(cx) { + Poll::Ready(()) => Poll::Ready(Ok(false)), + Poll::Pending => Poll::Pending, + }, } } -- 2.54.0