perf(client): reuse receive buffer

crates/irtt-client/src/client.rs

@@ -41,6 +41,7 @@ pub struct Client {
     negotiated: Option<NegotiatedParams>,
     phase: ClientPhase,
     session: Option<ActiveSession>,
+    recv_buffer: Vec<u8>,
 }
 
 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 datagram = match recv_datagram(&self.socket, &mut self.recv_buffer) {
-        let mut buf = vec![0_u8; buf_size];
-        let datagram = match recv_datagram(&self.socket, &mut buf) {
             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<Vec<ClientEvent>, ClientError> {
@@ -320,15 +327,9 @@ impl Client {
         Ok(all_events)
     }
 
+    #[cfg(test)]
     fn recv_buffer_size(&self) -> usize {
-        match self.negotiated.as_ref() {
+        recv_buffer_size(self.config.hmac_key.is_some(), 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,
-        }
     }
 
     pub fn poll_timeouts(&mut self, now: ClientTimestamp) -> Result<Vec<ClientEvent>, 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<Vec<ClientEvent>, 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<EchoReply> {
         let negotiated = self
             .negotiated
             .as_ref()
             .expect("negotiated must exist when Open");
 
-        let reply =
+        decode_echo_reply(packet, &negotiated.params, self.config.hmac_key.as_deref()).ok()
-            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(),
-                    }]);
     }
-            };
 
+    fn process_echo_reply(
+        &mut self,
+        reply: EchoReply,
+        packet_len: usize,
+        now: ClientTimestamp,
+        meta: ReceiveMeta,
+    ) -> Result<Vec<ClientEvent>, 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,

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, &params));
+        }
+        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();
 }