Always drain v1 request body before session reuse

Previously the pingora-proxy filter implementations would be responsible
for ensuring that the HTTP/1.1 streams are safe for reuse and contain no
remaining body bytes to drain.

Draining body bytes is now always done as part of finishing the
downstream v1 session before it is attempted for reuse, alongside a
configurable total body drain timeout.

Author: drcaramelsyrup

.bleep

@@ -1 +1 @@
-c8dbb31df87f456082be681af9987172ee8c9954
+2513fe1ff4219fbb218dfd1fe7ffb4bbfb96aa91

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

@@ -111,12 +111,9 @@ impl Session {
     /// This is useful for making streams reusable (in particular for HTTP/1.1) after returning an
     /// error before the whole body has been read.
     pub async fn drain_request_body(&mut self) -> Result<()> {
-        loop {
-            match self.read_request_body().await {
-                Ok(Some(_)) => { /* continue to drain */ }
-                Ok(None) => return Ok(()), // done
-                Err(e) => return Err(e),
-            }
+        match self {
+            Self::H1(s) => s.drain_request_body().await,
+            Self::H2(s) => s.drain_request_body().await,
         }
     }
 
@@ -177,7 +174,7 @@ impl Session {
             Self::H1(mut s) => {
                 // need to flush body due to buffering
                 s.finish_body().await?;
-                Ok(s.reuse().await)
+                s.reuse().await
             }
             Self::H2(mut s) => {
                 s.finish()?;
@@ -225,6 +222,19 @@ impl Session {
         }
     }
 
+    /// Sets the total drain timeout, which will be applied while discarding the
+    /// request body using `drain_request_body`.
+    ///
+    /// For HTTP/1.1, reusing a session requires ensuring that the request body
+    /// is consumed. If the timeout is exceeded, the caller should give up on
+    /// trying to reuse the session.
+    pub fn set_total_drain_timeout(&mut self, timeout: Duration) {
+        match self {
+            Self::H1(s) => s.set_total_drain_timeout(timeout),
+            Self::H2(s) => s.set_total_drain_timeout(timeout),
+        }
+    }
+
     /// Sets the minimum downstream send rate in bytes per second. This
     /// is used to calculate a write timeout in seconds based on the size
     /// of the buffer being written. If a `min_send_rate` is configured it

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

@@ -62,6 +62,8 @@ pub struct HttpSession {
     keepalive_timeout: KeepaliveStatus,
     read_timeout: Option<Duration>,
     write_timeout: Option<Duration>,
+    /// How long to wait to make downstream session reusable, if body needs to be drained.
+    total_drain_timeout: Option<Duration>,
     /// A copy of the response that is already written to the client
     response_written: Option<Box<ResponseHeader>>,
     /// The parsed request header
@@ -106,6 +108,7 @@ impl HttpSession {
             request_header: None,
             read_timeout: None,
             write_timeout: None,
+            total_drain_timeout: None,
             body_bytes_sent: 0,
             body_bytes_read: 0,
             retry_buffer: None,
@@ -399,6 +402,30 @@ impl HttpSession {
         }
     }
 
+    async fn do_drain_request_body(&mut self) -> Result<()> {
+        loop {
+            match self.read_body_bytes().await {
+                Ok(Some(_)) => { /* continue to drain */ }
+                Ok(None) => return Ok(()), // done
+                Err(e) => return Err(e),
+            }
+        }
+    }
+
+    /// Drain the request body. `Ok(())` when there is no (more) body to read.
+    pub async fn drain_request_body(&mut self) -> Result<()> {
+        if self.is_body_done() {
+            return Ok(());
+        }
+        match self.total_drain_timeout {
+            Some(t) => match timeout(t, self.do_drain_request_body()).await {
+                Ok(res) => res,
+                Err(_) => Error::e_explain(ReadTimedout, format!("draining body, timeout: {t:?}")),
+            },
+            None => self.do_drain_request_body().await,
+        }
+    }
+
     /// Whether there is no (more) body need to be read.
     pub fn is_body_done(&mut self) -> bool {
         self.init_body_reader();
@@ -862,6 +889,18 @@ impl HttpSession {
         self.write_timeout = Some(timeout);
     }
 
+    /// Sets the total drain timeout. For HTTP/1.1, reusing a session requires
+    /// ensuring that the request body is consumed. This `timeout` will be used
+    /// to determine how long to wait for the entirety of the downstream request
+    /// body to finish after the upstream response is completed to return the
+    /// session to the reuse pool. If the timeout is exceeded, we will give up
+    /// on trying to reuse the session.
+    ///
+    /// Note that the downstream read timeout still applies between body byte reads.
+    pub fn set_total_drain_timeout(&mut self, timeout: Duration) {
+        self.total_drain_timeout = Some(timeout);
+    }
+
     /// Sets the minimum downstream send rate in bytes per second. This
     /// is used to calculate a write timeout in seconds based on the size
     /// of the buffer being written. If a `min_send_rate` is configured it
@@ -911,19 +950,25 @@ impl HttpSession {
     }
 
     /// Consume `self`, if the connection can be reused, the underlying stream will be returned
-    /// to be fed to the next [`Self::new()`]. The next session can just call [`Self::read_request()`].
+    /// to be fed to the next [`Self::new()`]. This drains any remaining request body if it hasn't
+    /// yet been read and the stream is reusable.
+    ///
+    /// The next session can just call [`Self::read_request()`].
+    ///
     /// If the connection cannot be reused, the underlying stream will be closed and `None` will be
-    /// returned.
-    pub async fn reuse(mut self) -> Option<Stream> {
-        // TODO: this function is unnecessarily slow for keepalive case
-        // because that case does not need async
+    /// returned. If there was an error while draining any remaining request body that error will
+    /// be returned.
+    pub async fn reuse(mut self) -> Result<Option<Stream>> {
         match self.keepalive_timeout {
             KeepaliveStatus::Off => {
                 debug!("HTTP shutdown connection");
                 self.shutdown().await;
-                None
+                Ok(None)
+            }
+            _ => {
+                self.drain_request_body().await?;
+                Ok(Some(self.underlying_stream))
             }
-            _ => Some(self.underlying_stream),
         }
     }
 

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

@@ -24,7 +24,9 @@ use http::uri::PathAndQuery;
 use http::{header, HeaderMap, Response};
 use log::{debug, warn};
 use pingora_http::{RequestHeader, ResponseHeader};
+use pingora_timeout::timeout;
 use std::sync::Arc;
+use std::time::Duration;
 
 use crate::protocols::http::body_buffer::FixedBuffer;
 use crate::protocols::http::date::get_cached_date;
@@ -99,6 +101,8 @@ pub struct HttpSession {
     retry_buffer: Option<FixedBuffer>,
     // digest to record underlying connection info
     digest: Arc<Digest>,
+    // How long to wait when draining (discarding) request body
+    total_drain_timeout: Option<Duration>,
 }
 
 impl HttpSession {
@@ -138,6 +142,7 @@ impl HttpSession {
                 body_sent: 0,
                 retry_buffer: None,
                 digest,
+                total_drain_timeout: None,
             }
         }))
     }
@@ -178,6 +183,40 @@ impl HttpSession {
         Ok(data)
     }
 
+    async fn do_drain_request_body(&mut self) -> Result<()> {
+        loop {
+            match self.read_body_bytes().await {
+                Ok(Some(_)) => { /* continue to drain */ }
+                Ok(None) => return Ok(()), // done
+                Err(e) => return Err(e),
+            }
+        }
+    }
+
+    /// Drain the request body. `Ok(())` when there is no (more) body to read.
+    // NOTE for h2 it may be worth allowing cancellation of the stream via reset.
+    pub async fn drain_request_body(&mut self) -> Result<()> {
+        if self.is_body_done() {
+            return Ok(());
+        }
+        match self.total_drain_timeout {
+            Some(t) => match timeout(t, self.do_drain_request_body()).await {
+                Ok(res) => res,
+                Err(_) => Error::e_explain(
+                    ErrorType::ReadTimedout,
+                    format!("draining body, timeout: {t:?}"),
+                ),
+            },
+            None => self.do_drain_request_body().await,
+        }
+    }
+
+    /// Sets the total drain timeout. This `timeout` will be used while draining
+    /// the request body.
+    pub fn set_total_drain_timeout(&mut self, timeout: Duration) {
+        self.total_drain_timeout = Some(timeout);
+    }
+
     // the write_* don't have timeouts because the actual writing happens on the connection
     // not here.