fix(s2n-quic-dc): don't restrict TCP writes to limits for UDP (#2534)

Author: WesleyRosenblum

dc/s2n-quic-dc/src/msg/segment.rs

@@ -1,6 +1,7 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+use crate::stream::TransportFeatures;
 use arrayvec::ArrayVec;
 use core::ops::Deref;
 use s2n_quic_core::{ensure, inet::ExplicitCongestionNotification};
@@ -12,7 +13,7 @@ use std::io::IoSlice;
 /// > #define UIO_FASTIOV 8
 pub const MAX_COUNT: usize = if cfg!(target_os = "linux") { 8 } else { 1 };
 
-/// The maximum payload allowed in sendmsg calls
+/// The maximum payload allowed in sendmsg calls using UDP
 ///
 /// From <https://github.com/torvalds/linux/blob/8cd26fd90c1ad7acdcfb9f69ca99d13aa7b24561/net/ipv4/ip_output.c#L987-L995>
 /// > Linux enforces a u16::MAX - IP_HEADER_LEN - UDP_HEADER_LEN
@@ -36,13 +37,13 @@ impl<'a> Deref for Batch<'a> {
 
 impl<'a> Batch<'a> {
     #[inline]
-    pub fn new<Q>(queue: Q) -> Self
+    pub fn new<Q>(queue: Q, features: &TransportFeatures) -> Self
     where
         Q: IntoIterator<Item = (ExplicitCongestionNotification, &'a [u8])>,
     {
         // this value is replaced by the first segment
         let mut ecn = ExplicitCongestionNotification::Ect0;
-        let mut total_len = 0u16;
+        let mut total_len = 0u32;
         let mut segments = Segments::new();
 
         for segment in queue {
@@ -53,12 +54,12 @@ impl<'a> Batch<'a> {
             );
             let packet_len = packet_len as u16;
 
-            // make sure the packet fits in u16::MAX
-            let Some(new_total_len) = total_len.checked_add(packet_len) else {
-                break;
-            };
-            // make sure we don't exceed the max allowed payload size
-            ensure!(new_total_len < MAX_TOTAL, break);
+            let new_total_len = total_len + packet_len as u32;
+
+            if !features.is_stream() {
+                // make sure we don't exceed the max allowed payload size
+                ensure!(new_total_len < MAX_TOTAL as u32, break);
+            }
 
             // track if the current segment is undersized from the previous
             let mut undersized_segment = false;

dc/s2n-quic-dc/src/stream/send/application.rs

@@ -26,6 +26,7 @@ mod builder;
 pub mod state;
 pub mod transmission;
 
+use crate::stream::socket::Application;
 pub use builder::Builder;
 
 pub struct Writer<Sub: event::Subscriber>(Box<Inner<Sub>>);
@@ -176,16 +177,18 @@ where
 
         let path = self.shared.sender.path.load();
 
-        // clamp the flow request based on the path state
-        request.clamp(path.max_flow_credits(max_header_len, max_segments));
+        let features = self.sockets.features();
+
+        if !features.is_flow_controlled() {
+            // clamp the flow request based on the path state
+            request.clamp(path.max_flow_credits(max_header_len, max_segments));
+        }
 
         // acquire flow credits from the worker
-        let credits = ready!(self.shared.sender.flow.poll_acquire(cx, request))?;
+        let credits = ready!(self.shared.sender.flow.poll_acquire(cx, request, &features))?;
 
         trace!(?credits);
 
-        let features = self.sockets.features();
-
         let mut batch = if features.is_reliable() {
             // the protocol does recovery for us so no need to track the transmissions
             None
@@ -219,6 +222,7 @@ where
                             self.shared.credentials(),
                             &clock::Cached::new(&self.shared.clock),
                             message,
+                            &self.sockets.features(),
                         )
                     },
                     |sealer| {

dc/s2n-quic-dc/src/stream/send/application/state.rs

@@ -8,6 +8,7 @@ use crate::{
     stream::{
         packet_number,
         send::{application::transmission, error::Error, flow, path, probes},
+        TransportFeatures,
     },
 };
 use bytes::buf::UninitSlice;
@@ -47,6 +48,7 @@ impl State {
         credentials: &Credentials,
         clock: &Clk,
         message: &mut M,
+        features: &TransportFeatures,
     ) -> Result<(), Error>
     where
         E: seal::Application,
@@ -71,13 +73,22 @@ impl State {
         let max_header_len = self.max_header_len();
 
         let mut total_payload_len = 0;
+        let max_record_size = if features.is_stream() {
+            // If the underlying transport is stream based, it will perform its own packetization
+            // based on the MTU determined at that layer. Therefore, we do not need to restrict
+            // writes to the probed max datagram size, and can instead use a larger value, in this
+            // case 2^14, based on the TLS max record size.
+            1 << 14
+        } else {
+            path.max_datagram_size
+        };
 
         loop {
             let packet_number = packet_number.next()?;
 
             let buffer_len = {
                 let estimated_len = reader.buffered_len() + max_header_len;
-                (path.max_datagram_size as usize).min(estimated_len)
+                (max_record_size as usize).min(estimated_len)
             };
 
             message.push(buffer_len, |buffer| {
@@ -108,7 +119,7 @@ impl State {
                 );
 
                 // buffer is clamped to u16::MAX so this is safe to cast without loss
-                let _: u16 = path.max_datagram_size;
+                let _: u16 = max_record_size;
                 let packet_len = packet_len as u16;
                 let payload_len = reader.consumed_len() as u16;
                 total_payload_len += payload_len as usize;

dc/s2n-quic-dc/src/stream/send/flow/blocking.rs

@@ -2,9 +2,12 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use super::Credits;
-use crate::stream::send::{
-    error::{self, Error},
-    flow,
+use crate::stream::{
+    send::{
+        error::{self, Error},
+        flow,
+    },
+    TransportFeatures,
 };
 use s2n_quic_core::{ensure, varint::VarInt};
 use std::sync::{Condvar, Mutex};
@@ -65,7 +68,11 @@ impl State {
 
     /// Called by the application to acquire flow credits
     #[inline]
-    pub fn acquire(&self, mut request: flow::Request) -> Result<Credits, Error> {
+    pub fn acquire(
+        &self,
+        mut request: flow::Request,
+        features: &TransportFeatures,
+    ) -> Result<Credits, Error> {
         let mut guard = self
             .state
             .lock()
@@ -104,8 +111,10 @@ impl State {
                 continue;
             };
 
-            // clamp the request to the flow credits we have
-            request.clamp(flow_credits);
+            if !features.is_flow_controlled() {
+                // clamp the request to the flow credits we have
+                request.clamp(flow_credits);
+            }
 
             // update the stream offset with the given request
             guard.stream_offset = current_offset
@@ -154,6 +163,7 @@ mod tests {
         let total = AtomicU64::new(0);
         let workers = 5;
         let worker_counts = Vec::from_iter((0..workers).map(|_| AtomicU64::new(0)));
+        let features = TransportFeatures::UDP;
 
         thread::scope(|s| {
             let total = &total;
@@ -178,7 +188,7 @@ mod tests {
                         };
                         request.clamp(path_info.max_flow_credits(max_header_len, max_segments));
 
-                        let Ok(credits) = state.acquire(request) else {
+                        let Ok(credits) = state.acquire(request, &features) else {
                             break;
                         };
 

dc/s2n-quic-dc/src/stream/send/flow/non_blocking.rs

@@ -2,9 +2,12 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use super::Credits;
-use crate::stream::send::{
-    error::{self, Error},
-    flow,
+use crate::stream::{
+    send::{
+        error::{self, Error},
+        flow,
+    },
+    TransportFeatures,
 };
 use atomic_waker::AtomicWaker;
 use core::{
@@ -96,8 +99,12 @@ impl State {
 
     /// Called by the application to acquire flow credits
     #[inline]
-    pub async fn acquire(&self, request: flow::Request) -> Result<Credits, Error> {
-        core::future::poll_fn(|cx| self.poll_acquire(cx, request)).await
+    pub async fn acquire(
+        &self,
+        request: flow::Request,
+        features: &TransportFeatures,
+    ) -> Result<Credits, Error> {
+        core::future::poll_fn(|cx| self.poll_acquire(cx, request, features)).await
     }
 
     /// Called by the application to acquire flow credits
@@ -106,6 +113,7 @@ impl State {
         &self,
         cx: &mut Context,
         mut request: flow::Request,
+        features: &TransportFeatures,
     ) -> Poll<Result<Credits, Error>> {
         let mut current_offset = self.acquire_offset(&request)?;
 
@@ -138,8 +146,10 @@ impl State {
                 continue;
             };
 
-            // clamp the request to the flow credits we have
-            request.clamp(flow_credits);
+            if !features.is_flow_controlled() {
+                // clamp the request to the flow credits we have
+                request.clamp(flow_credits);
+            }
 
             let mut new_offset = (current_offset & OFFSET_MASK)
                 .checked_add(request.len as u64)
@@ -223,6 +233,7 @@ mod tests {
         // TODO support more than one Waker via intrusive list or something
         let workers = 1;
         let worker_counts = Vec::from_iter((0..workers).map(|_| Arc::new(AtomicU64::new(0))));
+        let features = TransportFeatures::UDP;
 
         let mut tasks = tokio::task::JoinSet::new();
 
@@ -246,7 +257,7 @@ mod tests {
                     };
                     request.clamp(path_info.max_flow_credits(max_header_len, max_segments));
 
-                    let Ok(credits) = state.acquire(request).await else {
+                    let Ok(credits) = state.acquire(request, &features).await else {
                         break;
                     };
 

dc/s2n-quic-dc/src/stream/send/queue.rs

@@ -250,11 +250,14 @@ impl Queue {
     {
         while !self.segments.is_empty() {
             let mut provided_len = 0;
-            let segments = segment::Batch::new(self.segments.iter().map(|v| {
-                let slice = v.as_slice();
-                provided_len += slice.len();
-                (v.ecn, v.as_slice())
-            }));
+            let segments = segment::Batch::new(
+                self.segments.iter().map(|v| {
+                    let slice = v.as_slice();
+                    provided_len += slice.len();
+                    (v.ecn, v.as_slice())
+                }),
+                &socket.features(),
+            );
 
             let ecn = segments.ecn();
 
@@ -369,6 +372,7 @@ impl Queue {
                         (v.ecn, slice)
                     })
                     .take(max_segments),
+                &socket.features(),
             );
 
             let ecn = segments.ecn();

dc/s2n-quic-dc/src/stream/send/worker.rs

@@ -416,8 +416,10 @@ where
             ready!(self.pacer.poll_pacing(cx, &self.shared.clock));
 
             // construct all of the segments we're going to send in this batch
-            let segments =
-                msg::segment::Batch::new(self.sender.transmit_queue_iter(clock).take(max_segments));
+            let segments = msg::segment::Batch::new(
+                self.sender.transmit_queue_iter(clock).take(max_segments),
+                &self.socket.features(),
+            );
 
             let ecn = segments.ecn();
             let res = ready!(self.socket.poll_send(cx, &addr, ecn, &segments));