Adapt stream retransmit pacing to RTT

src/transport.rs

@@ -29,6 +29,8 @@ use tokio::net::UdpSocket;
 
 pub const DEFAULT_INITIAL_WINDOW: usize = 1024 * 1024;
 pub const DEFAULT_RETRANSMIT_AFTER: Duration = Duration::from_millis(200);
+const ADAPTIVE_RETRANSMIT_PAD: Duration = Duration::from_millis(10);
+const ADAPTIVE_RETRANSMIT_MIN: Duration = Duration::from_millis(10);
 pub const DEFAULT_KEEPALIVE_AFTER: Duration = Duration::from_secs(2);
 pub const DEFAULT_RETIRED_STREAM_TOMBSTONES: usize = 16 * 1024;
 pub const SERVICE_TARGET_PREFIX: &str = "@dosh-";
@@ -178,6 +180,7 @@ pub struct StreamMux {
     recv_pending: HashMap<u64, BTreeMap<u64, Vec<u8>>>,
     retired_streams: HashSet<u64>,
     retired_stream_order: VecDeque<u64>,
+    srtt: Option<Duration>,
 }
 
 #[derive(Debug, Clone, Default)]
@@ -235,6 +238,7 @@ impl StreamMux {
             recv_pending: HashMap::new(),
             retired_streams: HashSet::new(),
             retired_stream_order: VecDeque::new(),
+            srtt: None,
         }
     }
 
@@ -326,8 +330,11 @@ impl StreamMux {
         &mut self,
         ok: StreamOpenOk,
     ) -> Result<Option<Vec<OutgoingStreamPacket>>> {
-        if self.pending_opens.remove(&ok.stream_id).is_none() {
+        let Some(pending) = self.pending_opens.remove(&ok.stream_id) else {
             return Ok(None);
+        };
+        if pending.attempts == 1 {
+            self.observe_rtt_sample(pending.last_sent.elapsed());
         }
         self.opened_streams.insert(ok.stream_id);
         self.send_credit
@@ -499,10 +506,11 @@ impl StreamMux {
 
     pub fn tick(&mut self) -> Result<Vec<OutgoingStreamPacket>> {
         let now = Instant::now();
+        let retransmit_after = self.effective_retransmit_after();
         let mut out = Vec::new();
         let mut retransmit_opens = Vec::new();
         for (stream_id, pending) in self.pending_opens.iter_mut() {
-            if now.duration_since(pending.last_sent) < self.config.retransmit_after {
+            if now.duration_since(pending.last_sent) < retransmit_after {
                 continue;
             }
             pending.last_sent = now;
@@ -523,7 +531,7 @@ impl StreamMux {
         let mut retransmit_data = Vec::new();
         for (stream_id, chunks) in self.sent_data.iter_mut() {
             for chunk in chunks.values_mut() {
-                if now.duration_since(chunk.last_sent) < self.config.retransmit_after {
+                if now.duration_since(chunk.last_sent) < retransmit_after {
                     continue;
                 }
                 chunk.last_sent = now;
@@ -544,6 +552,20 @@ impl StreamMux {
         Ok(out)
     }
 
+    pub fn effective_retransmit_after(&self) -> Duration {
+        let Some(srtt) = self.srtt else {
+            return self.config.retransmit_after;
+        };
+        let adaptive = srtt
+            .saturating_add(ADAPTIVE_RETRANSMIT_PAD)
+            .max(ADAPTIVE_RETRANSMIT_MIN);
+        if self.config.retransmit_after < ADAPTIVE_RETRANSMIT_MIN {
+            self.config.retransmit_after
+        } else {
+            adaptive.min(self.config.retransmit_after)
+        }
+    }
+
     pub fn is_open(&self, stream_id: u64) -> bool {
         self.opened_streams.contains(&stream_id)
     }
@@ -654,28 +676,48 @@ impl StreamMux {
     }
 
     fn ack_data(&mut self, stream_id: u64, received_offset: u64) -> usize {
-        let Some(sent) = self.sent_data.get_mut(&stream_id) else {
-            return 0;
-        };
-        let acked_offsets = sent
-            .iter()
-            .filter_map(|(offset, chunk)| {
-                let end = chunk.offset.saturating_add(chunk.bytes.len() as u64);
-                (end <= received_offset).then_some(*offset)
-            })
-            .collect::<Vec<_>>();
         let mut acked_bytes = 0usize;
-        for offset in acked_offsets {
-            if let Some(chunk) = sent.remove(&offset) {
-                acked_bytes = acked_bytes.saturating_add(chunk.bytes.len());
+        let mut samples = Vec::new();
+        let remove_stream = {
+            let Some(sent) = self.sent_data.get_mut(&stream_id) else {
+                return 0;
+            };
+            let acked_offsets = sent
+                .iter()
+                .filter_map(|(offset, chunk)| {
+                    let end = chunk.offset.saturating_add(chunk.bytes.len() as u64);
+                    (end <= received_offset).then_some(*offset)
+                })
+                .collect::<Vec<_>>();
+            for offset in acked_offsets {
+                if let Some(chunk) = sent.remove(&offset) {
+                    acked_bytes = acked_bytes.saturating_add(chunk.bytes.len());
+                    if chunk.attempts == 1 {
+                        samples.push(chunk.last_sent.elapsed());
+                    }
+                }
             }
-        }
-        if sent.is_empty() {
+            sent.is_empty()
+        };
+        if remove_stream {
             self.sent_data.remove(&stream_id);
         }
+        for sample in samples {
+            self.observe_rtt_sample(sample);
+        }
         acked_bytes
     }
 
+    fn observe_rtt_sample(&mut self, sample: Duration) {
+        self.srtt = Some(match self.srtt {
+            None => sample,
+            Some(srtt) => {
+                let smoothed = ((srtt.as_micros() * 7) + sample.as_micros()) / 8;
+                Duration::from_micros(smoothed.min(u64::MAX as u128) as u64)
+            }
+        });
+    }
+
     fn add_credit(&mut self, stream_id: u64, bytes: usize) {
         let credit = self.send_credit.entry(stream_id).or_default();
         *credit = credit.saturating_add(bytes).min(self.config.initial_window);
@@ -1184,6 +1226,90 @@ mod tests {
         assert_eq!(data.bytes, b"!");
     }
 
+    #[test]
+    fn stream_open_ack_lowers_retransmit_timeout_from_observed_rtt() {
+        let mut mux = StreamMux::new(TransportConfig {
+            retransmit_after: Duration::from_secs(1),
+            ..TransportConfig::default()
+        });
+        mux.open_stream(1, "@dosh-fast", 0).unwrap();
+        mux.pending_opens.get_mut(&1).unwrap().last_sent =
+            Instant::now() - Duration::from_millis(30);
+        mux.handle_open_ok(StreamOpenOk { stream_id: 1 })
+            .unwrap()
+            .unwrap();
+
+        assert!(mux.effective_retransmit_after() < Duration::from_millis(100));
+        assert!(mux.effective_retransmit_after() >= Duration::from_millis(10));
+
+        mux.open_stream(2, "@dosh-fast", 0).unwrap();
+        mux.pending_opens.get_mut(&2).unwrap().last_sent =
+            Instant::now() - Duration::from_millis(50);
+        let retransmits = mux.tick().unwrap();
+        assert_eq!(retransmits.len(), 1);
+        assert_eq!(retransmits[0].kind, PacketKind::StreamOpen);
+        let open: StreamOpen = decode(&retransmits[0]);
+        assert_eq!(open.stream_id, 2);
+    }
+
+    #[test]
+    fn data_ack_lowers_retransmit_timeout_from_first_attempt_rtt() {
+        let mut mux = StreamMux::new(TransportConfig {
+            retransmit_after: Duration::from_secs(1),
+            ..TransportConfig::default()
+        });
+        mux.open_stream(1, "@dosh-fast", 0).unwrap();
+        mux.handle_open_ok(StreamOpenOk { stream_id: 1 })
+            .unwrap()
+            .unwrap();
+        mux.send_data(1, b"hello".to_vec()).unwrap();
+        mux.sent_data
+            .get_mut(&1)
+            .unwrap()
+            .get_mut(&0)
+            .unwrap()
+            .last_sent = Instant::now() - Duration::from_millis(25);
+
+        mux.handle_window_adjust(StreamWindowAdjust {
+            stream_id: 1,
+            received_offset: 5,
+            bytes: 5,
+        })
+        .unwrap()
+        .unwrap();
+
+        assert!(mux.effective_retransmit_after() < Duration::from_millis(100));
+    }
+
+    #[test]
+    fn retransmitted_data_ack_does_not_update_rtt_estimate() {
+        let mut mux = StreamMux::new(TransportConfig {
+            retransmit_after: Duration::from_secs(1),
+            ..TransportConfig::default()
+        });
+        mux.open_stream(1, "@dosh-fast", 0).unwrap();
+        mux.pending_opens.get_mut(&1).unwrap().last_sent =
+            Instant::now() - Duration::from_millis(30);
+        mux.handle_open_ok(StreamOpenOk { stream_id: 1 })
+            .unwrap()
+            .unwrap();
+        let learned = mux.effective_retransmit_after();
+
+        mux.send_data(1, b"hello".to_vec()).unwrap();
+        let chunk = mux.sent_data.get_mut(&1).unwrap().get_mut(&0).unwrap();
+        chunk.attempts = 2;
+        chunk.last_sent = Instant::now() - Duration::from_millis(900);
+        mux.handle_window_adjust(StreamWindowAdjust {
+            stream_id: 1,
+            received_offset: 5,
+            bytes: 5,
+        })
+        .unwrap()
+        .unwrap();
+
+        assert_eq!(mux.effective_retransmit_after(), learned);
+    }
+
     #[test]
     fn cumulative_window_adjust_restores_credit_even_if_delta_was_lost() {
         let mut mux = StreamMux::new(TransportConfig {