core: allow a session to widen its retry buffer

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

@@ -53,6 +53,24 @@ impl FixedBuffer {
     pub fn is_empty(&self) -> bool {
         self.buffer.len() == 0
     }
+    /// Raise (or lower) the capacity of a buffer that has not started filling.
+    ///
+    /// Only meaningful before any body byte has been written: once data is in,
+    /// changing the ceiling cannot un-truncate what was already dropped, so
+    /// this leaves a non-empty buffer alone and reports `false`.
+    ///
+    /// Exists so a caller that needs the whole body retained — a proxy holding
+    /// a request back for content inspection, say — can widen the ceiling for
+    /// the requests it cares about, instead of every deployment paying a larger
+    /// default.
+    pub fn set_capacity(&mut self, capacity: usize) -> bool {
+        if !self.buffer.is_empty() || self.truncated {
+            return false;
+        }
+        self.capacity = capacity;
+        true
+    }
+
     pub fn is_truncated(&self) -> bool {
         self.truncated
     }
@@ -82,4 +100,33 @@ mod tests {
         assert_eq!(buffer.buffer.capacity(), 0);
         assert!(buffer.get_buffer().is_none());
     }
+
+    #[test]
+    fn raised_capacity_retains_a_body_the_default_would_truncate() {
+        let mut buf = FixedBuffer::new(4);
+        assert!(buf.set_capacity(16));
+
+        buf.write_to_buffer(&Bytes::from_static(b"0123456789"));
+        assert!(!buf.is_truncated());
+        assert_eq!(buf.get_buffer().unwrap(), Bytes::from_static(b"0123456789"));
+    }
+
+    /// Widening after bytes have landed cannot recover anything already
+    /// dropped, so the change is refused rather than silently reporting a whole
+    /// body that is actually a prefix.
+    #[test]
+    fn capacity_change_is_refused_once_the_buffer_has_data() {
+        let mut buf = FixedBuffer::new(8);
+        buf.write_to_buffer(&Bytes::from_static(b"abc"));
+        assert!(!buf.set_capacity(1024));
+    }
+
+    #[test]
+    fn capacity_change_is_refused_after_truncation() {
+        let mut buf = FixedBuffer::new(2);
+        buf.write_to_buffer(&Bytes::from_static(b"abcdef"));
+        assert!(buf.is_truncated());
+        assert!(!buf.set_capacity(1024));
+        assert!(buf.is_truncated(), "still truncated after a refused resize");
+    }
 }

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

@@ -118,6 +118,13 @@ pub trait Session: Send + Sync + Unpin + 'static {
 
     fn enable_retry_buffering(&mut self);
 
+    /// Widen the retry buffer for this session. Defaulted so existing custom
+    /// sessions keep compiling; `false` means the request cannot be fully
+    /// retained and the caller should not rely on replaying its body.
+    fn enable_retry_buffering_with_limit(&mut self, _limit: usize) -> bool {
+        false
+    }
+
     fn retry_buffer_truncated(&self) -> bool;
 
     fn get_retry_buffer(&self) -> Option<Bytes>;

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

@@ -672,6 +672,23 @@ impl Session {
         }
     }
 
+    /// Enable retry buffering with a capacity other than the 64 KB default.
+    ///
+    /// Returns whether the limit took effect — `false` if the buffer has
+    /// already started filling (or the session cannot retain a body), in which
+    /// case the body must not be assumed replayable. Callers that need the
+    /// whole body available after reading it — inspecting a request before
+    /// letting it reach the upstream, for instance — should check this and fall
+    /// back to streaming when it is `false`.
+    pub fn enable_retry_buffering_with_limit(&mut self, limit: usize) -> bool {
+        match self {
+            Self::H1(s) => s.enable_retry_buffering_with_limit(limit),
+            Self::H2(s) => s.enable_retry_buffering_with_limit(limit),
+            Self::Subrequest(s) => s.enable_retry_buffering_with_limit(limit),
+            Self::Custom(s) => s.enable_retry_buffering_with_limit(limit),
+        }
+    }
+
     pub fn get_retry_buffer(&self) -> Option<Bytes> {
         match self {
             Self::H1(s) => s.get_retry_buffer(),

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

@@ -587,6 +587,17 @@ impl HttpSession {
         }
     }
 
+    /// See [`crate::protocols::http::v1::server::HttpSession::enable_retry_buffering_with_limit`].
+    pub fn enable_retry_buffering_with_limit(&mut self, limit: usize) -> bool {
+        match self.retry_buffer.as_mut() {
+            Some(buffer) => buffer.set_capacity(limit),
+            None => {
+                self.retry_buffer = Some(FixedBuffer::new(limit));
+                true
+            }
+        }
+    }
+
     pub fn get_retry_buffer(&self) -> Option<Bytes> {
         self.retry_buffer.as_ref().and_then(|b| {
             if b.is_truncated() {

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

@@ -1158,6 +1158,26 @@ impl HttpSession {
         }
     }
 
+    /// Enable retry buffering with a capacity other than [`BODY_BUF_LIMIT`],
+    /// creating the buffer if it does not exist yet or widening one that has
+    /// not started filling. Returns whether the limit is now in effect.
+    ///
+    /// The default stays 64 KB: raising it globally would make every request
+    /// body up to the new ceiling be retained, since the proxy enables retry
+    /// buffering unconditionally. This lets a caller opt specific requests into
+    /// full retention — a proxy that must inspect a body before releasing it
+    /// upstream needs the whole thing, and a truncated buffer cannot be
+    /// replayed at all.
+    pub fn enable_retry_buffering_with_limit(&mut self, limit: usize) -> bool {
+        match self.retry_buffer.as_mut() {
+            Some(buffer) => buffer.set_capacity(limit),
+            None => {
+                self.retry_buffer = Some(FixedBuffer::new(limit));
+                true
+            }
+        }
+    }
+
     pub fn get_retry_buffer(&self) -> Option<Bytes> {
         self.retry_buffer.as_ref().and_then(|b| {
             if b.is_truncated() {

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

@@ -752,6 +752,17 @@ impl HttpSession {
         }
     }
 
+    /// See [`super::v1::server::HttpSession::enable_retry_buffering_with_limit`].
+    pub fn enable_retry_buffering_with_limit(&mut self, limit: usize) -> bool {
+        match self.retry_buffer.as_mut() {
+            Some(buffer) => buffer.set_capacity(limit),
+            None => {
+                self.retry_buffer = Some(FixedBuffer::new(limit));
+                true
+            }
+        }
+    }
+
     pub fn get_retry_buffer(&self) -> Option<Bytes> {
         self.retry_buffer.as_ref().and_then(|b| {
             if b.is_truncated() {