Make upstream_response_body_filter async

pingora-proxy/Cargo.toml

@@ -56,6 +56,7 @@ futures-util = "0.3"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 serde_yaml = "0.9"
+criterion = "0.5"
 
 [target.'cfg(unix)'.dev-dependencies]
 hyperlocal = "0.9"
@@ -82,6 +83,10 @@ trace = ["pingora-cache/trace"]
 name = "connection_filter"
 required-features = ["connection_filter"]
 
+[[bench]]
+name = "noop_body_filter"
+harness = false
+
 # or locally cargo doc --config "build.rustdocflags='--cfg doc_async_trait'"
 [package.metadata.docs.rs]
 rustdoc-args = ["--cfg", "doc_async_trait"]

pingora-proxy/benches/noop_body_filter.rs

@@ -0,0 +1,145 @@
+// Copyright 2026 Cloudflare, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Per-call cost of `ProxyHttp::upstream_response_body_filter`.
+//!
+//! Run with: `cargo bench -p pingora-proxy --bench noop_body_filter`
+
+use async_trait::async_trait;
+use bytes::Bytes;
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use pingora_core::protocols::l4::stream::Stream as L4Stream;
+use pingora_core::upstreams::peer::HttpPeer;
+use pingora_error::Result;
+use pingora_proxy::{ProxyHttp, Session};
+use std::time::Duration;
+use tokio::io::AsyncWriteExt;
+use tokio::net::{TcpListener, TcpStream};
+
+struct DefaultFilter;
+
+#[async_trait]
+impl ProxyHttp for DefaultFilter {
+    type CTX = ();
+
+    fn new_ctx(&self) -> Self::CTX {}
+
+    async fn upstream_peer(
+        &self,
+        _session: &mut Session,
+        _ctx: &mut Self::CTX,
+    ) -> Result<Box<HttpPeer>> {
+        unreachable!("not used by this benchmark")
+    }
+}
+
+struct OverriddenFilter;
+
+#[async_trait]
+impl ProxyHttp for OverriddenFilter {
+    type CTX = ();
+
+    fn new_ctx(&self) -> Self::CTX {}
+
+    async fn upstream_peer(
+        &self,
+        _session: &mut Session,
+        _ctx: &mut Self::CTX,
+    ) -> Result<Box<HttpPeer>> {
+        unreachable!("not used by this benchmark")
+    }
+
+    async fn upstream_response_body_filter(
+        &self,
+        _session: &mut Session,
+        _body: &mut Option<Bytes>,
+        _end_of_stream: bool,
+        _ctx: &mut Self::CTX,
+    ) -> Result<Option<Duration>> {
+        Ok(None)
+    }
+}
+
+async fn session() -> (Session, TcpStream) {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    let (client, server) = tokio::join!(TcpStream::connect(addr), listener.accept());
+    let mut client = client.unwrap();
+    let (server, _) = server.unwrap();
+
+    client
+        .write_all(b"GET / HTTP/1.1\r\nhost: example.com\r\n\r\n")
+        .await
+        .unwrap();
+    let mut session = Session::new_h1(Box::new(L4Stream::from(server)));
+    session.read_request().await.unwrap();
+    (session, client)
+}
+
+async fn call_filter<F: ProxyHttp<CTX = ()> + Send + Sync>(
+    filter: &F,
+    session: &mut Session,
+    chunk: &Bytes,
+) {
+    let mut body = Some(chunk.clone());
+    let mut ctx = ();
+    black_box(
+        filter
+            .upstream_response_body_filter(
+                black_box(session),
+                black_box(&mut body),
+                false,
+                &mut ctx,
+            )
+            .await
+            .unwrap(),
+    );
+}
+
+fn benchmark(c: &mut Criterion) {
+    let runtime = tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+        .unwrap();
+    let (mut session, _client) = runtime.block_on(session());
+    let chunk = Bytes::from(vec![0u8; 4096]);
+
+    let mut group = c.benchmark_group("upstream_response_body_filter");
+    group.bench_function("default", |b| {
+        b.iter_custom(|iterations| {
+            runtime.block_on(async {
+                let start = std::time::Instant::now();
+                for _ in 0..iterations {
+                    call_filter(&DefaultFilter, &mut session, &chunk).await;
+                }
+                start.elapsed()
+            })
+        })
+    });
+    group.bench_function("trivial_override", |b| {
+        b.iter_custom(|iterations| {
+            runtime.block_on(async {
+                let start = std::time::Instant::now();
+                for _ in 0..iterations {
+                    call_filter(&OverriddenFilter, &mut session, &chunk).await;
+                }
+                start.elapsed()
+            })
+        })
+    });
+    group.finish();
+}
+
+criterion_group!(benches, benchmark);
+criterion_main!(benches);

pingora-proxy/src/lib.rs

@@ -406,9 +406,11 @@ where
                     .await?;
                 None
             }
-            HttpTask::Body(data, eos) | HttpTask::UpgradedBody(data, eos) => self
-                .inner
-                .upstream_response_body_filter(session, data, *eos, ctx)?,
+            HttpTask::Body(data, eos) | HttpTask::UpgradedBody(data, eos) => {
+                self.inner
+                    .upstream_response_body_filter(session, data, *eos, ctx)
+                    .await?
+            }
             HttpTask::Trailer(Some(trailers)) => {
                 self.inner
                     .upstream_response_trailer_filter(session, trailers, ctx)?;
@@ -1838,6 +1840,67 @@ mod tests {
         session
     }
 
+    struct AsyncBodyFilter;
+
+    #[async_trait]
+    impl ProxyHttp for AsyncBodyFilter {
+        type CTX = bool;
+
+        fn new_ctx(&self) -> Self::CTX {
+            false
+        }
+
+        async fn upstream_peer(
+            &self,
+            _session: &mut Session,
+            _ctx: &mut Self::CTX,
+        ) -> Result<Box<HttpPeer>> {
+            unreachable!("not used by this test")
+        }
+
+        async fn upstream_response_body_filter(
+            &self,
+            _session: &mut Session,
+            body: &mut Option<Bytes>,
+            _end_of_stream: bool,
+            fail: &mut Self::CTX,
+        ) -> Result<Option<Duration>> {
+            tokio::task::yield_now().await;
+            if *fail {
+                return Error::e_explain(InternalError, "async body filter failed");
+            }
+            *body = Some(Bytes::from_static(b"filtered"));
+            Ok(Some(Duration::from_millis(7)))
+        }
+    }
+
+    #[tokio::test]
+    async fn upstream_filter_awaits_response_body_filter() {
+        let written = Arc::new(Mutex::new(Vec::new()));
+        let mut session = new_upgrade_request_session(written).await;
+        let proxy = HttpProxy::new(AsyncBodyFilter, Arc::new(ServerConf::default()));
+        let mut task = HttpTask::Body(Some(Bytes::from_static(b"original")), true);
+        let mut fail = false;
+
+        let delay = proxy
+            .upstream_filter(&mut session, &mut task, &mut fail)
+            .await
+            .unwrap();
+
+        assert_eq!(delay, Some(Duration::from_millis(7)));
+        let HttpTask::Body(body, end) = &task else {
+            panic!("expected a body task");
+        };
+        assert_eq!(body.as_deref(), Some(b"filtered".as_slice()));
+        assert!(*end);
+
+        fail = true;
+        assert!(proxy
+            .upstream_filter(&mut session, &mut task, &mut fail)
+            .await
+            .is_err());
+    }
+
     fn upgrade_response_header() -> ResponseHeader {
         let mut header =
             ResponseHeader::build(http::StatusCode::SWITCHING_PROTOCOLS, Some(2)).unwrap();

pingora-proxy/src/proxy_trait.rs

@@ -448,13 +448,21 @@ pub trait ProxyHttp {
     ///
     /// This function will be called every time a piece of response body is received. The `body` is
     /// **not the entire response body**.
-    fn upstream_response_body_filter(
+    ///
+    /// Like [Self::request_body_filter()], the async nature of this function allows executing
+    /// heavy computation logic (e.g. scanning or transforming response content) on offloaded
+    /// threads, or awaiting external services, without blocking the threads who process the
+    /// requests themselves.
+    async fn upstream_response_body_filter(
         &self,
         _session: &mut Session,
         _body: &mut Option<Bytes>,
         _end_of_stream: bool,
         _ctx: &mut Self::CTX,
-    ) -> Result<Option<Duration>> {
+    ) -> Result<Option<Duration>>
+    where
+        Self::CTX: Send + Sync,
+    {
         Ok(None)
     }