From fedb478ad60c0496298b6c4e9e4cdbf49c195c04 Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Thu, 7 May 2026 23:40:12 +0700 Subject: [PATCH] perf(client): reuse receive buffer --- crates/irtt-client/src/client.rs | 75 ++++++++++++------- .../src/client/tests/probes_replies.rs | 68 +++++++++++++++++ 2 files changed, 116 insertions(+), 27 deletions(-) diff --git a/crates/irtt-client/src/client.rs b/crates/irtt-client/src/client.rs index 4decdf6..b289f14 100644 --- a/crates/irtt-client/src/client.rs +++ b/crates/irtt-client/src/client.rs @@ -41,6 +41,7 @@ pub struct Client { negotiated: Option, phase: ClientPhase, session: Option, + recv_buffer: Vec, } impl Client { @@ -68,6 +69,7 @@ impl Client { negotiated: None, phase: ClientPhase::Connected, session: None, + recv_buffer: vec![0_u8; MIN_RECV_BUFFER_SIZE], }) } @@ -289,9 +291,7 @@ impl Client { ClientPhase::NoTestCompleted => return Err(ClientError::AlreadyCompleted), } - let buf_size = self.recv_buffer_size(); - let mut buf = vec![0_u8; buf_size]; - let datagram = match recv_datagram(&self.socket, &mut buf) { + let datagram = match recv_datagram(&self.socket, &mut self.recv_buffer) { Ok(datagram) => datagram, Err(err) if matches!( @@ -305,7 +305,14 @@ impl Client { }; let now = override_ts.unwrap_or(datagram.received_at); - self.process_received_packet(&buf[..datagram.len], now, datagram.meta) + let packet_len = datagram.len; + let Some(reply) = self.decode_received_packet(&self.recv_buffer[..packet_len]) else { + return Ok(vec![ClientEvent::Warning { + kind: WarningKind::MalformedOrUnrelatedPacket, + message: "dropped malformed or unrelated packet".to_owned(), + }]); + }; + self.process_echo_reply(reply, packet_len, now, datagram.meta) } pub fn recv_available(&mut self, budget: RecvBudget) -> Result, ClientError> { @@ -320,15 +327,9 @@ impl Client { Ok(all_events) } + #[cfg(test)] fn recv_buffer_size(&self) -> usize { - match self.negotiated.as_ref() { - Some(negotiated) => { - let has_hmac = self.config.hmac_key.is_some(); - let pkt_len = echo_packet_len(has_hmac, &negotiated.params); - pkt_len.max(MIN_RECV_BUFFER_SIZE) - } - None => MIN_RECV_BUFFER_SIZE, - } + recv_buffer_size(self.config.hmac_key.is_some(), self.negotiated.as_ref()) } pub fn poll_timeouts(&mut self, now: ClientTimestamp) -> Result, ClientError> { @@ -381,28 +382,39 @@ impl Client { .map_or(0, |session| session.packets_sent) } + #[cfg(test)] fn process_received_packet( &mut self, packet: &[u8], now: ClientTimestamp, meta: ReceiveMeta, ) -> Result, ClientError> { + let packet_len = packet.len(); + let Some(reply) = self.decode_received_packet(packet) else { + return Ok(vec![ClientEvent::Warning { + kind: WarningKind::MalformedOrUnrelatedPacket, + message: "dropped malformed or unrelated packet".to_owned(), + }]); + }; + self.process_echo_reply(reply, packet_len, now, meta) + } + + fn decode_received_packet(&self, packet: &[u8]) -> Option { let negotiated = self .negotiated .as_ref() .expect("negotiated must exist when Open"); - let reply = - match decode_echo_reply(packet, &negotiated.params, self.config.hmac_key.as_deref()) { - Ok(reply) => reply, - Err(_) => { - return Ok(vec![ClientEvent::Warning { - kind: WarningKind::MalformedOrUnrelatedPacket, - message: "dropped malformed or unrelated packet".to_owned(), - }]); - } - }; + decode_echo_reply(packet, &negotiated.params, self.config.hmac_key.as_deref()).ok() + } + fn process_echo_reply( + &mut self, + reply: EchoReply, + packet_len: usize, + now: ClientTimestamp, + meta: ReceiveMeta, + ) -> Result, ClientError> { let token = match self.phase { ClientPhase::Open { token } => token, _ => unreachable!(), @@ -447,7 +459,7 @@ impl Client { server_timing, one_way, received_stats, - bytes: packet.len(), + bytes: packet_len, packet_meta: meta.into(), }); } else { @@ -461,7 +473,7 @@ impl Client { server_timing, one_way, received_stats, - bytes: packet.len(), + bytes: packet_len, packet_meta: meta.into(), }); } @@ -472,7 +484,7 @@ impl Client { seq: wire_seq, remote: self.remote, received_at: now, - bytes: packet.len(), + bytes: packet_len, }]) } else if let Some(timed_out) = session.timed_out.remove(wire_seq) { let rtt = compute_rtt(&timed_out.sent_at, &now, &reply.timestamps); @@ -494,7 +506,7 @@ impl Client { server_timing, one_way, received_stats, - bytes: packet.len(), + bytes: packet_len, packet_meta: meta.into(), }]) } else if session @@ -512,7 +524,7 @@ impl Client { server_timing: build_server_timing(&reply.timestamps), one_way: None, received_stats: build_received_stats(&reply), - bytes: packet.len(), + bytes: packet_len, packet_meta: meta.into(), }]) } else { @@ -563,12 +575,14 @@ impl Client { let negotiated = NegotiatedParams { params: reply.params.clone(), }; + let recv_buffer_size = recv_buffer_size(self.config.hmac_key.is_some(), Some(&negotiated)); let negotiated_dscp = u8::try_from(negotiated.params.dscp).map_err(|_| ClientError::InvalidConfig { reason: "negotiated dscp must be in range 0..=63".to_owned(), })?; apply_dscp_to_socket(&self.socket, self.remote, negotiated_dscp)?; self.negotiated = Some(negotiated.clone()); + self.recv_buffer.resize(recv_buffer_size, 0); self.phase = ClientPhase::Open { token: reply.token }; let duration_ns = negotiated.params.duration_ns; @@ -688,6 +702,13 @@ fn instant_abs_diff(left: Instant, right: Instant) -> Duration { .unwrap_or(Duration::ZERO) } +fn recv_buffer_size(has_hmac: bool, negotiated: Option<&NegotiatedParams>) -> usize { + match negotiated { + Some(negotiated) => echo_packet_len(has_hmac, &negotiated.params).max(MIN_RECV_BUFFER_SIZE), + None => MIN_RECV_BUFFER_SIZE, + } +} + fn compute_rtt( sent_at: &ClientTimestamp, received_at: &ClientTimestamp, diff --git a/crates/irtt-client/src/client/tests/probes_replies.rs b/crates/irtt-client/src/client/tests/probes_replies.rs index ebb6b6a..239a47a 100644 --- a/crates/irtt-client/src/client/tests/probes_replies.rs +++ b/crates/irtt-client/src/client/tests/probes_replies.rs @@ -387,6 +387,73 @@ fn recv_once_decodes_echo_reply_and_emits_event() { server.join(); } +#[test] +fn recv_once_decodes_only_received_bytes_after_longer_datagram() { + let params = default_params(); + let server_params = params.clone(); + let server = start_fake_server(move |socket, tx| { + let (_, peer) = recv_request(&socket, &tx); + let reply = open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, &server_params, None); + socket.send_to(&reply, peer).unwrap(); + + let mut seqs = Vec::new(); + for _ in 0..2 { + let (packet, _) = recv_request(&socket, &tx); + seqs.push(u32::from_le_bytes(packet[12..16].try_into().unwrap())); + } + + let mut longer = echo_reply_packet( + TOKEN, + seqs[0], + &server_params, + &TimestampFields::default(), + None, + ); + longer.extend_from_slice(b"trailing bytes"); + socket.send_to(&longer, peer).unwrap(); + + let shorter = echo_reply_packet( + TOKEN, + seqs[1], + &server_params, + &TimestampFields::default(), + None, + ); + socket.send_to(&shorter, peer).unwrap(); + }); + let config = ClientConfig { + socket_config: crate::SocketConfig { + recv_timeout: Some(Duration::from_millis(200)), + ..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().unwrap(); + client.send_probe().unwrap(); + thread::sleep(Duration::from_millis(50)); + + let _ = client.recv_once().unwrap(); + let events = client.recv_once().unwrap(); + assert_eq!(events.len(), 1); + match &events[0] { + ClientEvent::EchoReply { + seq, + logical_seq, + bytes, + .. + } => { + assert_eq!(*seq, 1); + assert_eq!(*logical_seq, 1); + assert_eq!(*bytes, echo_packet_len(false, ¶ms)); + } + other => panic!("expected EchoReply after shorter datagram, got {other:?}"), + } + server.join(); +} + #[test] #[cfg(not(all(target_os = "linux", feature = "ancillary")))] fn echo_reply_metadata_is_unavailable_without_ancillary() { @@ -1785,6 +1852,7 @@ fn recv_buffer_uses_negotiated_packet_length() { buf_size >= 4096, "recv buffer should be at least negotiated length, got {buf_size}" ); + assert_eq!(client.recv_buffer.len(), buf_size); client.close(ClientTimestamp::now()).unwrap(); server.join(); }