Poll for H2 capacity before sending H2 body.

This change stops data being buffered while the remote peer lacks the
capacity to receive so that the back pressure is propagated.

All h2 body write functions are async now.

Author: eaufavor

.bleep

@@ -1 +1 @@
-81c56d987be33640906fb081facd6af46dc23721
+c57df409e909c786e37f46868146e711025a2b34

pingora-core/src/protocols/http/client.rs

@@ -62,7 +62,7 @@ impl HttpSession {
                 h1.write_body(&data).await?;
                 Ok(())
             }
-            HttpSession::H2(h2) => h2.write_request_body(data, end),
+            HttpSession::H2(h2) => h2.write_request_body(data, end).await,
         }
     }
 

pingora-core/src/protocols/http/server.rs

@@ -143,7 +143,7 @@ impl Session {
                 s.write_body(&data).await?;
                 Ok(())
             }
-            Self::H2(s) => s.write_body(data, end),
+            Self::H2(s) => s.write_body(data, end).await,
         }
     }
 
@@ -175,7 +175,7 @@ impl Session {
     pub async fn response_duplex_vec(&mut self, tasks: Vec<HttpTask>) -> Result<bool> {
         match self {
             Self::H1(s) => s.response_duplex_vec(tasks).await,
-            Self::H2(s) => s.response_duplex_vec(tasks),
+            Self::H2(s) => s.response_duplex_vec(tasks).await,
         }
     }
 

pingora-core/src/protocols/http/v2/client.rs

@@ -122,7 +122,7 @@ impl Http2Session {
     }
 
     /// Write a request body chunk
-    pub fn write_request_body(&mut self, data: Bytes, end: bool) -> Result<()> {
+    pub async fn write_request_body(&mut self, data: Bytes, end: bool) -> Result<()> {
         if self.ended {
             warn!("Try to write request body after end of stream, dropping the extra data");
             return Ok(());
@@ -133,7 +133,9 @@ impl Http2Session {
             .as_mut()
             .expect("Try to write request body before sending request header");
 
-        write_body(body_writer, data, end).map_err(|e| self.handle_err(e))?;
+        super::write_body(body_writer, data, end)
+            .await
+            .map_err(|e| self.handle_err(e))?;
         self.ended = self.ended || end;
         Ok(())
     }
@@ -404,15 +406,6 @@ impl Http2Session {
     }
 }
 
-/// A helper function to write the request body
-pub fn write_body(send_body: &mut SendStream<Bytes>, data: Bytes, end: bool) -> Result<()> {
-    let data_len = data.len();
-    send_body.reserve_capacity(data_len);
-    send_body
-        .send_data(data, end)
-        .or_err(WriteError, "while writing h2 request body")
-}
-
 /* helper functions */
 
 /* Types of errors during h2 header read

pingora-core/src/protocols/http/v2/mod.rs

@@ -14,5 +14,33 @@
 
 //! HTTP/2 implementation
 
+use crate::{Error, ErrorType::*, OrErr, Result};
+use bytes::Bytes;
+use h2::SendStream;
+
 pub mod client;
 pub mod server;
+
+/// A helper function to write the body of h2 streams.
+pub async fn write_body(writer: &mut SendStream<Bytes>, data: Bytes, end: bool) -> Result<()> {
+    let mut remaining = data;
+    // Note: this loop allows to send the data frame even if data is empty in order to signal
+    // END_OF_STREAM to the peer.
+    loop {
+        writer.reserve_capacity(remaining.len());
+        match std::future::poll_fn(|cx| writer.poll_capacity(cx)).await {
+            None => return Error::e_explain(H2Error, "cannot reserve capacity"),
+            Some(ready) => {
+                let n = ready.or_err(H2Error, "while waiting for capacity")?;
+                let remaining_size = remaining.len();
+                let data_to_send = remaining.split_to(std::cmp::min(remaining_size, n));
+                writer
+                    .send_data(data_to_send, remaining.is_empty() && end)
+                    .or_err(WriteError, "while writing h2 request body")?;
+                if remaining.is_empty() {
+                    return Ok(());
+                }
+            }
+        }
+    }
+}

pingora-core/src/protocols/http/v2/server.rs

@@ -232,7 +232,7 @@ impl HttpSession {
     }
 
     /// Write response body to the client. See [Self::write_response_header] for how to use `end`.
-    pub fn write_body(&mut self, data: Bytes, end: bool) -> Result<()> {
+    pub async fn write_body(&mut self, data: Bytes, end: bool) -> Result<()> {
         if self.ended {
             // NOTE: in h1, we also track to see if content-length matches the data
             // We have not tracked that in h2
@@ -246,11 +246,9 @@ impl HttpSession {
             ));
         };
         let data_len = data.len();
-        writer.reserve_capacity(data_len);
-        writer.send_data(data, end).or_err(
-            ErrorType::WriteError,
-            "while writing h2 response body to downstream",
-        )?;
+        super::write_body(writer, data, end)
+            .await
+            .map_err(|e| e.into_down())?;
         self.body_sent += data_len;
         self.ended = self.ended || end;
         Ok(())
@@ -308,7 +306,7 @@ impl HttpSession {
         Ok(())
     }
 
-    pub fn response_duplex_vec(&mut self, tasks: Vec<HttpTask>) -> Result<bool> {
+    pub async fn response_duplex_vec(&mut self, tasks: Vec<HttpTask>) -> Result<bool> {
         let mut end_stream = false;
         for task in tasks.into_iter() {
             end_stream = match task {
@@ -320,7 +318,7 @@ impl HttpSession {
                 HttpTask::Body(data, end) => match data {
                     Some(d) => {
                         if !d.is_empty() {
-                            self.write_body(d, end).map_err(|e| e.into_down())?;
+                            self.write_body(d, end).await.map_err(|e| e.into_down())?;
                         }
                         end
                     }
@@ -567,7 +565,7 @@ mod test {
                 }
 
                 // end: false here to verify finish() closes the stream nicely
-                http.write_body(server_body.into(), false).unwrap();
+                http.write_body(server_body.into(), false).await.unwrap();
                 assert_eq!(http.body_bytes_sent(), 16);
 
                 http.write_trailers(trailers).unwrap();
@@ -639,7 +637,7 @@ mod test {
                     .write_response_header(response_header.clone(), false)
                     .is_ok());
 
-                http.write_body(server_body.into(), false).unwrap();
+                http.write_body(server_body.into(), false).await.unwrap();
                 assert_eq!(http.body_bytes_sent(), 16);
 
                 // 3. Waiting for the reset from the client
@@ -711,7 +709,7 @@ mod test {
                     .write_response_header(response_header.clone(), false)
                     .is_ok());
 
-                http.write_body(server_body.into(), false).unwrap();
+                http.write_body(server_body.into(), false).await.unwrap();
                 assert_eq!(http.body_bytes_sent(), 16);
 
                 // 3. Waiting for the client to close stream.

pingora-proxy/src/proxy_h2.rs

@@ -16,7 +16,7 @@ use super::*;
 use crate::proxy_cache::{range_filter::RangeBodyFilter, ServeFromCache};
 use crate::proxy_common::*;
 use http::{header::CONTENT_LENGTH, Method, StatusCode};
-use pingora_core::protocols::http::v2::client::{write_body, Http2Session};
+use pingora_core::protocols::http::v2::{client::Http2Session, write_body};
 
 // add scheme and authority as required by h2 lib
 fn update_h2_scheme_authority(
@@ -150,7 +150,7 @@ impl<SV> HttpProxy<SV> {
 
         if !send_end_stream && body_empty {
             // send END_STREAM on empty DATA frame
-            match client_session.write_request_body(Bytes::new(), true) {
+            match client_session.write_request_body(Bytes::new(), true).await {
                 Ok(()) => debug!("sent empty DATA frame to h2"),
                 Err(e) => {
                     return (false, Some(e.into_up()));
@@ -547,11 +547,15 @@ impl<SV> HttpProxy<SV> {
 
         if let Some(data) = data {
             debug!("Write {} bytes body to h2 upstream", data.len());
-            write_body(client_body, data, end_of_body).map_err(|e| e.into_up())?;
+            write_body(client_body, data, end_of_body)
+                .await
+                .map_err(|e| e.into_up())?;
         } else {
             debug!("Read downstream body done");
             /* send a standalone END_STREAM flag */
-            write_body(client_body, Bytes::new(), true).map_err(|e| e.into_up())?;
+            write_body(client_body, Bytes::new(), true)
+                .await
+                .map_err(|e| e.into_up())?;
         }
 
         Ok(end_of_body)