Retransmit stream controls in transport

src/transport.rs

@@ -9,7 +9,8 @@
 //! 1. authenticate and maintain a Dosh session using [`crate::protocol`],
 //! 2. encrypt each [`OutgoingStreamPacket::body`] as the given packet kind,
 //! 3. pass received stream packet bodies back into [`StreamMux`],
-//! 4. call [`StreamMux::tick`] periodically to retransmit lost stream data.
+//! 4. call [`StreamMux::tick`] periodically to retransmit lost stream data and
+//!    idempotent stream control packets.
 //!
 //! That split is intentional: Dosh roaming/reconnect stays at the authenticated
 //! session layer, while this module provides the reusable fast persistent byte
@@ -33,6 +34,7 @@ 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 STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS: u32 = 8;
 pub const SERVICE_TARGET_PREFIX: &str = "@dosh-";
 pub const MAX_STREAM_DATA_BYTES: usize = 60 * 1024;
 
@@ -171,6 +173,20 @@ struct PendingStreamOpen {
     attempts: u32,
 }
 
+#[derive(Debug, Clone)]
+struct PendingStreamControl {
+    last_sent: Instant,
+    attempts: u32,
+}
+
+#[derive(Debug, Clone)]
+struct PendingWindowAdjust {
+    received_offset: u64,
+    bytes: usize,
+    last_sent: Instant,
+    attempts: u32,
+}
+
 #[derive(Debug, Clone)]
 pub struct StreamMux {
     config: TransportConfig,
@@ -182,6 +198,9 @@ pub struct StreamMux {
     next_send_offset: HashMap<u64, u64>,
     next_recv_offset: HashMap<u64, u64>,
     recv_pending: HashMap<u64, BTreeMap<u64, Vec<u8>>>,
+    pending_eofs: HashMap<u64, PendingStreamControl>,
+    pending_closes: HashMap<u64, PendingStreamControl>,
+    pending_window_adjusts: HashMap<u64, PendingWindowAdjust>,
     retired_streams: HashSet<u64>,
     retired_stream_order: VecDeque<u64>,
     srtt: Option<Duration>,
@@ -240,6 +259,9 @@ impl StreamMux {
             next_send_offset: HashMap::new(),
             next_recv_offset: HashMap::new(),
             recv_pending: HashMap::new(),
+            pending_eofs: HashMap::new(),
+            pending_closes: HashMap::new(),
+            pending_window_adjusts: HashMap::new(),
             retired_streams: HashSet::new(),
             retired_stream_order: VecDeque::new(),
             srtt: None,
@@ -411,6 +433,17 @@ impl StreamMux {
                 bytes: consumed.min(u32::MAX as usize) as u32,
             },
         )?;
+        if consumed > 0 {
+            self.pending_window_adjusts.insert(
+                stream_id,
+                PendingWindowAdjust {
+                    received_offset,
+                    bytes: consumed,
+                    last_sent: Instant::now(),
+                    attempts: 1,
+                },
+            );
+        }
         Ok(Some(IncomingStreamData {
             stream_id,
             chunks,
@@ -437,13 +470,27 @@ impl StreamMux {
             bail!("stream {stream_id} is not open");
         }
         self.retire_stream(stream_id);
+        self.pending_closes.insert(
+            stream_id,
+            PendingStreamControl {
+                last_sent: Instant::now(),
+                attempts: 1,
+            },
+        );
         encode_packet(PacketKind::StreamClose, &StreamClose { stream_id })
     }
 
-    pub fn eof_stream(&self, stream_id: u64) -> Result<OutgoingStreamPacket> {
+    pub fn eof_stream(&mut self, stream_id: u64) -> Result<OutgoingStreamPacket> {
         if !self.has_stream_state(stream_id) {
             bail!("stream {stream_id} is not open");
         }
+        self.pending_eofs.insert(
+            stream_id,
+            PendingStreamControl {
+                last_sent: Instant::now(),
+                attempts: 1,
+            },
+        );
         encode_packet(PacketKind::StreamEof, &StreamEof { stream_id })
     }
 
@@ -452,7 +499,9 @@ impl StreamMux {
     }
 
     pub fn handle_close(&mut self, close: StreamClose) -> bool {
-        if !self.has_stream_state(close.stream_id) {
+        let had_state = self.has_stream_state(close.stream_id);
+        self.pending_closes.remove(&close.stream_id);
+        if !had_state {
             return false;
         }
         self.retire_stream(close.stream_id);
@@ -579,6 +628,79 @@ impl StreamMux {
                 },
             )?);
         }
+
+        let mut retransmit_eofs = Vec::new();
+        for (stream_id, pending) in self.pending_eofs.iter_mut() {
+            if !self.opened_streams.contains(stream_id) {
+                continue;
+            }
+            if pending.attempts >= STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS {
+                continue;
+            }
+            if now.duration_since(pending.last_sent) < retransmit_after {
+                continue;
+            }
+            pending.last_sent = now;
+            pending.attempts = pending.attempts.saturating_add(1);
+            retransmit_eofs.push(*stream_id);
+        }
+        let opened_streams = self.opened_streams.clone();
+        self.pending_eofs.retain(|stream_id, pending| {
+            opened_streams.contains(stream_id)
+                && pending.attempts < STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS
+        });
+        for stream_id in retransmit_eofs {
+            out.push(encode_packet(
+                PacketKind::StreamEof,
+                &StreamEof { stream_id },
+            )?);
+        }
+
+        let mut retransmit_closes = Vec::new();
+        for (stream_id, pending) in self.pending_closes.iter_mut() {
+            if pending.attempts >= STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS {
+                continue;
+            }
+            if now.duration_since(pending.last_sent) < retransmit_after {
+                continue;
+            }
+            pending.last_sent = now;
+            pending.attempts = pending.attempts.saturating_add(1);
+            retransmit_closes.push(*stream_id);
+        }
+        self.pending_closes
+            .retain(|_, pending| pending.attempts < STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS);
+        for stream_id in retransmit_closes {
+            out.push(encode_packet(
+                PacketKind::StreamClose,
+                &StreamClose { stream_id },
+            )?);
+        }
+
+        let mut retransmit_adjusts = Vec::new();
+        for (stream_id, pending) in self.pending_window_adjusts.iter_mut() {
+            if pending.attempts >= STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS {
+                continue;
+            }
+            if now.duration_since(pending.last_sent) < retransmit_after {
+                continue;
+            }
+            pending.last_sent = now;
+            pending.attempts = pending.attempts.saturating_add(1);
+            retransmit_adjusts.push((*stream_id, pending.received_offset, pending.bytes));
+        }
+        self.pending_window_adjusts
+            .retain(|_, pending| pending.attempts < STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS);
+        for (stream_id, received_offset, bytes) in retransmit_adjusts {
+            out.push(encode_packet(
+                PacketKind::StreamWindowAdjust,
+                &StreamWindowAdjust {
+                    stream_id,
+                    received_offset,
+                    bytes: bytes.min(u32::MAX as usize) as u32,
+                },
+            )?);
+        }
         Ok(out)
     }
 
@@ -738,6 +860,8 @@ impl StreamMux {
             || self.next_send_offset.contains_key(&stream_id)
             || self.next_recv_offset.contains_key(&stream_id)
             || self.recv_pending.contains_key(&stream_id)
+            || self.pending_eofs.contains_key(&stream_id)
+            || self.pending_window_adjusts.contains_key(&stream_id)
     }
 
     fn retire_stream(&mut self, stream_id: u64) {
@@ -764,6 +888,8 @@ impl StreamMux {
         self.next_send_offset.remove(&stream_id);
         self.next_recv_offset.remove(&stream_id);
         self.recv_pending.remove(&stream_id);
+        self.pending_eofs.remove(&stream_id);
+        self.pending_window_adjusts.remove(&stream_id);
     }
 }
 
@@ -1691,6 +1817,160 @@ mod tests {
         assert_eq!(mux.send_data(1, b"!".to_vec()).unwrap().len(), 1);
     }
 
+    #[test]
+    fn tick_retransmits_pending_stream_eof() {
+        let mut mux = StreamMux::new(TransportConfig {
+            retransmit_after: Duration::from_millis(200),
+            ..TransportConfig::default()
+        });
+        mux.accept_open(StreamOpen {
+            stream_id: 1,
+            target_host: "@dosh-test".to_string(),
+            target_port: 0,
+        })
+        .unwrap();
+
+        let eof = mux.eof_stream(1).unwrap();
+        assert_eq!(eof.kind, PacketKind::StreamEof);
+        mux.pending_eofs.get_mut(&1).unwrap().last_sent =
+            Instant::now() - Duration::from_millis(250);
+
+        let retransmits = mux.tick().unwrap();
+        assert_eq!(retransmits.len(), 1);
+        assert_eq!(retransmits[0].kind, PacketKind::StreamEof);
+        let eof: StreamEof = decode(&retransmits[0]);
+        assert_eq!(eof.stream_id, 1);
+        assert_eq!(mux.pending_eofs[&1].attempts, 2);
+    }
+
+    #[test]
+    fn tick_retransmits_pending_stream_close() {
+        let mut mux = StreamMux::new(TransportConfig {
+            retransmit_after: Duration::from_millis(200),
+            ..TransportConfig::default()
+        });
+        mux.accept_open(StreamOpen {
+            stream_id: 1,
+            target_host: "@dosh-test".to_string(),
+            target_port: 0,
+        })
+        .unwrap();
+
+        let close = mux.close_stream(1).unwrap();
+        assert_eq!(close.kind, PacketKind::StreamClose);
+        mux.pending_closes.get_mut(&1).unwrap().last_sent =
+            Instant::now() - Duration::from_millis(250);
+
+        let retransmits = mux.tick().unwrap();
+        assert_eq!(retransmits.len(), 1);
+        assert_eq!(retransmits[0].kind, PacketKind::StreamClose);
+        let close: StreamClose = decode(&retransmits[0]);
+        assert_eq!(close.stream_id, 1);
+        assert_eq!(mux.pending_closes[&1].attempts, 2);
+    }
+
+    #[test]
+    fn tick_retransmits_pending_window_adjust() {
+        let mut mux = StreamMux::new(TransportConfig {
+            retransmit_after: Duration::from_millis(200),
+            ..TransportConfig::default()
+        });
+        mux.accept_open(StreamOpen {
+            stream_id: 1,
+            target_host: "@dosh-test".to_string(),
+            target_port: 0,
+        })
+        .unwrap();
+
+        let incoming = mux
+            .handle_data(StreamData {
+                stream_id: 1,
+                offset: 0,
+                bytes: b"hello".to_vec(),
+            })
+            .unwrap()
+            .unwrap();
+        assert_eq!(incoming.consumed, 5);
+        assert_eq!(incoming.received_offset, 5);
+        assert_eq!(incoming.window_adjust.kind, PacketKind::StreamWindowAdjust);
+        mux.pending_window_adjusts.get_mut(&1).unwrap().last_sent =
+            Instant::now() - Duration::from_millis(250);
+
+        let retransmits = mux.tick().unwrap();
+        assert_eq!(retransmits.len(), 1);
+        assert_eq!(retransmits[0].kind, PacketKind::StreamWindowAdjust);
+        let adjust: StreamWindowAdjust = decode(&retransmits[0]);
+        assert_eq!(adjust.stream_id, 1);
+        assert_eq!(adjust.received_offset, 5);
+        assert_eq!(adjust.bytes, 5);
+        assert_eq!(mux.pending_window_adjusts[&1].attempts, 2);
+    }
+
+    #[test]
+    fn stream_control_retransmits_expire_after_cap() {
+        let mut mux = StreamMux::new(TransportConfig {
+            retransmit_after: Duration::from_millis(200),
+            ..TransportConfig::default()
+        });
+        mux.accept_open(StreamOpen {
+            stream_id: 1,
+            target_host: "@dosh-test".to_string(),
+            target_port: 0,
+        })
+        .unwrap();
+        mux.close_stream(1).unwrap();
+        let pending = mux.pending_closes.get_mut(&1).unwrap();
+        pending.last_sent = Instant::now() - Duration::from_millis(250);
+        pending.attempts = STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS;
+
+        assert!(mux.tick().unwrap().is_empty());
+        assert!(!mux.pending_closes.contains_key(&1));
+    }
+
+    #[test]
+    fn cleanup_stream_removes_reliable_controls() {
+        let mut mux = StreamMux::new(TransportConfig::default());
+        mux.accept_open(StreamOpen {
+            stream_id: 1,
+            target_host: "@dosh-test".to_string(),
+            target_port: 0,
+        })
+        .unwrap();
+        mux.eof_stream(1).unwrap();
+        mux.handle_data(StreamData {
+            stream_id: 1,
+            offset: 0,
+            bytes: b"hello".to_vec(),
+        })
+        .unwrap()
+        .unwrap();
+
+        mux.retire_stream(1);
+
+        assert!(!mux.pending_eofs.contains_key(&1));
+        assert!(!mux.pending_window_adjusts.contains_key(&1));
+        assert!(!mux.is_open(1));
+    }
+
+    #[test]
+    fn peer_close_after_local_close_stops_close_retransmits() {
+        let mut mux = StreamMux::new(TransportConfig {
+            retransmit_after: Duration::from_millis(200),
+            ..TransportConfig::default()
+        });
+        mux.accept_open(StreamOpen {
+            stream_id: 1,
+            target_host: "@dosh-test".to_string(),
+            target_port: 0,
+        })
+        .unwrap();
+        mux.close_stream(1).unwrap();
+
+        assert!(!mux.handle_close(StreamClose { stream_id: 1 }));
+        assert!(!mux.pending_closes.contains_key(&1));
+        assert!(mux.tick().unwrap().is_empty());
+    }
+
     #[test]
     fn handle_packet_decodes_and_dispatches_stream_packets() {
         let mut mux = StreamMux::new(TransportConfig::default());