Merge pull request #47 from Lochnair/feat/tokio-async-client

feat(client): add Tokio async client

Cargo.lock

@@ -713,6 +713,7 @@ dependencies = [
  "libc",
  "socket2",
  "thiserror 2.0.19",
+ "tokio",
 ]
 
 [[package]]
@@ -1716,6 +1717,19 @@ version = "0.1.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109"
 
+[[package]]
+name = "tokio"
+version = "1.53.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed"
+dependencies = [
+ "libc",
+ "mio",
+ "pin-project-lite",
+ "socket2",
+ "windows-sys",
+]
+
 [[package]]
 name = "typenum"
 version = "1.20.1"

crates/irtt-client/Cargo.toml

@@ -16,10 +16,12 @@ irtt-proto.workspace = true
 libc = { version = "0.2.189", optional = true }
 socket2 = { version = "0.6.5", features = ["all"] }
 thiserror = "2.0.19"
+tokio = { version = "1", default-features = false, features = ["rt", "net", "time", "sync"], optional = true}
 
 [features]
 default = []
 ancillary = ["dep:libc"]
+tokio = ["dep:tokio"]
 
 [package.metadata.docs.rs]
 all-features = true

crates/irtt-client/src/async_client.rs

@@ -0,0 +1,741 @@
+use std::{
+    io,
+    net::SocketAddr,
+    time::{Duration, Instant},
+};
+
+#[cfg(test)]
+use std::{
+    cell::{Cell, RefCell},
+    collections::VecDeque,
+    future,
+};
+
+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::try_recv_tokio_datagram,
+    session::machine::{
+        recv_buffer_size, OpenDatagramDisposition, PreparedOpenAcceptance, PreparedOpenRequest,
+        PreparedProbe, SessionMachine, MAX_OPEN_PACKET_SIZE,
+    },
+    socket::{connect_tokio_udp_socket, resolve_remote_tokio, validate_open_timeouts},
+    socket_options::{apply_dscp_to_tokio_socket, clear_dscp_on_tokio_socket},
+    timing::ClientTimestamp,
+};
+
+#[cfg(test)]
+use crate::client::ProbeSendTimestamps;
+
+#[derive(Debug)]
+struct PreparedAsyncOpen {
+    machine: PreparedOpenAcceptance,
+    schedule: Option<ProbeSchedule>,
+    recv_buffer_len: Option<usize>,
+    negotiated_dscp: Option<u8>,
+}
+
+#[derive(Debug)]
+struct PreparedAsyncOpenFailure {
+    primary: ClientError,
+    machine: PreparedOpenAcceptance,
+}
+
+#[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>,
+    fail_open_dscp: Cell<bool>,
+    fail_cleanup_send: Cell<bool>,
+    fail_peer_close_dscp: Cell<bool>,
+    fail_dscp_restore: 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()
+    }
+}
+
+/// 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.
+#[derive(Debug)]
+pub struct AsyncClient {
+    socket: tokio::net::UdpSocket,
+    machine: SessionMachine,
+    schedule: Option<ProbeSchedule>,
+    remote: SocketAddr,
+    recv_buffer: Vec<u8>,
+    applied_dscp: Option<u8>,
+    prepared_open: Option<PreparedOpenRequest>,
+    prepared_probe: Option<PreparedProbe>,
+    #[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_dscp: None,
+            prepared_open: Some(prepared_open),
+            prepared_probe: None,
+            #[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> {
+        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)
+    }
+
+    /// 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> {
+        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);
+        }
+    }
+
+    /// 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> {
+        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 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.machine.poll_timeouts_at(now)
+    }
+
+    /// 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> {
+        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();
+
+        loop {
+            #[cfg(test)]
+            if self.test_hooks.pause_close_before_writable.replace(false) {
+                future::pending::<()>().await;
+            }
+            self.socket.writable().await?;
+            self.clear_close_dscp()?;
+            let mut rollback = DscpRollback::armed(
+                &self.socket,
+                self.remote,
+                previous_dscp,
+                #[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 => {
+                    rollback.restore()?;
+                    #[cfg(test)]
+                    if self.test_hooks.pause_close_after_would_block.replace(false) {
+                        future::pending::<()>().await;
+                    }
+                    continue;
+                }
+                Err(error) => return 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_dscp = None;
+
+            validate_datagram_length(expected_bytes, bytes)?;
+            events.push(event);
+            return 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()
+    }
+
+    /// 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()
+    }
+
+    async fn send_open_request(&self, deadline: Instant) -> Result<bool, ClientError> {
+        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<Option<crate::receive::ReceivedDatagram>, 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,
+        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_dscp: 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_dscp = match u8::try_from(negotiated.params.dscp) {
+            Ok(dscp) => dscp,
+            Err(_) => {
+                return Err(Box::new(PreparedAsyncOpenFailure {
+                    primary: ClientError::InvalidConfig {
+                        reason: "negotiated dscp must be in range 0..=63".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_dscp_to_tokio_socket(&self.socket, self.remote, negotiated_dscp)
+        };
+        #[cfg(not(test))]
+        let dscp_result = apply_dscp_to_tokio_socket(&self.socket, self.remote, negotiated_dscp);
+        if let Err(primary) = dscp_result {
+            self.restore_dscp_best_effort(self.applied_dscp);
+            return Err(Box::new(PreparedAsyncOpenFailure { primary, machine }));
+        }
+
+        Ok(PreparedAsyncOpen {
+            machine,
+            schedule: Some(schedule),
+            recv_buffer_len: Some(recv_buffer_len),
+            negotiated_dscp: Some(negotiated_dscp),
+        })
+    }
+
+    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_dscp = prepared.negotiated_dscp;
+        self.prepared_open = None;
+        outcome
+    }
+
+    async fn send_cleanup_close_best_effort(&self, packet: Option<&[u8]>, deadline: Instant) {
+        let Some(packet) = packet else {
+            return;
+        };
+        loop {
+            #[cfg(test)]
+            if self.test_hooks.fail_cleanup_send.replace(false) {
+                return;
+            }
+            #[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,
+                Err(error) if error.kind() == io::ErrorKind::WouldBlock => {}
+                Err(_) => return,
+            }
+            match wait_for_writable(&self.socket, deadline).await {
+                Ok(true) => {}
+                Ok(false) | Err(_) => return,
+            }
+        }
+    }
+
+    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_dscp: Option<u8>) {
+        let _ = match previous_dscp {
+            Some(dscp) => apply_dscp_to_tokio_socket(&self.socket, self.remote, dscp),
+            None => clear_dscp_on_tokio_socket(&self.socket, self.remote),
+        };
+    }
+}
+
+struct DscpRollback<'a> {
+    socket: &'a tokio::net::UdpSocket,
+    remote: SocketAddr,
+    previous_dscp: 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_dscp: Option<u8>,
+        #[cfg(test)] fail_restore: &'a Cell<bool>,
+    ) -> Self {
+        Self {
+            socket,
+            remote,
+            previous_dscp,
+            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_dscp {
+            Some(dscp) => apply_dscp_to_tokio_socket(self.socket, self.remote, dscp)?,
+            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();
+        }
+    }
+}
+
+async fn wait_for_writable(
+    socket: &tokio::net::UdpSocket,
+    deadline: Instant,
+) -> Result<bool, ClientError> {
+    match time::timeout_at(deadline.into(), socket.writable()).await {
+        Ok(result) => {
+            result?;
+            Ok(true)
+        }
+        Err(_) => Ok(false),
+    }
+}
+
+async fn wait_for_readable(
+    socket: &tokio::net::UdpSocket,
+    deadline: Instant,
+) -> Result<bool, ClientError> {
+    match time::timeout_at(deadline.into(), socket.readable()).await {
+        Ok(result) => {
+            result?;
+            Ok(true)
+        }
+        Err(_) => Ok(false),
+    }
+}
+
+#[cfg(test)]
+mod tests;

crates/irtt-client/src/error.rs

@@ -12,6 +12,10 @@ use thiserror::Error;
 #[derive(Debug, Error)]
 #[non_exhaustive]
 pub enum ClientError {
+    /// A Tokio-backed client operation was polled without a current runtime.
+    #[cfg(feature = "tokio")]
+    #[error("Tokio AsyncClient requires a current Tokio runtime")]
+    NoTokioRuntime,
     /// The configured server address could not be resolved to a socket address.
     ///
     /// Check the host, port, address-family restrictions, and name service

crates/irtt-client/src/lib.rs

@@ -3,6 +3,8 @@
 //! `irtt-client` opens IRTT-compatible sessions, sends echo probes, receives
 //! replies, classifies loss/late/duplicate packets, and emits [`ClientEvent`]
 //! values for callers to consume directly or aggregate with `irtt-stats`.
+//! With the `tokio` feature enabled, `AsyncClient` provides the same low-level
+//! lifecycle for callers that own a Tokio runtime and drive readiness directly.
 //!
 //! Timing values intentionally preserve signed measurement semantics. When
 //! server timing is available, [`RttSample::effective`] is adjusted for server
@@ -51,6 +53,8 @@
 )]
 #![cfg_attr(all(target_os = "linux", feature = "ancillary"), deny(unsafe_code))]
 
+#[cfg(feature = "tokio")]
+mod async_client;
 mod client;
 mod config;
 mod error;
@@ -64,6 +68,8 @@ mod socket;
 mod socket_options;
 mod timing;
 
+#[cfg(feature = "tokio")]
+pub use async_client::AsyncClient;
 pub use client::Client;
 pub use config::{
     ClientAuthConfig, ClientConfig, NegotiationPolicy, RecvBudget, RunMode, SocketConfig,

crates/irtt-client/src/receive.rs

@@ -62,6 +62,32 @@ pub(crate) fn recv_datagram(
     linux::recv_datagram(socket, buf)
 }
 
+#[cfg(all(
+    feature = "tokio",
+    not(all(target_os = "linux", feature = "ancillary"))
+))]
+pub(crate) fn try_recv_tokio_datagram(
+    socket: &tokio::net::UdpSocket,
+    buf: &mut [u8],
+) -> Result<ReceivedDatagram, io::Error> {
+    let len = socket.try_recv(buf)?;
+    let received_at = ClientTimestamp::now();
+
+    Ok(ReceivedDatagram {
+        len,
+        received_at,
+        meta: ReceiveMeta::default(),
+    })
+}
+
+#[cfg(all(feature = "tokio", target_os = "linux", feature = "ancillary"))]
+pub(crate) fn try_recv_tokio_datagram(
+    socket: &tokio::net::UdpSocket,
+    buf: &mut [u8],
+) -> Result<ReceivedDatagram, io::Error> {
+    linux::try_recv_tokio_datagram(socket, buf)
+}
+
 #[cfg(all(target_os = "linux", feature = "ancillary"))]
 pub(crate) fn recv_datagram_from(
     socket: &UdpSocket,
@@ -133,4 +159,29 @@ mod tests {
         assert!(datagram.received_at.mono >= before.mono);
         assert!(datagram.received_at.mono <= after.mono);
     }
+
+    #[cfg(feature = "tokio")]
+    #[test]
+    fn tokio_try_receive_preserves_fallback_metadata_and_would_block() {
+        let runtime = tokio::runtime::Builder::new_current_thread()
+            .enable_io()
+            .build()
+            .unwrap();
+
+        runtime.block_on(async {
+            let (sender, receiver) = connected_loopback_pair();
+            receiver.set_nonblocking(true).unwrap();
+            let receiver = tokio::net::UdpSocket::from_std(receiver).unwrap();
+            sender.send(b"tokio").unwrap();
+            receiver.readable().await.unwrap();
+
+            let mut buf = [0_u8; 16];
+            let datagram = super::try_recv_tokio_datagram(&receiver, &mut buf).unwrap();
+            assert_eq!(&buf[..datagram.len], b"tokio");
+            assert_eq!(datagram.meta, ReceiveMeta::default());
+
+            let error = super::try_recv_tokio_datagram(&receiver, &mut buf).unwrap_err();
+            assert_eq!(error.kind(), std::io::ErrorKind::WouldBlock);
+        });
+    }
 }

crates/irtt-client/src/receive/linux.rs

@@ -2,7 +2,7 @@
 use std::{
     io, mem,
     net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6, UdpSocket},
-    os::fd::AsRawFd,
+    os::fd::{AsRawFd, RawFd},
     ptr,
     time::{Duration, SystemTime, UNIX_EPOCH},
 };
@@ -49,6 +49,24 @@ fn enable_kernel_rx_timestamps(socket: &UdpSocket) -> io::Result<()> {
 pub(crate) fn recv_datagram(
     socket: &UdpSocket,
     buf: &mut [u8],
+) -> Result<ReceivedDatagram, io::Error> {
+    recv_datagram_fd(socket.as_raw_fd(), buf, 0)
+}
+
+#[cfg(feature = "tokio")]
+pub(crate) fn try_recv_tokio_datagram(
+    socket: &tokio::net::UdpSocket,
+    buf: &mut [u8],
+) -> Result<ReceivedDatagram, io::Error> {
+    socket.try_io(tokio::io::Interest::READABLE, || {
+        recv_datagram_fd(socket.as_raw_fd(), buf, libc::MSG_DONTWAIT)
+    })
+}
+
+fn recv_datagram_fd(
+    socket_fd: RawFd,
+    buf: &mut [u8],
+    flags: libc::c_int,
 ) -> Result<ReceivedDatagram, io::Error> {
     let mut control = ControlBuffer::new();
     let mut iov = libc::iovec {
@@ -64,8 +82,8 @@ pub(crate) fn recv_datagram(
     let len = unsafe {
         // SAFETY: `msg` points to one writable iovec backed by `buf`, and the
         // control buffer is writable and lives until `recvmsg` returns. The
-        // socket file descriptor is borrowed from a valid `UdpSocket`.
+        // socket file descriptor is borrowed from a valid UDP socket.
-        libc::recvmsg(socket.as_raw_fd(), &mut msg, 0)
+        libc::recvmsg(socket_fd, &mut msg, flags)
     };
     if len < 0 {
         return Err(io::Error::last_os_error());
@@ -366,6 +384,52 @@ mod tests {
         assert!(datagram.received_at.mono <= after.mono);
     }
 
+    #[cfg(feature = "tokio")]
+    #[test]
+    fn tokio_recvmsg_would_block_clears_readiness() {
+        let runtime = tokio::runtime::Builder::new_current_thread()
+            .enable_io()
+            .enable_time()
+            .build()
+            .unwrap();
+
+        runtime.block_on(async {
+            let (sender, receiver) = connected_ipv4_loopback_pair();
+            configure_receive_metadata(&receiver, sender.local_addr().unwrap()).unwrap();
+            receiver.set_nonblocking(true).unwrap();
+            let receiver = tokio::net::UdpSocket::from_std(receiver).unwrap();
+            let mut buf = [0_u8; 16];
+
+            sender.send(b"first").unwrap();
+            tokio::time::timeout(Duration::from_secs(1), receiver.readable())
+                .await
+                .expect("initial readability timed out")
+                .unwrap();
+
+            let first = super::try_recv_tokio_datagram(&receiver, &mut buf).unwrap();
+            assert_eq!(&buf[..first.len], b"first");
+
+            let error = super::try_recv_tokio_datagram(&receiver, &mut buf).unwrap_err();
+            assert_eq!(error.kind(), io::ErrorKind::WouldBlock);
+
+            assert!(
+                tokio::time::timeout(Duration::from_millis(20), receiver.readable())
+                    .await
+                    .is_err(),
+                "readable completed from stale readiness after recvmsg WouldBlock"
+            );
+
+            sender.send(b"second").unwrap();
+            tokio::time::timeout(Duration::from_secs(1), receiver.readable())
+                .await
+                .expect("second readability timed out")
+                .unwrap();
+
+            let second = super::try_recv_tokio_datagram(&receiver, &mut buf).unwrap();
+            assert_eq!(&buf[..second.len], b"second");
+        });
+    }
+
     #[test]
     fn kernel_rx_timestamp_metadata_is_observed_when_kernel_provides_it() {
         let (sender, receiver) = connected_ipv4_loopback_pair();

crates/irtt-client/src/socket.rs

@@ -28,18 +28,57 @@ pub(crate) fn validate_open_timeouts(timeouts: &[Duration]) -> Result<(), Client
 }
 
 pub(crate) fn resolve_remote(config: &ClientConfig) -> Result<SocketAddr, ClientError> {
+    #[cfg(all(test, feature = "tokio"))]
+    SYNC_RESOLVER_CALLS.with(|calls| calls.set(calls.get() + 1));
+
     let addr = normalize_server_addr(&config.server_addr);
     let mut addrs = addr
         .to_socket_addrs()
         .map_err(|_| ClientError::Resolve { addr: addr.clone() })?;
     addrs
-        .find(|addr| {
+        .find(|addr| address_family_allowed(config, *addr))
-            (!config.socket_config.ipv4_only || addr.is_ipv4())
-                && (!config.socket_config.ipv6_only || addr.is_ipv6())
-        })
         .ok_or(ClientError::Resolve { addr })
 }
 
+#[cfg(feature = "tokio")]
+pub(crate) async fn resolve_remote_tokio(config: &ClientConfig) -> Result<SocketAddr, ClientError> {
+    let addr = normalize_server_addr(&config.server_addr);
+    if let Ok(remote) = addr.parse::<SocketAddr>() {
+        return address_family_allowed(config, remote)
+            .then_some(remote)
+            .ok_or(ClientError::Resolve { addr });
+    }
+
+    #[cfg(all(test, feature = "tokio"))]
+    TOKIO_DNS_LOOKUPS.with(|calls| calls.set(calls.get() + 1));
+
+    let mut addrs = tokio::net::lookup_host(&addr)
+        .await
+        .map_err(|_| ClientError::Resolve { addr: addr.clone() })?;
+    addrs
+        .find(|remote| address_family_allowed(config, *remote))
+        .ok_or_else(|| ClientError::Resolve { addr: addr.clone() })
+}
+
+fn address_family_allowed(config: &ClientConfig, remote: SocketAddr) -> bool {
+    (!config.socket_config.ipv4_only || remote.is_ipv4())
+        && (!config.socket_config.ipv6_only || remote.is_ipv6())
+}
+
+#[cfg(all(test, feature = "tokio"))]
+std::thread_local! {
+    static SYNC_RESOLVER_CALLS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
+    static TOKIO_DNS_LOOKUPS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
+}
+
+#[cfg(all(test, feature = "tokio"))]
+pub(crate) fn resolution_call_counts() -> (usize, usize) {
+    (
+        SYNC_RESOLVER_CALLS.with(std::cell::Cell::get),
+        TOKIO_DNS_LOOKUPS.with(std::cell::Cell::get),
+    )
+}
+
 pub(crate) fn normalize_server_addr(addr: &str) -> String {
     if addr.parse::<SocketAddr>().is_ok() {
         return addr.to_owned();
@@ -66,16 +105,16 @@ pub(crate) fn connect_udp_socket(
     config: &SocketConfig,
     remote: SocketAddr,
 ) -> Result<UdpSocket, ClientError> {
-    let domain = if remote.is_ipv4() {
+    let socket = create_connected_udp_socket(config, remote)?;
-        Domain::IPV4
+    socket.set_read_timeout(config.recv_timeout)?;
-    } else {
+    Ok(socket)
-        Domain::IPV6
-    };
-    let socket = Socket::new(domain, Type::DGRAM, Some(Protocol::UDP))?;
-
-    if config.ipv6_only && remote.is_ipv6() {
-        socket.set_only_v6(true)?;
 }
+
+fn create_connected_udp_socket(
+    config: &SocketConfig,
+    remote: SocketAddr,
+) -> Result<UdpSocket, ClientError> {
+    let socket = create_prebind_udp_socket(config, remote)?;
     let bind_addr = config.bind_addr.unwrap_or_else(|| {
         if remote.is_ipv4() {
             SocketAddr::from(([0, 0, 0, 0], 0))
@@ -95,10 +134,36 @@ pub(crate) fn connect_udp_socket(
     if let Some(ttl) = config.ttl {
         apply_ttl_to_socket(&socket, remote, ttl)?;
     }
-    socket.set_read_timeout(config.recv_timeout)?;
     Ok(socket)
 }
 
+fn create_prebind_udp_socket(
+    config: &SocketConfig,
+    remote: SocketAddr,
+) -> Result<Socket, ClientError> {
+    let domain = if remote.is_ipv4() {
+        Domain::IPV4
+    } else {
+        Domain::IPV6
+    };
+    let socket = Socket::new(domain, Type::DGRAM, Some(Protocol::UDP))?;
+
+    if config.ipv6_only && remote.is_ipv6() {
+        socket.set_only_v6(true)?;
+    }
+    Ok(socket)
+}
+
+#[cfg(feature = "tokio")]
+pub(crate) fn connect_tokio_udp_socket(
+    config: &SocketConfig,
+    remote: SocketAddr,
+) -> Result<tokio::net::UdpSocket, ClientError> {
+    let socket = create_connected_udp_socket(config, remote)?;
+    socket.set_nonblocking(true)?;
+    Ok(tokio::net::UdpSocket::from_std(socket)?)
+}
+
 pub(crate) fn bind_unconnected_udp_socket(
     config: &SocketConfig,
     family_remote: SocketAddr,
@@ -151,4 +216,192 @@ mod tests {
         assert_eq!(normalize_server_addr("[::1]"), "[::1]:2112");
         assert_eq!(normalize_server_addr("[::1]:1234"), "[::1]:1234");
     }
+
+    #[test]
+    fn address_family_filtering_remains_unchanged() {
+        let ipv4 = SocketAddr::from(([127, 0, 0, 1], DEFAULT_PORT));
+        let ipv6 = SocketAddr::from(([0, 0, 0, 0, 0, 0, 0, 1], DEFAULT_PORT));
+        let mut config = ClientConfig::default();
+
+        assert!(address_family_allowed(&config, ipv4));
+        assert!(address_family_allowed(&config, ipv6));
+
+        config.socket_config.ipv4_only = true;
+        assert!(address_family_allowed(&config, ipv4));
+        assert!(!address_family_allowed(&config, ipv6));
+
+        config.socket_config.ipv4_only = false;
+        config.socket_config.ipv6_only = true;
+        assert!(!address_family_allowed(&config, ipv4));
+        assert!(address_family_allowed(&config, ipv6));
+    }
+
+    #[cfg(feature = "tokio")]
+    #[test]
+    fn tokio_ipv6_only_is_applied_before_bind() {
+        let config = SocketConfig {
+            ipv6_only: true,
+            ..SocketConfig::default()
+        };
+        let remote = SocketAddr::from(([0, 0, 0, 0, 0, 0, 0, 1], DEFAULT_PORT));
+
+        let socket = create_prebind_udp_socket(&config, remote).unwrap();
+
+        assert!(socket.only_v6().unwrap());
+    }
+
+    #[cfg(feature = "tokio")]
+    #[test]
+    fn tokio_ipv6_socket_preserves_v6only_and_explicit_bind() {
+        let bind_addr = SocketAddr::from(([0, 0, 0, 0, 0, 0, 0, 1], 0));
+        let peer = match UdpSocket::bind(bind_addr) {
+            Ok(peer) => peer,
+            Err(error) if ipv6_loopback_unavailable(&error) => {
+                eprintln!("skipping IPv6 Tokio socket test: IPv6 loopback unavailable: {error}");
+                return;
+            }
+            Err(error) => panic!("{error}"),
+        };
+        let remote = peer.local_addr().unwrap();
+        let config = SocketConfig {
+            bind_addr: Some(bind_addr),
+            ipv6_only: true,
+            ..SocketConfig::default()
+        };
+
+        tokio_runtime().block_on(async {
+            let socket = connect_tokio_udp_socket(&config, remote).unwrap();
+
+            assert!(socket2::SockRef::from(&socket).only_v6().unwrap());
+            assert!(socket.local_addr().unwrap().is_ipv6());
+            assert_eq!(socket.local_addr().unwrap().ip(), bind_addr.ip());
+            assert_eq!(socket.peer_addr().unwrap(), remote);
+        });
+    }
+
+    #[cfg(feature = "tokio")]
+    #[test]
+    fn tokio_ipv4_socket_preserves_family_bind_peer_and_ttl() {
+        let peer = UdpSocket::bind(SocketAddr::from(([127, 0, 0, 1], 0))).unwrap();
+        let remote = peer.local_addr().unwrap();
+        let bind_addr = SocketAddr::from(([127, 0, 0, 1], 0));
+        let config = SocketConfig {
+            bind_addr: Some(bind_addr),
+            ttl: Some(64),
+            ..SocketConfig::default()
+        };
+
+        tokio_runtime().block_on(async {
+            let socket = connect_tokio_udp_socket(&config, remote).unwrap();
+
+            assert!(socket.local_addr().unwrap().is_ipv4());
+            assert_eq!(socket.local_addr().unwrap().ip(), bind_addr.ip());
+            assert_eq!(socket.peer_addr().unwrap(), remote);
+            assert_eq!(socket2::SockRef::from(&socket).ttl_v4().unwrap(), 64);
+        });
+    }
+
+    #[cfg(feature = "tokio")]
+    #[test]
+    fn tokio_resolver_filters_families_without_using_blocking_resolution() {
+        let mut config = ClientConfig {
+            server_addr: "localhost:2112".to_owned(),
+            ..ClientConfig::default()
+        };
+        config.socket_config.ipv4_only = true;
+
+        tokio_runtime().block_on(async {
+            let before = resolution_call_counts();
+            let remote = resolve_remote_tokio(&config).await.unwrap();
+            let after = resolution_call_counts();
+
+            assert!(remote.is_ipv4());
+            assert_eq!(after.0, before.0);
+            assert_eq!(after.1, before.1 + 1);
+        });
+    }
+
+    #[cfg(feature = "tokio")]
+    #[test]
+    fn tokio_resolver_selects_ipv6_for_ipv6_only_hostname() {
+        let mut config = ClientConfig {
+            server_addr: "localhost:2112".to_owned(),
+            ..ClientConfig::default()
+        };
+        config.socket_config.ipv6_only = true;
+
+        tokio_runtime().block_on(async {
+            let remote = resolve_remote_tokio(&config).await.unwrap();
+            assert!(remote.is_ipv6());
+        });
+    }
+
+    #[cfg(feature = "tokio")]
+    #[test]
+    fn tokio_resolver_literal_path_bypasses_all_name_resolution() {
+        let config = ClientConfig {
+            server_addr: "127.0.0.1:2112".to_owned(),
+            ..ClientConfig::default()
+        };
+
+        tokio_runtime().block_on(async {
+            let before = resolution_call_counts();
+            let remote = resolve_remote_tokio(&config).await.unwrap();
+            let after = resolution_call_counts();
+
+            assert_eq!(remote, SocketAddr::from(([127, 0, 0, 1], 2112)));
+            assert_eq!(after, before);
+        });
+    }
+
+    #[cfg(feature = "tokio")]
+    #[test]
+    fn tokio_resolver_rejects_disallowed_literal_family() {
+        let mut config = ClientConfig {
+            server_addr: "127.0.0.1:2112".to_owned(),
+            ..ClientConfig::default()
+        };
+        config.socket_config.ipv6_only = true;
+
+        tokio_runtime().block_on(async {
+            assert!(matches!(
+                resolve_remote_tokio(&config).await,
+                Err(ClientError::Resolve { .. })
+            ));
+        });
+    }
+
+    #[cfg(feature = "tokio")]
+    #[test]
+    fn tokio_resolver_rejects_mixed_family_restrictions() {
+        let mut config = ClientConfig {
+            server_addr: "127.0.0.1:2112".to_owned(),
+            ..ClientConfig::default()
+        };
+        config.socket_config.ipv4_only = true;
+        config.socket_config.ipv6_only = true;
+
+        tokio_runtime().block_on(async {
+            assert!(matches!(
+                resolve_remote_tokio(&config).await,
+                Err(ClientError::Resolve { .. })
+            ));
+        });
+    }
+
+    #[cfg(feature = "tokio")]
+    fn tokio_runtime() -> tokio::runtime::Runtime {
+        tokio::runtime::Builder::new_current_thread()
+            .enable_io()
+            .build()
+            .unwrap()
+    }
+
+    #[cfg(feature = "tokio")]
+    fn ipv6_loopback_unavailable(error: &std::io::Error) -> bool {
+        matches!(
+            error.kind(),
+            std::io::ErrorKind::AddrNotAvailable | std::io::ErrorKind::Unsupported
+        )
+    }
 }

crates/irtt-client/src/socket_options.rs

@@ -3,6 +3,8 @@ use std::{
     net::{SocketAddr, UdpSocket},
 };
 
+use socket2::SockRef;
+
 use crate::{
     config::{MAX_DSCP_CODEPOINT, MAX_TTL},
     error::ClientError,
@@ -33,6 +35,24 @@ pub(crate) fn clear_dscp_on_socket(
     set_socket_traffic_class(socket, remote, 0, "clear DSCP before close")
 }
 
+#[cfg(feature = "tokio")]
+pub(crate) fn apply_dscp_to_tokio_socket(
+    socket: &tokio::net::UdpSocket,
+    remote: SocketAddr,
+    dscp_codepoint: u8,
+) -> Result<(), ClientError> {
+    let traffic_class = dscp_codepoint_to_traffic_class(dscp_codepoint)?;
+    set_tokio_socket_traffic_class(socket, remote, traffic_class, "set negotiated DSCP")
+}
+
+#[cfg(feature = "tokio")]
+pub(crate) fn clear_dscp_on_tokio_socket(
+    socket: &tokio::net::UdpSocket,
+    remote: SocketAddr,
+) -> Result<(), ClientError> {
+    set_tokio_socket_traffic_class(socket, remote, 0, "clear DSCP before close")
+}
+
 pub(crate) fn validate_ttl(ttl: u32) -> Result<(), ClientError> {
     if ttl == 0 || ttl > MAX_TTL {
         return Err(ClientError::InvalidConfig {
@@ -67,6 +87,20 @@ pub(crate) fn socket_traffic_class(
     })
 }
 
+#[cfg(all(test, feature = "tokio"))]
+pub(crate) fn tokio_socket_traffic_class(
+    socket: &tokio::net::UdpSocket,
+    remote: SocketAddr,
+) -> Result<u32, ClientError> {
+    get_socket_traffic_class_ref(SockRef::from(socket), remote).map_err(|source| {
+        ClientError::SocketOption {
+            operation: "read DSCP socket option",
+            remote,
+            source,
+        }
+    })
+}
+
 #[cfg(test)]
 pub(crate) fn socket_ttl(socket: &UdpSocket, remote: SocketAddr) -> Result<u32, ClientError> {
     get_socket_ttl(socket, remote).map_err(|source| ClientError::SocketOption {
@@ -77,19 +111,24 @@ pub(crate) fn socket_ttl(socket: &UdpSocket, remote: SocketAddr) -> Result<u32,
 }
 
 fn set_socket_ttl(socket: &UdpSocket, remote: SocketAddr, ttl: u32) -> io::Result<()> {
+    set_socket_ttl_ref(SockRef::from(socket), remote, ttl)
+}
+
+fn set_socket_ttl_ref(socket: SockRef<'_>, remote: SocketAddr, ttl: u32) -> io::Result<()> {
     if remote.is_ipv4() {
-        socket2::SockRef::from(socket).set_ttl_v4(ttl)
+        socket.set_ttl_v4(ttl)
     } else {
-        socket2::SockRef::from(socket).set_unicast_hops_v6(ttl)
+        socket.set_unicast_hops_v6(ttl)
     }
 }
 
 #[cfg(test)]
 fn get_socket_ttl(socket: &UdpSocket, remote: SocketAddr) -> io::Result<u32> {
+    let socket = SockRef::from(socket);
     if remote.is_ipv4() {
-        socket2::SockRef::from(socket).ttl_v4()
+        socket.ttl_v4()
     } else {
-        socket2::SockRef::from(socket).unicast_hops_v6()
+        socket.unicast_hops_v6()
     }
 }
 
@@ -112,6 +151,30 @@ fn set_socket_traffic_class_io(
     socket: &UdpSocket,
     remote: SocketAddr,
     traffic_class: u32,
+) -> io::Result<()> {
+    set_socket_traffic_class_ref(SockRef::from(socket), remote, traffic_class)
+}
+
+#[cfg(feature = "tokio")]
+fn set_tokio_socket_traffic_class(
+    socket: &tokio::net::UdpSocket,
+    remote: SocketAddr,
+    traffic_class: u32,
+    operation: &'static str,
+) -> Result<(), ClientError> {
+    set_socket_traffic_class_ref(SockRef::from(socket), remote, traffic_class).map_err(|source| {
+        ClientError::SocketOption {
+            operation,
+            remote,
+            source,
+        }
+    })
+}
+
+fn set_socket_traffic_class_ref(
+    socket: SockRef<'_>,
+    remote: SocketAddr,
+    traffic_class: u32,
 ) -> io::Result<()> {
     if remote.is_ipv4() {
         set_ipv4_traffic_class(socket, traffic_class)
@@ -127,8 +190,8 @@ fn set_socket_traffic_class_io(
     target_os = "illumos",
     target_os = "haiku",
 )))]
-fn set_ipv4_traffic_class(socket: &UdpSocket, traffic_class: u32) -> io::Result<()> {
+fn set_ipv4_traffic_class(socket: SockRef<'_>, traffic_class: u32) -> io::Result<()> {
-    socket2::SockRef::from(socket).set_tos_v4(traffic_class)
+    socket.set_tos_v4(traffic_class)
 }
 
 #[cfg(any(
@@ -138,7 +201,7 @@ fn set_ipv4_traffic_class(socket: &UdpSocket, traffic_class: u32) -> io::Result<
     target_os = "illumos",
     target_os = "haiku",
 ))]
-fn set_ipv4_traffic_class(_socket: &UdpSocket, traffic_class: u32) -> io::Result<()> {
+fn set_ipv4_traffic_class(_socket: SockRef<'_>, traffic_class: u32) -> io::Result<()> {
     unsupported_traffic_class(traffic_class, "IPv4 DSCP socket options")
 }
 
@@ -153,8 +216,8 @@ fn set_ipv4_traffic_class(_socket: &UdpSocket, traffic_class: u32) -> io::Result
     target_os = "openbsd",
     target_os = "cygwin",
 ))]
-fn set_ipv6_traffic_class(socket: &UdpSocket, traffic_class: u32) -> io::Result<()> {
+fn set_ipv6_traffic_class(socket: SockRef<'_>, traffic_class: u32) -> io::Result<()> {
-    socket2::SockRef::from(socket).set_tclass_v6(traffic_class)
+    socket.set_tclass_v6(traffic_class)
 }
 
 #[cfg(not(any(
@@ -168,7 +231,7 @@ fn set_ipv6_traffic_class(socket: &UdpSocket, traffic_class: u32) -> io::Result<
     target_os = "openbsd",
     target_os = "cygwin",
 )))]
-fn set_ipv6_traffic_class(_socket: &UdpSocket, traffic_class: u32) -> io::Result<()> {
+fn set_ipv6_traffic_class(_socket: SockRef<'_>, traffic_class: u32) -> io::Result<()> {
     unsupported_traffic_class(traffic_class, "IPv6 DSCP socket options")
 }
 
@@ -211,8 +274,8 @@ fn unsupported_traffic_class(traffic_class: u32, feature: &'static str) -> io::R
         target_os = "haiku",
     ))
 ))]
-fn get_ipv4_traffic_class(socket: &UdpSocket) -> io::Result<u32> {
+fn get_ipv4_traffic_class(socket: SockRef<'_>) -> io::Result<u32> {
-    socket2::SockRef::from(socket).tos_v4()
+    socket.tos_v4()
 }
 
 #[cfg(all(
@@ -225,7 +288,7 @@ fn get_ipv4_traffic_class(socket: &UdpSocket) -> io::Result<u32> {
         target_os = "haiku",
     )
 ))]
-fn get_ipv4_traffic_class(_socket: &UdpSocket) -> io::Result<u32> {
+fn get_ipv4_traffic_class(_socket: SockRef<'_>) -> io::Result<u32> {
     unsupported_readback("IPv4 DSCP socket option readback")
 }
 
@@ -243,8 +306,8 @@ fn get_ipv4_traffic_class(_socket: &UdpSocket) -> io::Result<u32> {
         target_os = "cygwin",
     )
 ))]
-fn get_ipv6_traffic_class(socket: &UdpSocket) -> io::Result<u32> {
+fn get_ipv6_traffic_class(socket: SockRef<'_>) -> io::Result<u32> {
-    socket2::SockRef::from(socket).tclass_v6()
+    socket.tclass_v6()
 }
 
 #[cfg(all(
@@ -261,12 +324,17 @@ fn get_ipv6_traffic_class(socket: &UdpSocket) -> io::Result<u32> {
         target_os = "cygwin",
     ))
 ))]
-fn get_ipv6_traffic_class(_socket: &UdpSocket) -> io::Result<u32> {
+fn get_ipv6_traffic_class(_socket: SockRef<'_>) -> io::Result<u32> {
     unsupported_readback("IPv6 DSCP socket option readback")
 }
 
 #[cfg(test)]
 fn get_socket_traffic_class(socket: &UdpSocket, remote: SocketAddr) -> io::Result<u32> {
+    get_socket_traffic_class_ref(SockRef::from(socket), remote)
+}
+
+#[cfg(test)]
+fn get_socket_traffic_class_ref(socket: SockRef<'_>, remote: SocketAddr) -> io::Result<u32> {
     if remote.is_ipv4() {
         get_ipv4_traffic_class(socket)
     } else {
@@ -395,6 +463,40 @@ mod tests {
         assert_eq!(socket_traffic_class(&socket, remote).unwrap(), 0);
     }
 
+    #[cfg(all(
+        feature = "tokio",
+        not(any(
+            target_os = "fuchsia",
+            target_os = "redox",
+            target_os = "solaris",
+            target_os = "illumos",
+            target_os = "haiku",
+        ))
+    ))]
+    #[test]
+    fn tokio_ipv4_socket_option_sets_and_clears_traffic_class() {
+        let runtime = tokio::runtime::Builder::new_current_thread()
+            .enable_io()
+            .build()
+            .unwrap();
+
+        runtime.block_on(async {
+            let remote = SocketAddr::from(([127, 0, 0, 1], 9));
+            let socket = UdpSocket::bind(SocketAddr::from(([127, 0, 0, 1], 0))).unwrap();
+            socket.connect(remote).unwrap();
+            socket.set_nonblocking(true).unwrap();
+            let socket = tokio::net::UdpSocket::from_std(socket).unwrap();
+
+            apply_dscp_to_tokio_socket(&socket, remote, 46).unwrap();
+            let traffic_class = tokio_socket_traffic_class(&socket, remote).unwrap();
+            assert_eq!(traffic_class & 0xfc, 184);
+            assert_eq!(traffic_class & 0b11, 0);
+
+            clear_dscp_on_tokio_socket(&socket, remote).unwrap();
+            assert_eq!(tokio_socket_traffic_class(&socket, remote).unwrap(), 0);
+        });
+    }
+
     #[test]
     #[cfg(any(
         target_os = "android",