From 97d61371cb2a6768ad178d4abc1bd5f708cdec0a Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Thu, 7 May 2026 17:35:54 +0700 Subject: [PATCH] feat(client): add receive datagram abstraction --- crates/irtt-client/src/client.rs | 18 +++-- .../src/client/tests/probes_replies.rs | 47 +++++++++++- crates/irtt-client/src/lib.rs | 1 + crates/irtt-client/src/receive.rs | 76 +++++++++++++++++++ 4 files changed, 133 insertions(+), 9 deletions(-) create mode 100644 crates/irtt-client/src/receive.rs diff --git a/crates/irtt-client/src/client.rs b/crates/irtt-client/src/client.rs index 23d7ec7..438141b 100644 --- a/crates/irtt-client/src/client.rs +++ b/crates/irtt-client/src/client.rs @@ -19,6 +19,7 @@ use crate::{ }, metadata::ReceiveMeta, probe::{CompletedSet, PendingMap, PendingProbe, TimedOutMap}, + receive::recv_datagram, session::{validate_negotiated_params, ActiveSession, ClientPhase, NegotiatedParams}, socket::{connect_udp_socket, resolve_remote, validate_open_timeouts}, socket_options::{apply_dscp_to_socket, clear_dscp_on_socket}, @@ -278,8 +279,8 @@ impl Client { let buf_size = self.recv_buffer_size(); let mut buf = vec![0_u8; buf_size]; - let size = match self.socket.recv(&mut buf) { - Ok(size) => size, + let datagram = match recv_datagram(&self.socket, &mut buf) { + Ok(datagram) => datagram, Err(err) if matches!( err.kind(), @@ -291,8 +292,8 @@ impl Client { Err(err) => return Err(ClientError::Socket(err)), }; - let now = override_ts.unwrap_or_else(ClientTimestamp::now); - self.process_received_packet(&buf[..size], now) + let now = override_ts.unwrap_or(datagram.received_at); + self.process_received_packet(&buf[..datagram.len], now, datagram.meta) } pub fn recv_available(&mut self, budget: RecvBudget) -> Result, ClientError> { @@ -372,6 +373,7 @@ impl Client { &mut self, packet: &[u8], now: ClientTimestamp, + meta: ReceiveMeta, ) -> Result, ClientError> { let negotiated = self .negotiated @@ -432,7 +434,7 @@ impl Client { one_way, received_stats, bytes: packet.len(), - packet_meta: ReceiveMeta::default().into(), + packet_meta: meta.into(), }); } else { events.push(ClientEvent::EchoReply { @@ -446,7 +448,7 @@ impl Client { one_way, received_stats, bytes: packet.len(), - packet_meta: ReceiveMeta::default().into(), + packet_meta: meta.into(), }); } Ok(events) @@ -479,7 +481,7 @@ impl Client { one_way, received_stats, bytes: packet.len(), - packet_meta: ReceiveMeta::default().into(), + packet_meta: meta.into(), }]) } else if session.highest_received_seq.is_some_and(|h| wire_seq < h) { // TODO: u32 ordering is acceptable for finite tests but continuous @@ -496,7 +498,7 @@ impl Client { one_way: None, received_stats: build_received_stats(&reply), bytes: packet.len(), - packet_meta: ReceiveMeta::default().into(), + packet_meta: meta.into(), }]) } else { Ok(vec![ClientEvent::Warning { diff --git a/crates/irtt-client/src/client/tests/probes_replies.rs b/crates/irtt-client/src/client/tests/probes_replies.rs index ea1b7af..c89c31f 100644 --- a/crates/irtt-client/src/client/tests/probes_replies.rs +++ b/crates/irtt-client/src/client/tests/probes_replies.rs @@ -1057,7 +1057,9 @@ fn matched_reply_with_reversed_monotonic_time_still_emits_event() { wall: send_ts.wall + Duration::from_millis(10), }; let reply = echo_reply_packet(TOKEN, 0, ¶ms, &TimestampFields::default(), None); - let events = client.process_received_packet(&reply, recv_ts).unwrap(); + let events = client + .process_received_packet(&reply, recv_ts, ReceiveMeta::default()) + .unwrap(); assert_eq!(events.len(), 1); match &events[0] { @@ -1072,6 +1074,49 @@ fn matched_reply_with_reversed_monotonic_time_still_emits_event() { server.join(); } +#[test] +fn process_received_packet_uses_supplied_receive_metadata() { + let params = default_params(); + let server = silent_open_server(params.clone()); + let config = ClientConfig { + socket_config: crate::SocketConfig { + recv_timeout: Some(Duration::from_millis(50)), + ..Default::default() + }, + ..default_test_config(server.addr) + }; + let mut client = Client::connect(config).unwrap(); + assert_open_started(client.open(ClientTimestamp::now()).unwrap()); + + client.send_probe_at(ClientTimestamp::now()).unwrap(); + let recv_ts = ClientTimestamp::now(); + let reply = echo_reply_packet(TOKEN, 0, ¶ms, &TimestampFields::default(), None); + let events = client + .process_received_packet( + &reply, + recv_ts, + ReceiveMeta { + traffic_class: Some(186), + kernel_rx_timestamp: None, + }, + ) + .unwrap(); + + assert_eq!(events.len(), 1); + match &events[0] { + ClientEvent::EchoReply { packet_meta, .. } => { + assert_eq!(packet_meta.traffic_class, Some(186)); + assert_eq!(packet_meta.dscp, Some(46)); + assert_eq!(packet_meta.ecn, Some(2)); + assert_eq!(packet_meta.kernel_rx_timestamp, None); + } + other => panic!("expected EchoReply, got {other:?}"), + } + + client.close(ClientTimestamp::now()).unwrap(); + server.join(); +} + #[test] fn recv_available_drains_burst_replies() { let params = default_params(); diff --git a/crates/irtt-client/src/lib.rs b/crates/irtt-client/src/lib.rs index 98e6183..06a36d6 100644 --- a/crates/irtt-client/src/lib.rs +++ b/crates/irtt-client/src/lib.rs @@ -7,6 +7,7 @@ mod event; mod managed; mod metadata; mod probe; +mod receive; mod session; mod socket; mod socket_options; diff --git a/crates/irtt-client/src/receive.rs b/crates/irtt-client/src/receive.rs new file mode 100644 index 0000000..fbfd39d --- /dev/null +++ b/crates/irtt-client/src/receive.rs @@ -0,0 +1,76 @@ +use std::{io, net::UdpSocket}; + +use crate::{metadata::ReceiveMeta, timing::ClientTimestamp}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct ReceivedDatagram { + pub(crate) len: usize, + pub(crate) received_at: ClientTimestamp, + pub(crate) meta: ReceiveMeta, +} + +pub(crate) fn recv_datagram( + socket: &UdpSocket, + buf: &mut [u8], +) -> Result { + let len = socket.recv(buf)?; + let received_at = ClientTimestamp::now(); + + Ok(ReceivedDatagram { + len, + received_at, + meta: ReceiveMeta::default(), + }) +} + +#[cfg(test)] +mod tests { + use std::net::UdpSocket; + + use crate::{metadata::ReceiveMeta, receive::recv_datagram, timing::ClientTimestamp}; + + fn connected_loopback_pair() -> (UdpSocket, UdpSocket) { + let a = UdpSocket::bind("127.0.0.1:0").unwrap(); + let b = UdpSocket::bind("127.0.0.1:0").unwrap(); + a.connect(b.local_addr().unwrap()).unwrap(); + b.connect(a.local_addr().unwrap()).unwrap(); + (a, b) + } + + #[test] + fn fallback_receive_returns_length() { + let (sender, receiver) = connected_loopback_pair(); + sender.send(b"hello").unwrap(); + + let mut buf = [0_u8; 16]; + let datagram = recv_datagram(&receiver, &mut buf).unwrap(); + + assert_eq!(datagram.len, 5); + assert_eq!(&buf[..datagram.len], b"hello"); + } + + #[test] + fn fallback_receive_returns_default_metadata() { + let (sender, receiver) = connected_loopback_pair(); + sender.send(b"meta").unwrap(); + + let mut buf = [0_u8; 16]; + let datagram = recv_datagram(&receiver, &mut buf).unwrap(); + + assert_eq!(datagram.meta, ReceiveMeta::default()); + } + + #[test] + fn fallback_receive_captures_timestamp_after_successful_receive() { + let (sender, receiver) = connected_loopback_pair(); + sender.send(b"time").unwrap(); + + let before = ClientTimestamp::now(); + let mut buf = [0_u8; 16]; + let datagram = recv_datagram(&receiver, &mut buf).unwrap(); + let after = ClientTimestamp::now(); + + assert!(datagram.received_at.mono >= before.mono); + assert!(datagram.received_at.mono <= after.mono); + } +}