diff --git a/src/transport.rs b/src/transport.rs index 5674f82..7e69f6b 100644 --- a/src/transport.rs +++ b/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>>, retired_streams: HashSet, retired_stream_order: VecDeque, + srtt: Option, } #[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>> { - 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> { 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::>(); 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::>(); + 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 {