From 4b82334e021f37e7e744b311e4bab017a58658fa Mon Sep 17 00:00:00 2001 From: DuProcess <273172371+DuProcess@users.noreply.github.com> Date: Sun, 12 Jul 2026 23:11:25 -0400 Subject: [PATCH] Adapt runtime stream retransmit pacing --- src/bin/dosh-client.rs | 117 +++++++++++++++++++++++++++++------- src/bin/dosh-server.rs | 131 +++++++++++++++++++++++++++++++++-------- src/transport.rs | 48 ++++++++------- 3 files changed, 232 insertions(+), 64 deletions(-) diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index 9edaa62..5ade700 100644 --- a/src/bin/dosh-client.rs +++ b/src/bin/dosh-client.rs @@ -5850,7 +5850,8 @@ async fn run_terminal( let mut last_packet_at = Instant::now(); let mut status_tick = tokio::time::interval(Duration::from_secs(1)); let mut resize_tick = tokio::time::interval(Duration::from_millis(250)); - let mut stream_retransmit_tick = tokio::time::interval(Duration::from_millis(200)); + let mut stream_retransmit_tick = + tokio::time::interval(dosh::transport::ADAPTIVE_RETRANSMIT_MIN); let mut frame_gap_tick = tokio::time::interval(Duration::from_millis(250)); let mut last_size = terminal_size(); // React to terminal resize the instant it happens via SIGWINCH (mosh-style), @@ -5896,6 +5897,7 @@ async fn run_terminal( let mut stream_pending_opens: HashMap = HashMap::new(); let mut stream_next_send_offset: HashMap = HashMap::new(); let mut stream_sent_data: HashMap> = HashMap::new(); + let mut stream_retransmit_srtt: Option = None; let mut stream_next_recv_offset: HashMap = HashMap::new(); let mut stream_recv_pending: HashMap>> = HashMap::new(); let mut stream_retired: HashSet = HashSet::new(); @@ -6746,8 +6748,14 @@ async fn run_terminal( continue; }; last_packet_at = Instant::now(); - if stream_pending_opens.remove(&ok.stream_id).is_none() { + let Some(pending_open) = stream_pending_opens.remove(&ok.stream_id) else { continue; + }; + if pending_open.attempts == 1 { + dosh::transport::observe_retransmit_rtt( + &mut stream_retransmit_srtt, + pending_open.last_sent.elapsed(), + ); } opened_streams.insert(ok.stream_id); stream_send_credit.entry(ok.stream_id).or_insert(STREAM_INITIAL_WINDOW); @@ -6851,6 +6859,7 @@ async fn run_terminal( ack_stream_data( &mut stream_sent_data, &mut stream_send_credit, + &mut stream_retransmit_srtt, adjust.stream_id, adjust.received_offset, ); @@ -7105,6 +7114,7 @@ async fn run_terminal( &cred, &mut send_seq, &mut stream_sent_data, + stream_retransmit_srtt, ).await?; retransmit_stream_opens( &socket, @@ -7112,6 +7122,7 @@ async fn run_terminal( &cred, &mut send_seq, &mut stream_pending_opens, + stream_retransmit_srtt, ).await?; } } @@ -8027,12 +8038,17 @@ async fn retransmit_stream_data( cred: &CachedCredential, send_seq: &mut u64, stream_sent_data: &mut HashMap>, + stream_retransmit_srtt: Option, ) -> Result<()> { let now = Instant::now(); + let retransmit_after = dosh::transport::adaptive_retransmit_after( + dosh::transport::DEFAULT_RETRANSMIT_AFTER, + stream_retransmit_srtt, + ); let mut retransmit = Vec::new(); for (stream_id, chunks) in stream_sent_data.iter_mut() { for chunk in chunks.values_mut() { - if now.duration_since(chunk.last_sent) < Duration::from_millis(200) { + if now.duration_since(chunk.last_sent) < retransmit_after { continue; } chunk.last_sent = now; @@ -8052,11 +8068,16 @@ async fn retransmit_stream_opens( cred: &CachedCredential, send_seq: &mut u64, stream_pending_opens: &mut HashMap, + stream_retransmit_srtt: Option, ) -> Result<()> { let now = Instant::now(); + let retransmit_after = dosh::transport::adaptive_retransmit_after( + dosh::transport::DEFAULT_RETRANSMIT_AFTER, + stream_retransmit_srtt, + ); let mut retransmit = Vec::new(); for (stream_id, pending) in stream_pending_opens.iter_mut() { - if now.duration_since(pending.last_sent) < Duration::from_millis(200) { + if now.duration_since(pending.last_sent) < retransmit_after { continue; } pending.last_sent = now; @@ -8215,28 +8236,39 @@ fn retire_stream_state( fn ack_stream_data( stream_sent_data: &mut HashMap>, stream_send_credit: &mut HashMap, + stream_retransmit_srtt: &mut Option, stream_id: u64, received_offset: u64, ) { - let Some(sent) = stream_sent_data.get_mut(&stream_id) else { - return; - }; - let acked_offsets: Vec = 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) = stream_sent_data.get_mut(&stream_id) else { + return; + }; + let acked_offsets: Vec = 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 { stream_sent_data.remove(&stream_id); } + for sample in samples { + dosh::transport::observe_retransmit_rtt(stream_retransmit_srtt, sample); + } add_stream_credit(stream_send_credit, stream_id, acked_bytes); } @@ -10794,7 +10826,7 @@ mod tests { }, )]); - retransmit_stream_opens(&sender, addr, &cred, &mut send_seq, &mut pending) + retransmit_stream_opens(&sender, addr, &cred, &mut send_seq, &mut pending, None) .await .unwrap(); @@ -10815,6 +10847,51 @@ mod tests { assert_eq!(pending[&44].attempts, 2); } + #[tokio::test] + async fn client_stream_retransmit_uses_observed_rtt() { + let receiver = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let sender = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let addr = receiver.local_addr().unwrap(); + let cred = CachedCredential { + server: "host".to_string(), + session: "term".to_string(), + mode: "forward-only".to_string(), + udp_host: "127.0.0.1".to_string(), + udp_port: addr.port(), + client_id: [1; 16], + session_key: [2; 32], + session_key_id: protocol::session_key_id(&[2; 32]), + attach_ticket: Vec::new(), + attach_ticket_psk: [3; 32], + last_rendered_seq: 0, + }; + let mut send_seq = 10; + let mut pending = HashMap::from([( + 44, + PendingStreamOpen { + target_host: "127.0.0.1".to_string(), + target_port: 8080, + last_sent: Instant::now() - Duration::from_millis(50), + attempts: 1, + }, + )]); + let mut srtt = None; + dosh::transport::observe_retransmit_rtt(&mut srtt, Duration::from_millis(25)); + + retransmit_stream_opens(&sender, addr, &cred, &mut send_seq, &mut pending, srtt) + .await + .unwrap(); + + let mut buf = [0u8; 2048]; + let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf)) + .await + .unwrap() + .unwrap(); + let packet = protocol::decode(&buf[..n]).unwrap(); + assert_eq!(packet.header.kind, PacketKind::StreamOpen); + assert_eq!(pending[&44].attempts, 2); + } + #[test] fn update_options_accept_check_and_role_flags() { assert_eq!( diff --git a/src/bin/dosh-server.rs b/src/bin/dosh-server.rs index d12a778..649a861 100644 --- a/src/bin/dosh-server.rs +++ b/src/bin/dosh-server.rs @@ -408,7 +408,6 @@ struct Session { struct RestoredScreenSnapshot { cols: u16, rows: u16, - output_seq: u64, snapshot: Vec, } @@ -436,6 +435,7 @@ struct ClientState { stream_recv_pending: HashMap>>, stream_retired: HashSet, stream_retired_order: VecDeque, + stream_retransmit_srtt: Option, /// Current transport key epoch (0 = original handshake key). Bumped on rekey. epoch: u64, /// When the current epoch began, for the wall-clock rekey trigger. @@ -674,7 +674,6 @@ impl ServerState { restored_screen = Some(RestoredScreenSnapshot { cols, rows, - output_seq, snapshot: saved.snapshot, }); } @@ -804,7 +803,6 @@ fn attach_snapshot(session: &Session, cols: u16, rows: u16) -> Vec { let cols = cols.max(1); let rows = rows.max(1); if let Some(restored) = &session.restored_screen - && (restored.output_seq == session.output_seq || session.clients.is_empty()) && restored.cols == cols && restored.rows == rows { @@ -1272,6 +1270,7 @@ async fn handle_native_user_auth( stream_recv_pending: HashMap::new(), stream_retired: HashSet::new(), stream_retired_order: VecDeque::new(), + stream_retransmit_srtt: None, epoch: 0, epoch_started: Instant::now(), epoch_packets: 0, @@ -1421,6 +1420,7 @@ async fn handle_bootstrap_attach( stream_recv_pending: HashMap::new(), stream_retired: HashSet::new(), stream_retired_order: VecDeque::new(), + stream_retransmit_srtt: None, epoch: 0, epoch_started: Instant::now(), epoch_packets: 0, @@ -1570,6 +1570,7 @@ async fn handle_ticket_attach( stream_recv_pending: HashMap::new(), stream_retired: HashSet::new(), stream_retired_order: VecDeque::new(), + stream_retransmit_srtt: None, epoch: 0, epoch_started: Instant::now(), epoch_packets: 0, @@ -2289,8 +2290,14 @@ async fn handle_stream_open_ok( } client.endpoint = peer; client.last_seen = Instant::now(); - if client.stream_pending_opens.remove(&ok.stream_id).is_none() { + let Some(pending_open) = client.stream_pending_opens.remove(&ok.stream_id) else { return Ok(()); + }; + if pending_open.attempts == 1 { + dosh::transport::observe_retransmit_rtt( + &mut client.stream_retransmit_srtt, + pending_open.last_sent.elapsed(), + ); } client.opened_streams.insert(ok.stream_id); client @@ -3580,25 +3587,35 @@ fn retire_client_stream(client: &mut ClientState, stream_id: u64) { } fn ack_stream_data(client: &mut ClientState, stream_id: u64, received_offset: u64) { - let Some(sent) = client.stream_sent_data.get_mut(&stream_id) else { - return; - }; - let acked_offsets: Vec = 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) = client.stream_sent_data.get_mut(&stream_id) else { + return; + }; + let acked_offsets: Vec = 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 { client.stream_sent_data.remove(&stream_id); } + for sample in samples { + dosh::transport::observe_retransmit_rtt(&mut client.stream_retransmit_srtt, sample); + } add_stream_credit(&mut client.stream_send_credit, stream_id, acked_bytes); } @@ -3882,11 +3899,15 @@ async fn retransmit_pending( let pending_open_ids: Vec = client.stream_pending_opens.keys().copied().collect(); let mut retransmit_opens = Vec::new(); + let stream_retransmit_after = dosh::transport::adaptive_retransmit_after( + dosh::transport::DEFAULT_RETRANSMIT_AFTER, + client.stream_retransmit_srtt, + ); for stream_id in pending_open_ids { let Some(open) = client.stream_pending_opens.get_mut(&stream_id) else { continue; }; - if now.duration_since(open.last_sent) < Duration::from_millis(200) { + if now.duration_since(open.last_sent) < stream_retransmit_after { continue; } open.last_sent = now; @@ -3918,7 +3939,7 @@ async fn retransmit_pending( continue; }; for chunk in chunks.values_mut() { - if now.duration_since(chunk.last_sent) < Duration::from_millis(200) { + if now.duration_since(chunk.last_sent) < stream_retransmit_after { continue; } chunk.last_sent = now; @@ -4312,7 +4333,6 @@ mod tests { restored_screen: Some(RestoredScreenSnapshot { cols: 80, rows: 24, - output_seq: 7, snapshot: restored.clone(), }), clients: HashMap::new(), @@ -4340,7 +4360,6 @@ mod tests { restored_screen: Some(RestoredScreenSnapshot { cols: 80, rows: 24, - output_seq: 7, snapshot: restored.clone(), }), clients: HashMap::new(), @@ -4360,6 +4379,9 @@ mod tests { session .clients .insert([1u8; 16], test_client_state([2u8; 32])); + assert_eq!(attach_snapshot(&session, 80, 24), restored); + + session.restored_screen = None; assert_ne!(attach_snapshot(&session, 80, 24), restored); } @@ -4385,7 +4407,6 @@ mod tests { restored_screen: Some(RestoredScreenSnapshot { cols: 80, rows: 24, - output_seq: 7, snapshot: restored.clone(), }), clients: HashMap::new(), @@ -4507,6 +4528,7 @@ mod tests { stream_recv_pending: HashMap::new(), stream_retired: HashSet::new(), stream_retired_order: VecDeque::new(), + stream_retransmit_srtt: None, epoch: 0, epoch_started: Instant::now(), epoch_packets: 0, @@ -4974,6 +4996,65 @@ mod tests { assert_eq!(open.target_port, 8080); } + #[tokio::test] + async fn server_stream_retransmit_uses_observed_rtt() { + let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); + let client_id = [15u8; 16]; + let session_key = [16u8; 32]; + let stream_id = 124; + let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let sender = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap()); + let mut client = test_client_state(session_key); + client.endpoint = receiver.local_addr().unwrap(); + dosh::transport::observe_retransmit_rtt( + &mut client.stream_retransmit_srtt, + Duration::from_millis(25), + ); + client.stream_pending_opens.insert( + stream_id, + PendingStreamOpen { + target_host: "127.0.0.1".to_string(), + target_port: 8081, + last_sent: Instant::now() - Duration::from_millis(50), + attempts: 1, + }, + ); + state.sessions.insert( + "test".to_string(), + Session { + pty: None, + parser: vt100::Parser::new(24, 80, 100), + restored_screen: None, + clients: HashMap::from([(client_id, client)]), + output_seq: 0, + recent: VecDeque::new(), + empty_since: None, + holder_control: None, + persistent: false, + bytes_since_persist: 0, + last_persisted_seq: 0, + last_screen_persist_at: Instant::now(), + }, + ); + state.client_index.insert(client_id, "test".to_string()); + let state = Arc::new(Mutex::new(state)); + + retransmit_pending(&state, &sender).await.unwrap(); + + let mut buf = [0u8; 2048]; + let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf)) + .await + .unwrap() + .unwrap(); + let packet = protocol::decode(&buf[..n]).unwrap(); + assert_eq!(packet.header.kind, PacketKind::StreamOpen); + let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap(); + let open: StreamOpen = protocol::from_body(&plain).unwrap(); + assert_eq!(open.stream_id, stream_id); + assert_eq!(open.target_port, 8081); + } + #[test] fn forward_only_session_does_not_allocate_pty() { let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); @@ -5111,6 +5192,7 @@ mod tests { stream_recv_pending: HashMap::new(), stream_retired: HashSet::new(), stream_retired_order: VecDeque::new(), + stream_retransmit_srtt: None, epoch: 0, epoch_started: Instant::now(), epoch_packets: 0, @@ -5209,6 +5291,7 @@ mod tests { stream_recv_pending: HashMap::new(), stream_retired: HashSet::new(), stream_retired_order: VecDeque::new(), + stream_retransmit_srtt: None, epoch: 0, epoch_started: Instant::now(), epoch_packets: 0, diff --git a/src/transport.rs b/src/transport.rs index 7e69f6b..fa401e7 100644 --- a/src/transport.rs +++ b/src/transport.rs @@ -29,8 +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 ADAPTIVE_RETRANSMIT_PAD: Duration = Duration::from_millis(10); +pub 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-"; @@ -553,17 +553,7 @@ impl StreamMux { } 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) - } + adaptive_retransmit_after(self.config.retransmit_after, self.srtt) } pub fn is_open(&self, stream_id: u64) -> bool { @@ -709,13 +699,7 @@ impl StreamMux { } 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) - } - }); + observe_retransmit_rtt(&mut self.srtt, sample); } fn add_credit(&mut self, stream_id: u64, bytes: usize) { @@ -761,6 +745,30 @@ impl StreamMux { } } +pub fn adaptive_retransmit_after(configured: Duration, srtt: Option) -> Duration { + let Some(srtt) = srtt else { + return configured; + }; + let adaptive = srtt + .saturating_add(ADAPTIVE_RETRANSMIT_PAD) + .max(ADAPTIVE_RETRANSMIT_MIN); + if configured < ADAPTIVE_RETRANSMIT_MIN { + configured + } else { + adaptive.min(configured) + } +} + +pub fn observe_retransmit_rtt(srtt: &mut Option, sample: Duration) { + *srtt = Some(match *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) + } + }); +} + pub struct DoshTransport { socket: Arc, role: SessionRole,