Adapt runtime stream retransmit pacing

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<u64, PendingStreamOpen> = HashMap::new();
     let mut stream_next_send_offset: HashMap<u64, u64> = HashMap::new();
     let mut stream_sent_data: HashMap<u64, BTreeMap<u64, PendingStreamChunk>> = HashMap::new();
+    let mut stream_retransmit_srtt: Option<Duration> = None;
     let mut stream_next_recv_offset: HashMap<u64, u64> = HashMap::new();
     let mut stream_recv_pending: HashMap<u64, BTreeMap<u64, Vec<u8>>> = HashMap::new();
     let mut stream_retired: HashSet<u64> = 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<u64, BTreeMap<u64, PendingStreamChunk>>,
+    stream_retransmit_srtt: Option<Duration>,
 ) -> 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<u64, PendingStreamOpen>,
+    stream_retransmit_srtt: Option<Duration>,
 ) -> 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<u64, BTreeMap<u64, PendingStreamChunk>>,
     stream_send_credit: &mut HashMap<u64, usize>,
+    stream_retransmit_srtt: &mut Option<Duration>,
     stream_id: u64,
     received_offset: u64,
 ) {
-    let Some(sent) = stream_sent_data.get_mut(&stream_id) else {
-        return;
-    };
-    let acked_offsets: Vec<u64> = 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<u64> = 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!(

src/bin/dosh-server.rs

@@ -408,7 +408,6 @@ struct Session {
 struct RestoredScreenSnapshot {
     cols: u16,
     rows: u16,
-    output_seq: u64,
     snapshot: Vec<u8>,
 }
 
@@ -436,6 +435,7 @@ struct ClientState {
     stream_recv_pending: HashMap<u64, BTreeMap<u64, Vec<u8>>>,
     stream_retired: HashSet<u64>,
     stream_retired_order: VecDeque<u64>,
+    stream_retransmit_srtt: Option<Duration>,
     /// 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<u8> {
     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<u64> = 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<u64> = 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<u64> =
                     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,

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>) -> 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<Duration>, 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<UdpSocket>,
     role: SessionRole,