Implement server::Builder::max_send_buffer_size

Author: nox

src/proto/connection.rs

@@ -140,6 +140,10 @@ where
         self.inner.streams.set_target_connection_window_size(size);
     }
 
+    pub(crate) fn set_max_send_buffer_size(&mut self, max: usize) {
+        self.inner.streams.set_max_send_buffer_size(max);
+    }
+
     /// Send a new SETTINGS frame with an updated initial window size.
     pub(crate) fn set_initial_window_size(&mut self, size: WindowSize) -> Result<(), UserError> {
         let mut settings = frame::Settings::default();

src/proto/streams/prioritize.rs

@@ -51,6 +51,12 @@ pub(super) struct Prioritize {
 
     /// What `DATA` frame is currently being sent in the codec.
     in_flight_data_frame: InFlightData,
+
+    /// The max send buffer size allowed.
+    max_send_buffer_size: usize,
+
+    /// The current send buffer size.
+    current_send_buffer_size: usize,
 }
 
 #[derive(Debug, Eq, PartialEq)]
@@ -93,9 +99,17 @@ impl Prioritize {
             flow,
             last_opened_id: StreamId::ZERO,
             in_flight_data_frame: InFlightData::Nothing,
+            max_send_buffer_size: usize::MAX,
+            current_send_buffer_size: 0,
         }
     }
 
+    pub fn set_max_send_buffer_size(&mut self, max: usize, store: &mut Store, counts: &mut Counts) {
+        self.max_send_buffer_size = max;
+
+        self.assign_connection_capacity(0, store, counts);
+    }
+
     /// Queue a frame to be sent to the remote
     pub fn queue_frame<B>(
         &mut self,
@@ -175,6 +189,8 @@ impl Prioritize {
             self.try_assign_capacity(stream);
         }
 
+        self.current_send_buffer_size += sz as usize;
+
         if frame.is_end_stream() {
             stream.state.send_close();
             self.reserve_capacity(0, stream, counts);
@@ -350,7 +366,7 @@ impl Prioritize {
         self.flow.assign_capacity(inc);
 
         // Assign newly acquired capacity to streams pending capacity.
-        while self.flow.available() > 0 {
+        while self.available() > 0 {
             let stream = match self.pending_capacity.pop(store) {
                 Some(stream) => stream,
                 None => return,
@@ -373,6 +389,17 @@ impl Prioritize {
         }
     }
 
+    fn available(&self) -> WindowSize {
+        cmp::min(
+            self.flow.available().as_size() as usize,
+            cmp::min(
+                self.max_send_buffer_size
+                    .saturating_sub(self.current_send_buffer_size),
+                WindowSize::MAX as usize,
+            ),
+        ) as WindowSize
+    }
+
     /// Request capacity to send data
     fn try_assign_capacity(&mut self, stream: &mut store::Ptr) {
         let total_requested = stream.requested_send_capacity;
@@ -395,7 +422,8 @@ impl Prioritize {
             additional,
             buffered = stream.buffered_send_data,
             window = stream.send_flow.window_size(),
-            conn = %self.flow.available()
+            conn_window = %self.flow.available(),
+            conn = self.available(),
         );
 
         if additional == 0 {
@@ -413,7 +441,7 @@ impl Prioritize {
         );
 
         // The amount of currently available capacity on the connection
-        let conn_available = self.flow.available().as_size();
+        let conn_available = self.available();
 
         // First check if capacity is immediately available
         if conn_available > 0 {
@@ -509,6 +537,8 @@ impl Prioritize {
 
                     // Because, always try to reclaim...
                     self.reclaim_frame(buffer, store, dst);
+
+                    self.assign_connection_capacity(0, store, counts);
                 }
                 None => {
                     // Try to flush the codec.
@@ -630,6 +660,8 @@ impl Prioritize {
             tracing::trace!(?frame, "dropping");
         }
 
+        self.current_send_buffer_size -= stream.buffered_send_data;
+
         stream.buffered_send_data = 0;
         stream.requested_send_capacity = 0;
         if let InFlightData::DataFrame(key) = self.in_flight_data_frame {
@@ -736,6 +768,8 @@ impl Prioritize {
                             tracing::trace_span!("updating stream flow").in_scope(|| {
                                 stream.send_flow.send_data(len);
 
+                                self.current_send_buffer_size -= len as usize;
+
                                 // Decrement the stream's buffered data counter
                                 debug_assert!(stream.buffered_send_data >= len as usize);
                                 stream.buffered_send_data -= len as usize;

src/proto/streams/send.rs

@@ -118,6 +118,10 @@ impl Send {
         Ok(())
     }
 
+    pub fn set_max_send_buffer_size(&mut self, max: usize, store: &mut Store, counts: &mut Counts) {
+        self.prioritize.set_max_send_buffer_size(max, store, counts);
+    }
+
     pub fn send_headers<B>(
         &mut self,
         frame: frame::Headers,

src/proto/streams/streams.rs

@@ -127,6 +127,15 @@ where
             .set_target_connection_window(size, &mut me.actions.task)
     }
 
+    pub fn set_max_send_buffer_size(&mut self, max: usize) {
+        let mut me = self.inner.lock().unwrap();
+        let me = &mut *me;
+
+        me.actions
+            .send
+            .set_max_send_buffer_size(max, &mut me.store, &mut me.counts);
+    }
+
     pub fn next_incoming(&mut self) -> Option<StreamRef<B>> {
         let mut me = self.inner.lock().unwrap();
         let me = &mut *me;

src/server.rs

@@ -245,6 +245,9 @@ pub struct Builder {
 
     /// Initial target window size for new connections.
     initial_target_connection_window_size: Option<u32>,
+
+    /// Max send buffer size.
+    max_send_buffer_size: Option<usize>,
 }
 
 /// Send a response back to the client
@@ -451,6 +454,11 @@ where
         self.connection.set_target_window_size(size);
     }
 
+    /// Sets the max send buffer size for the whole connection.
+    pub fn set_max_send_buffer_size(&mut self, max: usize) {
+        self.connection.set_max_send_buffer_size(max);
+    }
+
     /// Set a new `INITIAL_WINDOW_SIZE` setting (in octets) for stream-level
     /// flow control for received data.
     ///
@@ -620,6 +628,7 @@ impl Builder {
             reset_stream_max: proto::DEFAULT_RESET_STREAM_MAX,
             settings: Settings::default(),
             initial_target_connection_window_size: None,
+            max_send_buffer_size: None,
         }
     }
 
@@ -763,6 +772,35 @@ impl Builder {
         self
     }
 
+    /// Sets the max size of the send buffer.
+    ///
+    /// This setting is also used to limit the maximum amount of data
+    /// buffered to be sent.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use tokio::io::{AsyncRead, AsyncWrite};
+    /// # use h2::server::*;
+    /// #
+    /// # fn doc<T: AsyncRead + AsyncWrite + Unpin>(my_io: T)
+    /// # -> Handshake<T>
+    /// # {
+    /// // `server_fut` is a future representing the completion of the HTTP/2
+    /// // handshake.
+    /// let server_fut = Builder::new()
+    ///     .max_send_buffer_size(16 * 1024)
+    ///     .handshake(my_io);
+    /// # server_fut
+    /// # }
+    /// #
+    /// # pub fn main() {}
+    /// ```
+    pub fn max_send_buffer_size(&mut self, max: usize) -> &mut Self {
+        self.max_send_buffer_size = Some(max);
+        self
+    }
+
     /// Sets the maximum number of concurrent streams.
     ///
     /// The maximum concurrent streams setting only controls the maximum number
@@ -1280,6 +1318,9 @@ where
             if let Some(sz) = self.builder.initial_target_connection_window_size {
                 c.set_target_window_size(sz);
             }
+            if let Some(sz) = self.builder.max_send_buffer_size {
+                c.set_max_send_buffer_size(sz);
+            }
             Ok(c)
         })
     }