Split stream writes into packet chunks

src/transport.rs

@@ -34,6 +34,7 @@ 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-";
+pub const MAX_STREAM_DATA_BYTES: usize = 60 * 1024;
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct TransportConfig {
@@ -362,29 +363,35 @@ impl StreamMux {
         if bytes.is_empty() {
             return Ok(Vec::new());
         }
-        if !self.opened_streams.contains(&stream_id) {
-            if self.pending_opens.contains_key(&stream_id) {
+        let chunks = split_stream_data_bytes(bytes, self.config.initial_window);
+        let mut out = Vec::new();
+        if !self.opened_streams.contains(&stream_id) && !self.pending_opens.contains_key(&stream_id)
+        {
+            bail!("stream {stream_id} is not open");
+        }
+        for chunk in chunks {
+            if !self.opened_streams.contains(&stream_id) {
                 self.pending_data
                     .entry(stream_id)
                     .or_default()
-                    .push_back(bytes);
-                return Ok(Vec::new());
+                    .push_back(chunk);
+                continue;
             }
-            bail!("stream {stream_id} is not open");
+            if self.send_credit.get(&stream_id).copied().unwrap_or(0) < chunk.len()
+                || self
+                    .pending_data
+                    .get(&stream_id)
+                    .is_some_and(|pending| !pending.is_empty())
+            {
+                self.pending_data
+                    .entry(stream_id)
+                    .or_default()
+                    .push_back(chunk);
+                continue;
+            }
+            out.push(self.send_data_now(stream_id, chunk)?);
         }
-        if self.send_credit.get(&stream_id).copied().unwrap_or(0) < bytes.len()
-            || self
-                .pending_data
-                .get(&stream_id)
-                .is_some_and(|pending| !pending.is_empty())
-        {
-            self.pending_data
-                .entry(stream_id)
-                .or_default()
-                .push_back(bytes);
-            return Ok(Vec::new());
-        }
-        Ok(vec![self.send_data_now(stream_id, bytes)?])
+        Ok(out)
     }
 
     pub fn handle_data(&mut self, data: StreamData) -> Result<Option<IncomingStreamData>> {
@@ -600,11 +607,30 @@ impl StreamMux {
 
     fn flush_pending_data(&mut self, stream_id: u64) -> Result<Vec<OutgoingStreamPacket>> {
         let mut out = Vec::new();
+        let chunk_limit = stream_data_chunk_limit(self.config.initial_window);
         while let Some(front_len) = self
             .pending_data
             .get(&stream_id)
             .and_then(|pending| pending.front().map(Vec::len))
         {
+            if front_len > chunk_limit {
+                let Some(bytes) = self
+                    .pending_data
+                    .get_mut(&stream_id)
+                    .and_then(VecDeque::pop_front)
+                else {
+                    break;
+                };
+                if let Some(pending) = self.pending_data.get_mut(&stream_id) {
+                    for chunk in split_stream_data_bytes(bytes, self.config.initial_window)
+                        .into_iter()
+                        .rev()
+                    {
+                        pending.push_front(chunk);
+                    }
+                }
+                continue;
+            }
             if self.send_credit.get(&stream_id).copied().unwrap_or(0) < front_len {
                 break;
             }
@@ -796,6 +822,21 @@ pub fn stream_data_within_receive_window(
     end <= expected_offset.saturating_add(window)
 }
 
+pub fn stream_data_chunk_limit(receive_window: usize) -> usize {
+    MAX_STREAM_DATA_BYTES.min(receive_window.max(1))
+}
+
+pub fn split_stream_data_bytes(bytes: Vec<u8>, receive_window: usize) -> Vec<Vec<u8>> {
+    if bytes.is_empty() {
+        return Vec::new();
+    }
+    let limit = stream_data_chunk_limit(receive_window);
+    if bytes.len() <= limit {
+        return vec![bytes];
+    }
+    bytes.chunks(limit).map(Vec::from).collect()
+}
+
 pub struct DoshTransport {
     socket: Arc<UdpSocket>,
     role: SessionRole,
@@ -1328,6 +1369,63 @@ mod tests {
         assert_eq!(data.bytes, b"!");
     }
 
+    #[test]
+    fn send_data_splits_large_writes_into_packet_sized_chunks() {
+        let mut mux = StreamMux::new(TransportConfig::default());
+        mux.open_stream(1, "@dosh-test", 0).unwrap();
+        mux.handle_open_ok(StreamOpenOk { stream_id: 1 })
+            .unwrap()
+            .unwrap();
+
+        let bytes = vec![7; MAX_STREAM_DATA_BYTES * 2 + 13];
+        let sent = mux.send_data(1, bytes).unwrap();
+
+        assert_eq!(sent.len(), 3);
+        let first: StreamData = decode(&sent[0]);
+        let second: StreamData = decode(&sent[1]);
+        let third: StreamData = decode(&sent[2]);
+        assert_eq!(first.offset, 0);
+        assert_eq!(first.bytes.len(), MAX_STREAM_DATA_BYTES);
+        assert_eq!(second.offset, MAX_STREAM_DATA_BYTES as u64);
+        assert_eq!(second.bytes.len(), MAX_STREAM_DATA_BYTES);
+        assert_eq!(third.offset, (MAX_STREAM_DATA_BYTES * 2) as u64);
+        assert_eq!(third.bytes.len(), 13);
+    }
+
+    #[test]
+    fn queued_large_write_flushes_in_window_sized_chunks_after_open() {
+        let mut mux = StreamMux::new(TransportConfig {
+            initial_window: 5,
+            ..TransportConfig::default()
+        });
+        mux.open_stream(1, "@dosh-test", 0).unwrap();
+        assert!(mux.send_data(1, b"hello!".to_vec()).unwrap().is_empty());
+
+        let flushed = mux
+            .handle_open_ok(StreamOpenOk { stream_id: 1 })
+            .unwrap()
+            .unwrap();
+
+        assert_eq!(flushed.len(), 1);
+        let first: StreamData = decode(&flushed[0]);
+        assert_eq!(first.offset, 0);
+        assert_eq!(first.bytes, b"hello");
+        assert_eq!(mux.pending_bytes(1), 1);
+
+        let flushed = mux
+            .handle_window_adjust(StreamWindowAdjust {
+                stream_id: 1,
+                received_offset: 5,
+                bytes: 5,
+            })
+            .unwrap()
+            .unwrap();
+        assert_eq!(flushed.len(), 1);
+        let second: StreamData = decode(&flushed[0]);
+        assert_eq!(second.offset, 5);
+        assert_eq!(second.bytes, b"!");
+    }
+
     #[test]
     fn stream_open_ack_lowers_retransmit_timeout_from_observed_rtt() {
         let mut mux = StreamMux::new(TransportConfig {