Files
irtt-rs/crates/irtt-client/src/async_client.rs
T

1208 lines
44 KiB
Rust

use std::{
future::{poll_fn, Future},
io,
net::SocketAddr,
pin::Pin,
task::{Context, Poll},
time::{Duration, Instant},
};
#[cfg(test)]
use std::{
cell::{Cell, RefCell},
collections::VecDeque,
};
use tokio::time;
use crate::{
client::{
echo_sent_event,
schedule::{instant_abs_diff, ProbeSchedule},
validate_datagram_length,
},
config::ClientConfig,
error::ClientError,
event::{ClientEvent, OpenOutcome},
receive::{drain_tx_timestamps, try_enable_tx_timestamping, try_recv_tokio_datagram},
session::machine::{
recv_buffer_size, OpenDatagramDisposition, PreparedOpenAcceptance, PreparedOpenRequest,
PreparedProbe, SessionMachine, TimeoutBatch, MAX_OPEN_PACKET_SIZE,
},
socket::{connect_tokio_udp_socket, resolve_remote_tokio, validate_open_timeouts},
socket_options::{apply_traffic_class_to_tokio_socket, clear_dscp_on_tokio_socket},
timing::ClientTimestamp,
};
/// Maximum amount of immediately available opening work performed by one poll.
///
/// Opening may consume ignored UDP datagrams, roll over an expired attempt, or
/// encounter readiness false positives without ever awaiting. Bound that work
/// so one managed target cannot retain the current-thread runtime.
const OPEN_POLL_WORK_BUDGET: usize = 64;
/// Maximum immediately available receive work performed by one
/// [`AsyncClient::poll_recv`] poll.
///
/// A `WouldBlock` readiness false positive is naturally bounded: `try_io`
/// clears the reactor's readiness on `WouldBlock`, so the next
/// `poll_recv_ready` genuinely awaits a fresh edge. An `Interrupted`
/// (`EINTR`) nonblocking `recvmsg` does not clear that readiness — the
/// syscall was interrupted, not told there was nothing to read — so without
/// an explicit bound a sustained signal delivery rate could otherwise retry
/// in a tight loop on this executor thread. Bounding both cases the same way
/// this crate already bounds open-attempt work keeps that impossible.
const RECV_POLL_WORK_BUDGET: usize = 64;
#[cfg(test)]
use crate::client::ProbeSendTimestamps;
#[derive(Debug)]
struct PreparedAsyncOpen {
machine: PreparedOpenAcceptance,
schedule: Option<ProbeSchedule>,
recv_buffer_len: Option<usize>,
negotiated_traffic_class: Option<u8>,
}
#[derive(Debug)]
struct PreparedAsyncOpenFailure {
primary: ClientError,
machine: PreparedOpenAcceptance,
}
#[derive(Debug)]
enum AsyncOpenCleanup {
Acceptance {
primary: ClientError,
packet: Option<Box<[u8]>>,
},
Adapter(Box<PreparedAsyncOpenFailure>),
}
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<Instant>,
deadline_timer: Option<Pin<Box<time::Sleep>>>,
request_submitted: bool,
stop_after_current_attempt: bool,
stopped_after_current_attempt: bool,
buffer: [u8; MAX_OPEN_PACKET_SIZE],
cleanup: Option<AsyncOpenCleanup>,
}
impl AsyncOpenState {
pub(crate) fn new() -> Self {
Self {
attempt: 0,
deadline: None,
deadline_timer: None,
request_submitted: false,
stop_after_current_attempt: false,
stopped_after_current_attempt: 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;
self.stopped_after_current_attempt = self.stop_after_current_attempt;
}
pub(crate) fn has_in_flight_work(&self) -> bool {
self.request_submitted || self.cleanup.is_some()
}
pub(crate) fn request_stop_after_current_attempt(&mut self) {
self.stop_after_current_attempt = true;
}
pub(crate) fn stopped_after_current_attempt(&self) -> bool {
self.stopped_after_current_attempt
}
#[cfg(test)]
pub(crate) fn has_retained_cleanup(&self) -> bool {
self.cleanup.is_some()
}
fn deadline(&self) -> Instant {
self.deadline
.expect("an open attempt has an active deadline")
}
fn deadline_timer(&mut self) -> &mut Pin<Box<time::Sleep>> {
self.deadline_timer
.as_mut()
.expect("an open attempt has an active deadline timer")
}
}
#[cfg(test)]
#[derive(Debug)]
enum InjectedSend {
WouldBlock,
Error,
ReportedLength(usize),
}
#[cfg(test)]
#[derive(Debug, Default)]
struct AsyncClientTestHooks {
sends: RefCell<VecDeque<InjectedSend>>,
send_attempts: Cell<usize>,
receive_would_block: Cell<usize>,
/// Number of subsequent receive attempts (`poll_open` or `poll_recv`)
/// that should report `io::ErrorKind::Interrupted` instead of performing
/// a real nonblocking receive, deterministically exercising interrupted-
/// syscall retry behavior.
inject_recv_interrupted: Cell<usize>,
fail_open_dscp: Cell<bool>,
fail_cleanup_send: Cell<bool>,
fail_peer_close_dscp: Cell<bool>,
fail_dscp_restore: Cell<bool>,
pause_open_cleanup_before_send: Cell<bool>,
pause_open_before_writable: Cell<bool>,
pause_open_before_readable: Cell<bool>,
pause_probe_before_writable: Cell<bool>,
pause_close_before_writable: Cell<bool>,
pause_close_after_would_block: Cell<bool>,
probe_timestamps: RefCell<VecDeque<ProbeSendTimestamps>>,
close_sent_at: Cell<Option<ClientTimestamp>>,
}
#[cfg(test)]
impl AsyncClientTestHooks {
fn try_send(&self, socket: &tokio::net::UdpSocket, packet: &[u8]) -> io::Result<usize> {
self.send_attempts.set(self.send_attempts.get() + 1);
if let Some(injected) = self.sends.borrow_mut().pop_front() {
return match injected {
InjectedSend::WouldBlock => Err(io::Error::from(io::ErrorKind::WouldBlock)),
InjectedSend::Error => Err(io::Error::other("injected send failure")),
InjectedSend::ReportedLength(bytes) => Ok(bytes),
};
}
socket.try_send(packet)
}
fn take_probe_timestamps(&self) -> Option<ProbeSendTimestamps> {
self.probe_timestamps.borrow_mut().pop_front()
}
/// Consumes one pending injected interruption, if any, reporting whether
/// the caller should report `io::ErrorKind::Interrupted` instead of
/// attempting a real receive.
fn take_injected_interrupted(&self) -> bool {
let remaining = self.inject_recv_interrupted.get();
if remaining == 0 {
return false;
}
self.inject_recv_interrupted.set(remaining - 1);
true
}
}
/// Low-level Tokio IRTT client for one connected UDP target.
///
/// `AsyncClient` does not construct, own, or store a Tokio runtime. Its async
/// methods are polled by the caller and require a current Tokio runtime with I/O
/// and time enabled.
///
/// # Example
///
/// Use this low-level adapter when the application already owns Tokio and wants
/// to drive the session lifecycle itself. This example type-checks without
/// contacting a server.
///
/// ```no_run
/// use irtt_client::{AsyncClient, ClientConfig};
///
/// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
/// let mut client = AsyncClient::connect(ClientConfig::default()).await?;
/// let outcome = client.open().await?;
/// println!("opened: {outcome:?}");
/// let sent = client.send_probe().await?;
/// println!("sent: {sent:?}");
/// let received = client.recv().await?;
/// println!("received: {received:?}");
/// let timed_out = client.poll_timeouts()?;
/// println!("timeouts: {timed_out:?}");
/// let closed = client.close().await?;
/// println!("closed: {closed:?}");
/// # Ok(())
/// # }
/// ```
#[derive(Debug)]
pub struct AsyncClient {
socket: tokio::net::UdpSocket,
machine: SessionMachine,
schedule: Option<ProbeSchedule>,
remote: SocketAddr,
recv_buffer: Vec<u8>,
applied_traffic_class: Option<u8>,
prepared_open: Option<PreparedOpenRequest>,
prepared_probe: Option<PreparedProbe>,
/// See `Client::tx_timestamping_enabled` in `client.rs` for the
/// invariant this tracks; it is the same adapter-owned capability flag.
tx_timestamping_enabled: bool,
#[cfg(test)]
test_hooks: AsyncClientTestHooks,
}
impl AsyncClient {
/// Resolve the configured server and construct one connected Tokio UDP
/// socket.
///
/// Polling this future without a current Tokio runtime returns
/// [`ClientError::NoTokioRuntime`]. A runtime without enabled I/O or time
/// drivers is outside this type's runtime contract.
pub async fn connect(config: ClientConfig) -> Result<Self, ClientError> {
tokio::runtime::Handle::try_current().map_err(|_| ClientError::NoTokioRuntime)?;
validate_open_timeouts(&config.open_timeouts)?;
let remote = resolve_remote_tokio(&config).await?;
let machine = SessionMachine::new(config.clone(), remote)?;
let prepared_open = machine.prepare_open_request()?;
let socket = connect_tokio_udp_socket(&config.socket_config, remote)?;
Ok(Self {
socket,
machine,
schedule: None,
remote,
recv_buffer: vec![0_u8; recv_buffer_size(false, None)?],
applied_traffic_class: None,
prepared_open: Some(prepared_open),
prepared_probe: None,
tx_timestamping_enabled: false,
#[cfg(test)]
test_hooks: AsyncClientTestHooks::default(),
})
}
/// Perform the IRTT open exchange.
///
/// Each configured attempt uses one absolute deadline for one successful
/// request submission and all replies inspected during that attempt.
/// Malformed, unrelated, wrong-direction, and unauthenticated datagrams are
/// ignored without resending the request.
///
/// Dropping this future before trusted acceptance leaves the local machine
/// connected with no committed schedule or negotiated adapter state. If it
/// 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<OpenOutcome, ClientError> {
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
/// commits.
///
/// A prepared packet is retained across readiness false positives,
/// 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<Vec<ClientEvent>, ClientError> {
poll_fn(|cx| self.poll_send_probe(cx)).await
}
/// Await and classify one complete UDP datagram.
///
/// State is validated before readiness is awaited. Readiness false positives
/// 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<Vec<ClientEvent>, ClientError> {
poll_fn(|cx| self.poll_recv(cx)).await
}
/// Poll protocol timeouts using a newly captured monotonic timestamp.
pub fn poll_timeouts(&mut self) -> Result<Vec<ClientEvent>, ClientError> {
self.poll_timeouts_at(Instant::now())
}
/// Poll protocol timeouts using the caller's monotonic timestamp.
pub fn poll_timeouts_at(&mut self, now: Instant) -> Result<Vec<ClientEvent>, ClientError> {
self.drain_tx_timestamps()?;
self.machine.poll_timeouts_at(now)
}
pub(crate) fn poll_timeouts_bounded_at(
&mut self,
now: Instant,
limit: usize,
) -> Result<TimeoutBatch, ClientError> {
self.drain_tx_timestamps()?;
self.machine.poll_timeouts_bounded_at(now, limit)
}
/// Send the retained close packet and commit local close exactly once.
///
/// DSCP is cleared only after write readiness and immediately before the
/// nonblocking send attempt. A `WouldBlock` result restores it before the
/// 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<Vec<ClientEvent>, ClientError> {
poll_fn(|cx| self.poll_close(cx)).await
}
pub(crate) fn poll_open(
&mut self,
state: &mut AsyncOpenState,
cx: &mut Context<'_>,
) -> Poll<Result<OpenOutcome, ClientError>> {
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");
}
for _ in 0..OPEN_POLL_WORK_BUDGET {
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.stop_after_current_attempt && state.deadline.is_none() {
state.stopped_after_current_attempt = true;
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) => {
state.request_submitted = true;
if let Err(error) = validate_datagram_length(request.bytes.len(), bytes) {
return Poll::Ready(Err(error));
}
}
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();
if state.stopped_after_current_attempt() {
return Poll::Ready(Err(ClientError::OpenTimeout));
}
continue;
}
Poll::Ready(Ok(true)) => {}
}
#[cfg(test)]
if self.test_hooks.take_injected_interrupted() {
continue;
}
let datagram = match try_recv_tokio_datagram(&self.socket, &mut state.buffer) {
Ok(datagram) => datagram,
// WouldBlock is a readiness false positive; Interrupted
// (EINTR) is a transient nonblocking-recvmsg interruption.
// Both retry through the surrounding `OPEN_POLL_WORK_BUDGET`-
// bounded loop above, which re-checks readiness before the
// next attempt and yields back to the executor once that
// budget is exhausted.
Err(error)
if matches!(
error.kind(),
io::ErrorKind::WouldBlock | io::ErrorKind::Interrupted
) =>
{
continue
}
Err(error) => return Poll::Ready(Err(ClientError::Socket(error))),
};
if datagram.received_at.mono > state.deadline() {
state.finish_attempt();
if state.stopped_after_current_attempt() {
return Poll::Ready(Err(ClientError::OpenTimeout));
}
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)));
}
// We deliberately truncated known immediate opening work. Arrange for
// another poll even if no new socket-readiness edge occurs.
cx.waker().wake_by_ref();
Poll::Pending
}
pub(crate) fn poll_send_probe(
&mut self,
cx: &mut Context<'_>,
) -> Poll<Result<Vec<ClientEvent>, 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;
// Private pre-send anchor: finalizes all fallible commit work
// (timeout deadline arithmetic, kernel TX lower bound) before
// the send attempt. A WouldBlock retry samples a fresh one on
// the next loop iteration; only the successful attempt's anchor
// is retained (via `machine_commit`/`commit_probe_sent` below).
#[cfg(not(test))]
let send_anchor = ClientTimestamp::now();
#[cfg(test)]
let send_anchor = test_timestamps
.map(|timestamps| timestamps.send_anchor)
.unwrap_or_else(ClientTimestamp::now);
let machine_commit = match self
.machine
.finalize_probe_commit(machine_preflight, send_anchor)
{
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,
// No sent_at measurement, no machine commit, no schedule
// commit: the transaction is fully discarded and retried
// with a fresh send_anchor on the next loop iteration.
Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
self.machine.invalidate_kernel_tx_correlation();
continue;
}
Err(error) => {
self.machine.invalidate_kernel_tx_correlation();
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);
// Public post-send measurement, captured only after the
// successful send above. The subsequent machine/schedule commit
// is infallible.
#[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 sent = self
.machine
.commit_probe_sent(machine_commit, sent_at, 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));
}
if let Err(error) = self.drain_tx_timestamps() {
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<Result<Vec<ClientEvent>, ClientError>> {
if let Err(error) = self.machine.ensure_open() {
return Poll::Ready(Err(error));
}
for _ in 0..RECV_POLL_WORK_BUDGET {
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;
}
#[cfg(test)]
if self.test_hooks.take_injected_interrupted() {
continue;
}
let datagram = match try_recv_tokio_datagram(&self.socket, &mut self.recv_buffer) {
Ok(datagram) => datagram,
// See `RECV_POLL_WORK_BUDGET`: both a readiness false
// positive and an interrupted nonblocking recvmsg retry
// within this bounded loop rather than failing the receive.
Err(error)
if matches!(
error.kind(),
io::ErrorKind::WouldBlock | io::ErrorKind::Interrupted
) =>
{
continue
}
Err(error) => return Poll::Ready(Err(ClientError::Socket(error))),
};
// One final bounded drain before processing the reply, so a TX
// timestamp that raced the reply is still associated with its
// probe before that probe is looked up and possibly removed
// below.
if let Err(error) = self.drain_tx_timestamps() {
return Poll::Ready(Err(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_traffic_class = None;
}
}
return Poll::Ready(Ok(events));
}
// Truncated known immediate receive work, exactly as `poll_open`
// does: arrange for another poll even if no new socket-readiness
// edge occurs, rather than spinning this executor thread further.
cx.waker().wake_by_ref();
Poll::Pending
}
pub(crate) fn poll_close(
&mut self,
cx: &mut Context<'_>,
) -> Poll<Result<Vec<ClientEvent>, 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_traffic_class = self.applied_traffic_class;
let expected_bytes = prepared.bytes.len();
#[cfg(test)]
if self.test_hooks.pause_close_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)));
}
}
if let Err(error) = self.clear_close_dscp() {
return Poll::Ready(Err(error));
}
let mut rollback = DscpRollback::armed(
&self.socket,
self.remote,
previous_traffic_class,
#[cfg(test)]
&self.test_hooks.fail_dscp_restore,
);
#[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 => {
if let Err(error) = rollback.restore() {
return Poll::Ready(Err(error));
}
#[cfg(test)]
if self.test_hooks.pause_close_after_would_block.replace(false) {
return Poll::Pending;
}
continue;
}
Err(error) => return Poll::Ready(Err(ClientError::Socket(error))),
};
#[cfg(not(test))]
let close_sent_at = ClientTimestamp::now();
#[cfg(test)]
let close_sent_at = self
.test_hooks
.close_sent_at
.take()
.unwrap_or_else(ClientTimestamp::now);
let event = self
.machine
.commit_local_close(prepared.commit, close_sent_at);
rollback.disarm();
self.schedule = None;
self.prepared_probe = None;
self.applied_traffic_class = None;
if let Err(error) = validate_datagram_length(expected_bytes, bytes) {
return Poll::Ready(Err(error));
}
events.push(event);
return Poll::Ready(Ok(events));
}
}
/// Return the next scheduled probe deadline.
pub fn next_send_deadline(&self) -> Option<Instant> {
if !self.machine.is_open() {
return None;
}
self.schedule.as_ref()?.next_send_deadline()
}
pub(crate) fn probe_interval(&self) -> Option<Duration> {
if !self.machine.is_open() {
return None;
}
self.schedule.as_ref().map(ProbeSchedule::interval)
}
/// Return the configured local probe timeout.
pub fn probe_timeout(&self) -> Duration {
self.machine.probe_timeout()
}
/// Return whether the current run has completed.
pub fn is_run_complete(&self) -> bool {
self.machine.is_terminal()
|| self
.schedule
.as_ref()
.is_some_and(|schedule| schedule.is_finished() && self.machine.pending_is_empty())
}
/// Return whether an authenticated peer close ended the session.
pub fn is_peer_closed(&self) -> bool {
self.machine.is_peer_closed()
}
pub(crate) fn remote_addr(&self) -> SocketAddr {
self.remote
}
pub(crate) fn packets_sent(&self) -> u64 {
self.machine.packets_sent()
}
pub(crate) fn next_probe_timeout_deadline(&self) -> Option<Instant> {
self.machine.next_probe_timeout_deadline()
}
pub(crate) fn latest_probe_timeout_deadline(&self) -> Option<Instant> {
self.machine.latest_probe_timeout_deadline()
}
pub(crate) fn discard_prepared_probe(&mut self) {
self.prepared_probe = None;
}
#[cfg(test)]
pub(crate) fn retain_open_cleanup_once(&self) {
self.test_hooks.fail_open_dscp.set(true);
self.test_hooks.pause_open_cleanup_before_send.set(true);
}
#[cfg(test)]
pub(crate) fn poll_recv_ready_for_test(&self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.socket.poll_recv_ready(cx)
}
#[cfg(test)]
pub(crate) fn remove_pending_for_test(
&mut self,
wire_seq: u32,
) -> Option<crate::probe::PendingProbe> {
self.machine.remove_pending_for_test(wire_seq)
}
#[cfg(test)]
pub(crate) fn replace_pending_for_test(&mut self, probe: crate::probe::PendingProbe) {
self.machine.replace_pending_for_test(probe);
}
pub(crate) fn skip_missed_probe_slots_at(&mut self, now: Instant) -> Result<(), ClientError> {
let Some(schedule) = self.schedule.as_mut() else {
return Ok(());
};
let Some(deadline) = schedule.next_send_deadline() else {
return Ok(());
};
if deadline > now {
return Ok(());
}
let commit = schedule.preflight_managed_commit(deadline, now)?;
schedule.commit(commit);
Ok(())
}
fn prepare_async_open(
&mut self,
machine: PreparedOpenAcceptance,
opened_at: ClientTimestamp,
) -> Result<PreparedAsyncOpen, Box<PreparedAsyncOpenFailure>> {
let Some(negotiated) = machine.normal_negotiated() else {
return Ok(PreparedAsyncOpen {
machine,
schedule: None,
recv_buffer_len: None,
negotiated_traffic_class: None,
});
};
let schedule = match ProbeSchedule::new(opened_at.mono, negotiated) {
Ok(schedule) => schedule,
Err(primary) => {
return Err(Box::new(PreparedAsyncOpenFailure { primary, machine }));
}
};
let recv_buffer_len = match recv_buffer_size(self.machine.has_hmac(), Some(negotiated)) {
Ok(size) => size,
Err(primary) => {
return Err(Box::new(PreparedAsyncOpenFailure { primary, machine }));
}
};
let negotiated_traffic_class = match u8::try_from(negotiated.params.dscp) {
Ok(traffic_class) => traffic_class,
Err(_) => {
return Err(Box::new(PreparedAsyncOpenFailure {
primary: ClientError::InvalidConfig {
reason: "negotiated dscp must be in range 0..=255".to_owned(),
},
machine,
}));
}
};
let additional = recv_buffer_len.saturating_sub(self.recv_buffer.len());
if let Err(source) = self.recv_buffer.try_reserve(additional) {
return Err(Box::new(PreparedAsyncOpenFailure {
primary: ClientError::AllocationFailed {
operation: "negotiated receive buffer",
source,
},
machine,
}));
}
#[cfg(test)]
let dscp_result = if self.test_hooks.fail_open_dscp.replace(false) {
Err(ClientError::SocketOption {
operation: "set negotiated DSCP",
remote: self.remote,
source: io::Error::other("injected negotiated DSCP failure"),
})
} else {
apply_traffic_class_to_tokio_socket(&self.socket, self.remote, negotiated_traffic_class)
};
#[cfg(not(test))]
let dscp_result = apply_traffic_class_to_tokio_socket(
&self.socket,
self.remote,
negotiated_traffic_class,
);
if let Err(primary) = dscp_result {
self.restore_dscp_best_effort(self.applied_traffic_class);
return Err(Box::new(PreparedAsyncOpenFailure { primary, machine }));
}
Ok(PreparedAsyncOpen {
machine,
schedule: Some(schedule),
recv_buffer_len: Some(recv_buffer_len),
negotiated_traffic_class: Some(negotiated_traffic_class),
})
}
fn commit_async_open(&mut self, prepared: PreparedAsyncOpen) -> OpenOutcome {
if let Some(recv_buffer_len) = prepared.recv_buffer_len {
self.recv_buffer.resize(recv_buffer_len, 0);
}
let outcome = self.machine.commit_open(prepared.machine);
self.schedule = prepared.schedule;
self.applied_traffic_class = prepared.negotiated_traffic_class;
self.prepared_open = None;
self.tx_timestamping_enabled = try_enable_tx_timestamping(&self.socket);
outcome
}
/// See `Client::drain_tx_timestamps` in `client.rs`: same bounded,
/// nonblocking, best-effort drain, adapted to the Tokio socket.
fn drain_tx_timestamps(&mut self) -> Result<(), ClientError> {
if !self.tx_timestamping_enabled {
return Ok(());
}
let machine = &mut self.machine;
drain_tx_timestamps(&self.socket, |id, timestamp| {
machine.record_kernel_tx_timestamp(id, timestamp);
})
.map_err(ClientError::Socket)
}
fn poll_open_cleanup(
&self,
state: &mut AsyncOpenState,
cx: &mut Context<'_>,
) -> Poll<Result<OpenOutcome, ClientError>> {
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
.pause_open_cleanup_before_send
.replace(false)
{
return Poll::Pending;
}
#[cfg(test)]
if self.test_hooks.fail_cleanup_send.replace(false) {
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 Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())),
Err(error) if error.kind() == io::ErrorKind::WouldBlock => {}
Err(_) => return Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())),
}
match poll_writable_until(&self.socket, state.deadline_timer(), cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(Ok(true)) => {
// Retain cleanup state and resume the send in a fresh poll. A
// writable readiness false positive must not spin this future.
cx.waker().wake_by_ref();
Poll::Pending
}
Poll::Ready(Ok(false) | Err(_)) => {
Poll::Ready(Err(state.cleanup.take().unwrap().into_primary()))
}
}
}
fn clear_peer_close_dscp(&self) -> Result<(), ClientError> {
#[cfg(test)]
if self.test_hooks.fail_peer_close_dscp.replace(false) {
return Err(ClientError::SocketOption {
operation: "clear DSCP before close",
remote: self.remote,
source: io::Error::other("injected peer-close DSCP cleanup failure"),
});
}
clear_dscp_on_tokio_socket(&self.socket, self.remote)
}
fn clear_close_dscp(&self) -> Result<(), ClientError> {
clear_dscp_on_tokio_socket(&self.socket, self.remote)
}
fn restore_dscp_best_effort(&self, previous_traffic_class: Option<u8>) {
let _ = match previous_traffic_class {
Some(traffic_class) => {
apply_traffic_class_to_tokio_socket(&self.socket, self.remote, traffic_class)
}
None => clear_dscp_on_tokio_socket(&self.socket, self.remote),
};
}
}
struct DscpRollback<'a> {
socket: &'a tokio::net::UdpSocket,
remote: SocketAddr,
previous_traffic_class: Option<u8>,
armed: bool,
#[cfg(test)]
fail_restore: &'a Cell<bool>,
}
impl<'a> DscpRollback<'a> {
fn armed(
socket: &'a tokio::net::UdpSocket,
remote: SocketAddr,
previous_traffic_class: Option<u8>,
#[cfg(test)] fail_restore: &'a Cell<bool>,
) -> Self {
Self {
socket,
remote,
previous_traffic_class,
armed: true,
#[cfg(test)]
fail_restore,
}
}
fn restore(&mut self) -> Result<(), ClientError> {
if !self.armed {
return Ok(());
}
#[cfg(test)]
if self.fail_restore.replace(false) {
return Err(ClientError::SocketOption {
operation: "restore negotiated DSCP",
remote: self.remote,
source: io::Error::other("injected DSCP restore failure"),
});
}
match self.previous_traffic_class {
Some(traffic_class) => {
apply_traffic_class_to_tokio_socket(self.socket, self.remote, traffic_class)?
}
None => clear_dscp_on_tokio_socket(self.socket, self.remote)?,
}
self.armed = false;
Ok(())
}
fn disarm(&mut self) {
self.armed = false;
}
}
impl Drop for DscpRollback<'_> {
fn drop(&mut self) {
if self.armed {
let _ = self.restore();
}
}
}
fn poll_writable_until(
socket: &tokio::net::UdpSocket,
deadline: &mut Pin<Box<time::Sleep>>,
cx: &mut Context<'_>,
) -> Poll<Result<bool, ClientError>> {
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,
},
}
}
fn poll_readable_until(
socket: &tokio::net::UdpSocket,
deadline: &mut Pin<Box<time::Sleep>>,
cx: &mut Context<'_>,
) -> Poll<Result<bool, ClientError>> {
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,
},
}
}
#[cfg(test)]
mod tests;