feat(client): add optional Tokio socket support

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();