Allow modules to receive HttpTask::Done, flush response compression

This change adds a `response_done_filter` hook for downstream HTTP modules
to handle the HttpTask::Done type, which may be sent in lieu of a body +
end flag, for example in certain caching request paths.

In the `ResponseCompression` module this could manifest as un-flushed
body bytes. This change duplicates the `response_filter` logic currently
used for upstream response compression.

Author: drcaramelsyrup

.bleep

@@ -1 +1 @@
-f95698e4cc0825eab85a3c515914eb7c24c7ea83
+0cc5dbe2652b7f8ebe8ac981b3360c2f0966eeee

pingora-core/src/modules/http/compression.rs

@@ -72,6 +72,15 @@ impl HttpModule for ResponseCompression {
         }
         Ok(())
     }
+
+    fn response_done_filter(&mut self) -> Result<Option<Bytes>> {
+        if !self.0.is_enabled() {
+            return Ok(None);
+        }
+        // Flush or finish any remaining encoded bytes upon HTTP response completion
+        // (if it was not already ended in the body filter).
+        Ok(self.0.response_body_filter(None, true))
+    }
 }
 
 /// The builder for HTTP response compression module

pingora-core/src/modules/http/mod.rs

@@ -72,6 +72,10 @@ pub trait HttpModule {
         Ok(None)
     }
 
+    fn response_done_filter(&mut self) -> Result<Option<Bytes>> {
+        Ok(None)
+    }
+
     fn as_any(&self) -> &dyn Any;
     fn as_any_mut(&mut self) -> &mut dyn Any;
 }
@@ -258,6 +262,24 @@ impl HttpModuleCtx {
         }
         Ok(encoded)
     }
+
+    /// Run the `response_done_filter` for all the modules according to their orders.
+    ///
+    /// This filter may be invoked in certain response paths to signal end of response
+    /// if not already done so via trailers or body (with end flag set).
+    ///
+    /// Returns an `Option<Bytes>` which can be used to write additional response body
+    /// bytes. Note, if multiple modules attempt to write body bytes, only the last one
+    /// will be used.
+    pub fn response_done_filter(&mut self) -> Result<Option<Bytes>> {
+        let mut encoded = None;
+        for filter in self.module_ctx.iter_mut() {
+            if let Some(buf) = filter.response_done_filter()? {
+                encoded = Some(buf);
+            }
+        }
+        Ok(encoded)
+    }
 }
 
 #[cfg(test)]

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

@@ -19,7 +19,7 @@
 use super::HttpTask;
 
 use bytes::Bytes;
-use log::warn;
+use log::{debug, warn};
 use pingora_error::{ErrorType, Result};
 use pingora_http::{RequestHeader, ResponseHeader};
 use std::time::Duration;
@@ -248,6 +248,7 @@ impl ResponseCompressionCtx {
                 }
 
                 let action = decide_action(resp, accept_encoding);
+                debug!("compression action: {action:?}");
                 let (encoder, preserve_etag) = match action {
                     Action::Noop => (None, false),
                     Action::Compress(algorithm) => {
@@ -271,9 +272,10 @@ impl ResponseCompressionCtx {
         }
     }
 
-    /// Stream the response body chunks into this ctx. The return value will be the compressed data
+    /// Stream the response body chunks into this ctx. The return value will be the compressed
+    /// data.
     ///
-    /// Return None if the compressed is not enabled
+    /// Return None if compression is not enabled.
     pub fn response_body_filter(&mut self, data: Option<&Bytes>, end: bool) -> Option<Bytes> {
         match &mut self.0 {
             CtxInner::HeaderPhase { .. } => panic!("Wrong phase: HeaderPhase"),

pingora-proxy/src/lib.rs

@@ -431,7 +431,20 @@ impl Session {
                         *task = HttpTask::Body(Some(buf), true);
                     }
                 }
-                _ => { /* Done or Failed */ }
+                HttpTask::Done => {
+                    // `Done` can be sent in certain response paths to mark end
+                    // of response if not already done via trailers or body with
+                    // end flag set.
+                    // If the filter returns body bytes on Done,
+                    // write them into the response.
+                    //
+                    // Note, this will not work if end of stream has already
+                    // been seen or we've written content-length bytes.
+                    if let Some(buf) = self.downstream_modules_ctx.response_done_filter()? {
+                        *task = HttpTask::Body(Some(buf), true);
+                    }
+                }
+                _ => { /* Failed */ }
             }
         }
         self.downstream_session.response_duplex_vec(tasks).await

pingora-proxy/tests/test_upstream.rs

@@ -342,6 +342,53 @@ mod test_cache {
         assert_eq!(res.text().await.unwrap(), "no if headers detected\n");
     }
 
+    #[tokio::test]
+    async fn test_cache_downstream_compression() {
+        init();
+
+        // disable reqwest gzip support to check compression headers and body
+        // otherwise reqwest will decompress and strip the headers
+        let client = reqwest::ClientBuilder::new().gzip(false).build().unwrap();
+        let res = client
+            .get("http://127.0.0.1:6148/unique/test_cache_downstream_compression/no_compression")
+            .header("x-downstream-compression", "1")
+            .header("accept-encoding", "gzip")
+            .send()
+            .await
+            .unwrap();
+        assert_eq!(res.status(), StatusCode::OK);
+        let headers = res.headers();
+        assert_eq!(headers["Content-Encoding"], "gzip");
+        assert_eq!(headers["x-cache-status"], "miss");
+        let body = res.bytes().await.unwrap();
+        assert!(body.len() < 32);
+    }
+
+    #[tokio::test]
+    async fn test_cache_downstream_decompression() {
+        init();
+
+        // disable reqwest gzip support to check compression headers and body
+        // otherwise reqwest will decompress and strip the headers
+        let client = reqwest::ClientBuilder::new().gzip(false).build().unwrap();
+        let res = client
+            .get("http://127.0.0.1:6148/unique/test_cache_downstream_decompression/gzip/index.html")
+            .header("x-downstream-decompression", "1")
+            .header("x-upstream-accept-encoding", "gzip")
+            .send()
+            .await
+            .unwrap();
+
+        assert_eq!(res.status(), StatusCode::OK);
+        let headers = res.headers();
+        // upstream should have received gzip, should decompress for downstream
+        assert_eq!(headers["received-accept-encoding"], "gzip");
+        assert!(headers.get("Content-Encoding").is_none());
+        assert_eq!(headers["x-cache-status"], "miss");
+        let body = res.bytes().await.unwrap();
+        assert_eq!(body, "Hello World!\n");
+    }
+
     #[tokio::test]
     async fn test_network_error_mid_response() {
         init();

pingora-proxy/tests/utils/conf/origin/conf/nginx.conf

@@ -420,6 +420,14 @@ http {
             }
         }
 
+        location /gzip {
+            alias ./html;
+            gzip on;
+            gzip_min_length   0;
+            gzip_types        *;
+            add_header received-accept-encoding $http_accept_encoding;
+        }
+
         location /sleep {
             rewrite_by_lua_block {
                 local sleep_sec = tonumber(ngx.var.http_x_set_sleep) or 1

pingora-proxy/tests/utils/server_utils.rs

@@ -16,7 +16,7 @@
 use super::cert;
 use async_trait::async_trait;
 use clap::Parser;
-use http::header::VARY;
+use http::header::{ACCEPT_ENCODING, VARY};
 use http::HeaderValue;
 use once_cell::sync::Lazy;
 use pingora_cache::cache_control::CacheControl;
@@ -350,6 +350,38 @@ impl ProxyHttp for ExampleProxyCache {
         }
     }
 
+    async fn early_request_filter(
+        &self,
+        session: &mut Session,
+        _ctx: &mut Self::CTX,
+    ) -> Result<()> {
+        if session
+            .req_header()
+            .headers
+            .get("x-downstream-compression")
+            .is_some()
+        {
+            session
+                .downstream_modules_ctx
+                .get_mut::<ResponseCompression>()
+                .unwrap()
+                .adjust_level(6);
+        }
+        if session
+            .req_header()
+            .headers
+            .get("x-downstream-decompression")
+            .is_some()
+        {
+            session
+                .downstream_modules_ctx
+                .get_mut::<ResponseCompression>()
+                .unwrap()
+                .adjust_decompression(true);
+        }
+        Ok(())
+    }
+
     async fn upstream_peer(
         &self,
         session: &mut Session,
@@ -474,6 +506,22 @@ impl ProxyHttp for ExampleProxyCache {
         key.finalize()
     }
 
+    async fn upstream_request_filter(
+        &self,
+        session: &mut Session,
+        upstream_request: &mut RequestHeader,
+        _ctx: &mut Self::CTX,
+    ) -> Result<()> {
+        if let Some(up_accept_encoding) = session
+            .req_header()
+            .headers
+            .get("x-upstream-accept-encoding")
+        {
+            upstream_request.insert_header(&ACCEPT_ENCODING, up_accept_encoding)?;
+        }
+        Ok(())
+    }
+
     fn response_cache_filter(
         &self,
         _session: &Session,