feat(client): add optional Tokio socket support

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/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`.
-        libc::recvmsg(socket.as_raw_fd(), &mut msg, 0)
+        // socket file descriptor is borrowed from a valid UDP socket.
+        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| {
-            (!config.socket_config.ipv4_only || addr.is_ipv4())
-                && (!config.socket_config.ipv6_only || addr.is_ipv6())
-        })
+        .find(|addr| address_family_allowed(config, *addr))
         .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() {
-        Domain::IPV4
-    } else {
-        Domain::IPV6
-    };
-    let socket = Socket::new(domain, Type::DGRAM, Some(Protocol::UDP))?;
+    let socket = create_connected_udp_socket(config, remote)?;
+    socket.set_read_timeout(config.recv_timeout)?;
+    Ok(socket)
+}
 
-    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<()> {
-    socket2::SockRef::from(socket).set_tos_v4(traffic_class)
+fn set_ipv4_traffic_class(socket: SockRef<'_>, traffic_class: u32) -> io::Result<()> {
+    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<()> {
-    socket2::SockRef::from(socket).set_tclass_v6(traffic_class)
+fn set_ipv6_traffic_class(socket: SockRef<'_>, traffic_class: u32) -> io::Result<()> {
+    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> {
-    socket2::SockRef::from(socket).tos_v4()
+fn get_ipv4_traffic_class(socket: SockRef<'_>) -> io::Result<u32> {
+    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> {
-    socket2::SockRef::from(socket).tclass_v6()
+fn get_ipv6_traffic_class(socket: SockRef<'_>) -> io::Result<u32> {
+    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",