From 4c5165c9be19cba140fdfc5287f10c7d5008fb0d Mon Sep 17 00:00:00 2001 From: rxdiscovery <> Date: Sun, 21 Jun 2026 21:27:41 +0100 Subject: [PATCH 1/5] feat: make response filter functions asynchronous in pingora Changed `upstream_response_body_filter`, `upstream_response_trailer_filter`, and `response_body_filter` to be `async` functions. --- .bleep | 2 +- pingora-proxy/examples/modify_response.rs | 2 +- pingora-proxy/src/lib.rs | 11 +++++++---- pingora-proxy/src/proxy_cache.rs | 1 + pingora-proxy/src/proxy_custom.rs | 6 ++++-- pingora-proxy/src/proxy_h1.rs | 6 ++++-- pingora-proxy/src/proxy_h2.rs | 3 ++- pingora-proxy/src/proxy_trait.rs | 6 +++--- 8 files changed, 23 insertions(+), 14 deletions(-) diff --git a/.bleep b/.bleep index 2eff8ca09..bd08489ff 100644 --- a/.bleep +++ b/.bleep @@ -1 +1 @@ -3615b9f6e16a0a0c2d7523c8e8b812d1681dcf10 \ No newline at end of file +69857cc99cbb2e232b30e3b61fa060dafa1dfb7b \ No newline at end of file diff --git a/pingora-proxy/examples/modify_response.rs b/pingora-proxy/examples/modify_response.rs index ea10f03f8..6da9701ec 100644 --- a/pingora-proxy/examples/modify_response.rs +++ b/pingora-proxy/examples/modify_response.rs @@ -84,7 +84,7 @@ impl ProxyHttp for Json2Yaml { Ok(()) } - fn response_body_filter( + async fn response_body_filter( &self, _session: &mut Session, body: &mut Option, diff --git a/pingora-proxy/src/lib.rs b/pingora-proxy/src/lib.rs index c6b256bc7..261674dc3 100644 --- a/pingora-proxy/src/lib.rs +++ b/pingora-proxy/src/lib.rs @@ -501,12 +501,15 @@ 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)?; + .upstream_response_trailer_filter(session, trailers, ctx) + .await?; None } _ => { diff --git a/pingora-proxy/src/proxy_cache.rs b/pingora-proxy/src/proxy_cache.rs index 0d032a691..162dcb794 100644 --- a/pingora-proxy/src/proxy_cache.rs +++ b/pingora-proxy/src/proxy_cache.rs @@ -470,6 +470,7 @@ where match self .inner .response_body_filter(session, &mut body, end, ctx) + .await { Ok(Some(duration)) => { trace!("delaying response for {duration:?}"); diff --git a/pingora-proxy/src/proxy_custom.rs b/pingora-proxy/src/proxy_custom.rs index f664e14c6..f21974b8c 100644 --- a/pingora-proxy/src/proxy_custom.rs +++ b/pingora-proxy/src/proxy_custom.rs @@ -823,7 +823,8 @@ where let mut data = range_body_filter.filter_body(data); if let Some(duration) = self .inner - .response_body_filter(session, &mut data, eos, ctx)? + .response_body_filter(session, &mut data, eos, ctx) + .await? { trace!("delaying response for {duration:?}"); time::sleep(duration).await; @@ -840,7 +841,8 @@ where // range body filter doesn't apply to upgraded body if let Some(duration) = self .inner - .response_body_filter(session, &mut data, eos, ctx)? + .response_body_filter(session, &mut data, eos, ctx) + .await? { trace!("delaying upgraded response for {duration:?}"); time::sleep(duration).await; diff --git a/pingora-proxy/src/proxy_h1.rs b/pingora-proxy/src/proxy_h1.rs index 7d0881f9d..a332f8287 100644 --- a/pingora-proxy/src/proxy_h1.rs +++ b/pingora-proxy/src/proxy_h1.rs @@ -973,7 +973,8 @@ where let mut data = range_body_filter.filter_body(data); if let Some(duration) = self .inner - .response_body_filter(session, &mut data, end, ctx)? + .response_body_filter(session, &mut data, end, ctx) + .await? { trace!("delaying downstream response for {:?}", duration); time::sleep(duration).await; @@ -991,7 +992,8 @@ where // range doesn't apply to upgraded body if let Some(duration) = self .inner - .response_body_filter(session, &mut data, end, ctx)? + .response_body_filter(session, &mut data, end, ctx) + .await? { trace!("delaying downstream upgraded response for {:?}", duration); time::sleep(duration).await; diff --git a/pingora-proxy/src/proxy_h2.rs b/pingora-proxy/src/proxy_h2.rs index 3842d19c6..0b13d8a44 100644 --- a/pingora-proxy/src/proxy_h2.rs +++ b/pingora-proxy/src/proxy_h2.rs @@ -978,7 +978,8 @@ where let mut data = range_body_filter.filter_body(data); if let Some(duration) = self .inner - .response_body_filter(session, &mut data, eos, ctx)? + .response_body_filter(session, &mut data, eos, ctx) + .await? { trace!("delaying downstream response for {duration:?}"); time::sleep(duration).await; diff --git a/pingora-proxy/src/proxy_trait.rs b/pingora-proxy/src/proxy_trait.rs index d675c33cf..d8648918a 100644 --- a/pingora-proxy/src/proxy_trait.rs +++ b/pingora-proxy/src/proxy_trait.rs @@ -448,7 +448,7 @@ 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( + async fn upstream_response_body_filter( &self, _session: &mut Session, _body: &mut Option, @@ -459,7 +459,7 @@ pub trait ProxyHttp { } /// Similar to [Self::upstream_response_filter()] but for response trailers - fn upstream_response_trailer_filter( + async fn upstream_response_trailer_filter( &self, _session: &mut Session, _upstream_trailers: &mut header::HeaderMap, @@ -469,7 +469,7 @@ pub trait ProxyHttp { } /// Similar to [Self::response_filter()] but for response body chunks - fn response_body_filter( + async fn response_body_filter( &self, _session: &mut Session, _body: &mut Option, From 9c4747fd364d57f17349657ea063676311cac742 Mon Sep 17 00:00:00 2001 From: Andrew Hauck Date: Fri, 28 Aug 2026 12:24:08 -0700 Subject: [PATCH 2/5] Document async response filter requirements Require the request context to be Send and Sync for the newly asynchronous response filters, matching the other async ProxyHttp callbacks. Explain that these hooks can await I/O or offload expensive work without blocking request processing. --- .bleep | 2 +- pingora-proxy/src/proxy_trait.rs | 19 +++++++++++++++++-- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/.bleep b/.bleep index bd08489ff..b451bc5ce 100644 --- a/.bleep +++ b/.bleep @@ -1 +1 @@ -69857cc99cbb2e232b30e3b61fa060dafa1dfb7b \ No newline at end of file +16298af65271c10ff7e72abd76775c8163690d8c \ No newline at end of file diff --git a/pingora-proxy/src/proxy_trait.rs b/pingora-proxy/src/proxy_trait.rs index d8648918a..5e62a4ba1 100644 --- a/pingora-proxy/src/proxy_trait.rs +++ b/pingora-proxy/src/proxy_trait.rs @@ -448,27 +448,42 @@ 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**. + /// + /// The async nature of this function allows implementations to await I/O or offload expensive + /// work without blocking the task processing the request. async fn upstream_response_body_filter( &self, _session: &mut Session, _body: &mut Option, _end_of_stream: bool, _ctx: &mut Self::CTX, - ) -> Result> { + ) -> Result> + where + Self::CTX: Send + Sync, + { Ok(None) } /// Similar to [Self::upstream_response_filter()] but for response trailers + /// + /// The async nature of this function allows implementations to await I/O or offload expensive + /// work without blocking the task processing the request. async fn upstream_response_trailer_filter( &self, _session: &mut Session, _upstream_trailers: &mut header::HeaderMap, _ctx: &mut Self::CTX, - ) -> Result<()> { + ) -> Result<()> + where + Self::CTX: Send + Sync, + { Ok(()) } /// Similar to [Self::response_filter()] but for response body chunks + /// + /// The async nature of this function allows implementations to await I/O or offload expensive + /// work without blocking the task processing the request. async fn response_body_filter( &self, _session: &mut Session, From b8ce717fff1deeddef2e22b8aaaae2d04fee9a55 Mon Sep 17 00:00:00 2001 From: ewang Date: Thu, 27 Aug 2026 15:58:06 -0700 Subject: [PATCH 3/5] Parameterize custom downstream sessions --- .bleep | 2 +- pingora-cache/src/put.rs | 7 +- pingora-core/src/apps/mod.rs | 39 +- .../src/protocols/http/custom/server.rs | 56 ++- pingora-core/src/protocols/http/server.rs | 55 ++- .../src/protocols/http/subrequest/server.rs | 6 +- pingora-core/src/services/listening.rs | 53 ++- pingora-proxy/src/lib.rs | 343 +++++++++++------- pingora-proxy/src/proxy_cache.rs | 70 ++-- pingora-proxy/src/proxy_custom.rs | 51 +-- pingora-proxy/src/proxy_h1.rs | 55 +-- pingora-proxy/src/proxy_h2.rs | 51 +-- pingora-proxy/src/proxy_purge.rs | 11 +- pingora-proxy/src/proxy_trait.rs | 85 +++-- pingora-proxy/src/subrequest/mod.rs | 16 +- pingora-proxy/src/subrequest/pipe.rs | 18 +- 16 files changed, 560 insertions(+), 358 deletions(-) diff --git a/.bleep b/.bleep index b451bc5ce..aec9211d7 100644 --- a/.bleep +++ b/.bleep @@ -1 +1 @@ -16298af65271c10ff7e72abd76775c8163690d8c \ No newline at end of file +75f781c5d07dc5f9e292668db1d8872043d49b83 \ No newline at end of file diff --git a/pingora-cache/src/put.rs b/pingora-cache/src/put.rs index 0d118733f..6a7e6b537 100644 --- a/pingora-cache/src/put.rs +++ b/pingora-cache/src/put.rs @@ -20,7 +20,8 @@ use bytes::Bytes; use http::header; use log::warn; use pingora_core::protocols::http::{ - v1::common::header_value_content_length, HttpTask, ServerSession, + custom::server::Session as DownstreamSession, v1::common::header_value_content_length, + HttpTask, ServerSession, }; use pingora_error::Error; @@ -216,9 +217,9 @@ impl CachePutCtx { /// Return: /// - `Ok(None)` when the payload will be cache. /// - `Ok(Some(reason))` when the payload is not cacheable - pub async fn cache_put( + pub async fn cache_put( &mut self, - session: &mut ServerSession, + session: &mut ServerSession, ) -> Result> { let mut no_cache_reason = None; while let Some(data) = session.read_request_body().await? { diff --git a/pingora-core/src/apps/mod.rs b/pingora-core/src/apps/mod.rs index bd37037d2..8d4f5fef5 100644 --- a/pingora-core/src/apps/mod.rs +++ b/pingora-core/src/apps/mod.rs @@ -24,6 +24,7 @@ use std::any::Any; use std::sync::Arc; use std::time::Duration; +use crate::protocols::http::custom::server::Session as CustomServerSession; use crate::protocols::http::v2::server; use crate::protocols::http::{ReusableHttpStream, ServerSession}; use crate::protocols::Digest; @@ -35,7 +36,10 @@ const H2_PREFACE: &[u8] = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n"; #[async_trait] /// This trait defines the interface of a transport layer (TCP or TLS) application. -pub trait ServerApp { +pub trait ServerApp +where + DS: CustomServerSession, +{ /// Whenever a new connection is established, this function will be called with the established /// [`Stream`] object provided. /// @@ -120,7 +124,10 @@ pub struct HttpPersistentSettings { } impl HttpPersistentSettings { - pub fn for_session(session: &ServerSession) -> Self { + pub fn for_session(session: &ServerSession) -> Self + where + CS: CustomServerSession, + { HttpPersistentSettings { keepalive_timeout: session.get_keepalive(), keepalive_reuses_remaining: session.get_keepalive_reuses_remaining(), @@ -148,7 +155,10 @@ impl HttpPersistentSettings { self.pipelined_prefix = Some(prefix); } - pub fn apply_to_session(self, session: &mut ServerSession) { + pub fn apply_to_session(self, session: &mut ServerSession) + where + CS: CustomServerSession, + { let Self { keepalive_timeout, mut keepalive_reuses_remaining, @@ -214,7 +224,10 @@ impl ReusedHttpStream { /// This trait defines the interface of an HTTP application. #[async_trait] -pub trait HttpServerApp { +pub trait HttpServerApp +where + DS: CustomServerSession, +{ /// Similar to the [`ServerApp`], this function is called whenever a new HTTP session is established. /// /// After successful processing, [`ServerSession::finish()`] can be @@ -224,7 +237,7 @@ pub trait HttpServerApp { /// a `None` should be returned. async fn process_new_http( self: &Arc, - mut session: ServerSession, + mut session: ServerSession, // TODO: make this ShutdownWatch so that all task can await on this event shutdown: &ShutdownWatch, ) -> Option; @@ -258,9 +271,10 @@ pub trait HttpServerApp { } #[async_trait] -impl ServerApp for T +impl ServerApp for T where - T: HttpServerApp + Send + Sync + 'static, + T: HttpServerApp + Send + Sync + 'static, + DS: CustomServerSession, { async fn process_new( self: &Arc, @@ -334,8 +348,11 @@ where // loop's idle timeout sees this connection as busy. let _guard = guard; // Note, `PersistentSettings` not currently relevant for h2 - app.process_new_http(ServerSession::new_http2(h2_stream), &shutdown) - .await; + app.process_new_http( + ServerSession::::new_http2_with_custom_session(h2_stream), + &shutdown, + ) + .await; }); }, ) @@ -344,7 +361,7 @@ where return self.clone().process_custom_session(stream, shutdown).await; } else { // No ALPN or ALPN::H1 and h2c was not configured, fallback to HTTP/1.1 - let mut session = ServerSession::new_http1(stream); + let mut session = ServerSession::::new_http1_with_custom_session(stream); if *shutdown.borrow() { // stop downstream from reusing if this service is shutting down soon session.set_keepalive(None); @@ -359,7 +376,7 @@ where let mut result = self.process_new_http(session, shutdown).await; while let Some((stream, persistent_settings)) = result.map(|r| r.consume()) { - let mut session = ServerSession::new_http1(stream); + let mut session = ServerSession::::new_http1_with_custom_session(stream); if let Some(persistent_settings) = persistent_settings { persistent_settings.apply_to_session(&mut session); } diff --git a/pingora-core/src/protocols/http/custom/server.rs b/pingora-core/src/protocols/http/custom/server.rs index 65270b001..7a65d2aad 100644 --- a/pingora-core/src/protocols/http/custom/server.rs +++ b/pingora-core/src/protocols/http/custom/server.rs @@ -12,9 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::time::Duration; +use std::{future::Future, time::Duration}; -use async_trait::async_trait; use bytes::Bytes; use futures::Stream; use http::HeaderMap; @@ -26,25 +25,40 @@ use crate::protocols::{http::HttpTask, l4::socket::SocketAddr, Digest}; use super::CustomMessageWrite; #[doc(hidden)] -#[async_trait] +/// A concrete custom downstream session. +/// +/// Async operations return concrete `Send` futures so dispatch does not require +/// an allocation per call. Implementations may use `async fn` to satisfy these +/// methods. This trait is intentionally not dyn-compatible. pub trait Session: Send + Sync + Unpin + 'static { fn req_header(&self) -> &RequestHeader; fn req_header_mut(&mut self) -> &mut RequestHeader; - async fn read_body_bytes(&mut self) -> Result>; + fn read_body_bytes(&mut self) -> impl Future>> + Send; - async fn drain_request_body(&mut self) -> Result<()>; + fn drain_request_body(&mut self) -> impl Future> + Send; - async fn write_response_header(&mut self, resp: Box, end: bool) -> Result<()>; + fn write_response_header( + &mut self, + resp: Box, + end: bool, + ) -> impl Future> + Send; - async fn write_response_header_ref(&mut self, resp: &ResponseHeader, end: bool) -> Result<()>; + fn write_response_header_ref( + &mut self, + resp: &ResponseHeader, + end: bool, + ) -> impl Future> + Send; - async fn write_body(&mut self, data: Bytes, end: bool) -> Result<()>; + fn write_body(&mut self, data: Bytes, end: bool) -> impl Future> + Send; - async fn write_trailers(&mut self, trailers: HeaderMap) -> Result<()>; + fn write_trailers(&mut self, trailers: HeaderMap) -> impl Future> + Send; - async fn response_duplex_vec(&mut self, tasks: Vec) -> Result; + fn response_duplex_vec( + &mut self, + tasks: Vec, + ) -> impl Future> + Send; /// Whether the cancel-safe proxy task API is enabled for this session. fn proxy_tasks_enabled(&self) -> bool { @@ -72,8 +86,8 @@ pub trait Session: Send + Sync + Unpin + 'static { /// /// # Panics /// Panics if the Custom session does not implement the proxy task API. - async fn write_proxy_tasks(&mut self) -> Result { - panic!("Custom proxy task API not implemented") + fn write_proxy_tasks(&mut self) -> impl Future> + Send { + async { panic!("Custom proxy task API not implemented") } } fn set_read_timeout(&mut self, timeout: Option); @@ -92,7 +106,7 @@ pub trait Session: Send + Sync + Unpin + 'static { fn response_written(&self) -> Option<&ResponseHeader>; - async fn shutdown(&mut self, code: u32, ctx: &str); + fn shutdown(&mut self, code: u32, ctx: &str) -> impl Future + Send; /// Abandon the response mid-message, in a way the peer can tell apart from a /// response that was completed. @@ -106,17 +120,22 @@ pub trait Session: Send + Sync + Unpin + 'static { /// implementations that predate this method. Implementations whose protocol /// can distinguish an abandoned message from a completed one should override /// this; otherwise a peer may read the abandoned message as successful. - async fn abandon(&mut self, ctx: &str) { - self.shutdown(0, ctx).await; + fn abandon(&mut self, ctx: &str) -> impl Future + Send { + async move { + self.shutdown(0, ctx).await; + } } fn is_body_done(&mut self) -> bool; - async fn finish(&mut self) -> Result<()>; + fn finish(&mut self) -> impl Future> + Send; fn is_body_empty(&mut self) -> bool; - async fn read_body_or_idle(&mut self, no_body_expected: bool) -> Result>; + fn read_body_or_idle( + &mut self, + no_body_expected: bool, + ) -> impl Future>> + Send; fn body_bytes_sent(&self) -> usize; @@ -138,7 +157,7 @@ pub trait Session: Send + Sync + Unpin + 'static { fn get_retry_buffer(&self) -> Option; - async fn finish_custom(&mut self) -> Result<()>; + fn finish_custom(&mut self) -> impl Future> + Send; fn take_custom_message_reader( &mut self, @@ -169,7 +188,6 @@ pub trait Session: Send + Sync + Unpin + 'static { } #[doc(hidden)] -#[async_trait] impl Session for () { fn req_header(&self) -> &RequestHeader { unreachable!("server session: req_header") diff --git a/pingora-core/src/protocols/http/server.rs b/pingora-core/src/protocols/http/server.rs index a043e6426..f955ce26b 100644 --- a/pingora-core/src/protocols/http/server.rs +++ b/pingora-core/src/protocols/http/server.rs @@ -53,31 +53,57 @@ impl ReusableHttpStream { } /// HTTP server session object for both HTTP/1.x and HTTP/2 -pub enum Session { +pub enum Session +where + CS: SessionCustom, +{ H1(SessionV1), H2(SessionV2), Subrequest(SessionSubrequest), - Custom(Box), + Custom(CS), } -impl Session { +impl Session<()> { /// Create a new [`Session`] from an established connection for HTTP/1.x pub fn new_http1(stream: Stream) -> Self { - Self::H1(SessionV1::new(stream)) + Self::new_http1_with_custom_session(stream) } /// Create a new [`Session`] from an established HTTP/2 stream pub fn new_http2(session: SessionV2) -> Self { - Self::H2(session) + Self::new_http2_with_custom_session(session) } /// Create a new [`Session`] from a subrequest session pub fn new_subrequest(session: SessionSubrequest) -> Self { + Self::new_subrequest_with_custom_session(session) + } +} + +impl Session +where + CS: SessionCustom, +{ + /// Create a new [`Session`] with a concrete custom-session type from an + /// established connection for HTTP/1.x. + pub fn new_http1_with_custom_session(stream: Stream) -> Self { + Self::H1(SessionV1::new(stream)) + } + + /// Create a new [`Session`] with a concrete custom-session type from an + /// established HTTP/2 stream. + pub fn new_http2_with_custom_session(session: SessionV2) -> Self { + Self::H2(session) + } + + /// Create a new [`Session`] with a concrete custom-session type from a + /// subrequest session. + pub fn new_subrequest_with_custom_session(session: SessionSubrequest) -> Self { Self::Subrequest(session) } /// Create a new [`Session`] from a custom session - pub fn new_custom(session: Box) -> Self { + pub fn new_custom(session: CS) -> Self { Self::Custom(session) } @@ -745,16 +771,16 @@ impl Session { } } - pub fn as_custom(&self) -> Option<&dyn SessionCustom> { + pub fn as_custom(&self) -> Option<&CS> { match self { Self::H1(_) => None, Self::H2(_) => None, Self::Subrequest(_) => None, - Self::Custom(c) => Some(c.as_ref()), + Self::Custom(c) => Some(c), } } - pub fn as_custom_mut(&mut self) -> Option<&mut Box> { + pub fn as_custom_mut(&mut self) -> Option<&mut CS> { match self { Self::H1(_) => None, Self::H2(_) => None, @@ -1005,14 +1031,13 @@ impl Session { mod tests { use super::*; use crate::protocols::http::custom::CustomMessageWrite; - use async_trait::async_trait; use futures::Stream; use std::panic::{catch_unwind, AssertUnwindSafe}; use std::sync::{Arc, Mutex}; #[tokio::test] async fn custom_proxy_task_defaults_are_opted_out_and_fail_loudly() { - let mut session = Session::new_custom(Box::new(())); + let mut session = Session::new_custom(()); assert!(!session.supports_proxy_task_api()); session.set_proxy_tasks_enabled(true); @@ -1030,7 +1055,7 @@ mod tests { #[tokio::test] async fn custom_proxy_task_methods_delegate_to_the_custom_session() { - let mut session = Session::new_custom(Box::new(ProxyTaskCustom::new())); + let mut session = Session::new_custom(ProxyTaskCustom::new()); assert!(!session.supports_proxy_task_api()); session.set_proxy_tasks_enabled(true); @@ -1050,9 +1075,8 @@ mod tests { #[tokio::test] async fn custom_session_shutdown_signals_an_incomplete_message() { let shutdown_calls = Arc::new(Mutex::new(Vec::new())); - let mut session = Session::new_custom(Box::new(ProxyTaskCustom::with_shutdown_calls( - shutdown_calls.clone(), - ))); + let mut session = + Session::new_custom(ProxyTaskCustom::with_shutdown_calls(shutdown_calls.clone())); session.shutdown().await; @@ -1084,7 +1108,6 @@ mod tests { } } - #[async_trait] impl SessionCustom for ProxyTaskCustom { fn req_header(&self) -> &RequestHeader { &self.header diff --git a/pingora-core/src/protocols/http/subrequest/server.rs b/pingora-core/src/protocols/http/subrequest/server.rs index 938c8c97a..49e336fc1 100644 --- a/pingora-core/src/protocols/http/subrequest/server.rs +++ b/pingora-core/src/protocols/http/subrequest/server.rs @@ -46,6 +46,7 @@ use tokio::sync::{mpsc, oneshot}; use super::body::{BodyMode, BodyReader, BodyWriter, PREMATURE_BODY_END}; use crate::protocols::http::{ body_buffer::FixedBuffer, + custom::server::Session as CustomServerSession, server::Session as GenericHttpSession, subrequest::dummy::DummyIO, v1::common::{header_value_content_length, is_chunked_encoding_from_headers, BODY_BUF_LIMIT}, @@ -152,7 +153,10 @@ impl HttpSession { /// Create a new http server session for a subrequest. /// The created session needs to call [`Self::read_request()`] first before performing /// any other operations. - pub fn new_from_session(session: &GenericHttpSession) -> (Self, SubrequestHandle) { + pub fn new_from_session(session: &GenericHttpSession) -> (Self, SubrequestHandle) + where + DS: CustomServerSession, + { let v1_inner = SessionV1::new(Box::new(DummyIO::new(&session.to_h1_raw()))); let digest = session.digest().cloned(); // allow buffering a small number of tasks, otherwise exert backpressure diff --git a/pingora-core/src/services/listening.rs b/pingora-core/src/services/listening.rs index 016ed4ede..16d9684d9 100644 --- a/pingora-core/src/services/listening.rs +++ b/pingora-core/src/services/listening.rs @@ -25,6 +25,7 @@ use crate::listeners::AcceptAllFilter; use crate::listeners::{ ConnectionFilter, ListenerConfig, Listeners, ServerAddress, TcpSocketOptions, TransportStack, }; +use crate::protocols::http::custom::server::Session as CustomServerSession; use crate::protocols::Stream; #[cfg(unix)] use crate::server::ListenFds; @@ -37,6 +38,7 @@ use pingora_error::Result; use pingora_runtime::current_handle; use pingora_timeout::timeout; use std::fs::Permissions; +use std::marker::PhantomData; use std::sync::Arc; use std::time::Duration; @@ -44,10 +46,14 @@ use std::time::Duration; pub type RuntimeOptsOverride = Arc Option + Send + Sync>; /// The type of service that is associated with a list of listening endpoints and a particular application -pub struct Service { +pub struct Service +where + DS: CustomServerSession, +{ name: String, listeners: Listeners, app_logic: Option, + _custom_session: PhantomData DS>, /// The number of preferred threads. `None` to follow global setting. pub threads: Option, runtime_opts_override: Option, @@ -55,13 +61,17 @@ pub struct Service { connection_filter: Arc, } -impl Service { - /// Create a new [`Service`] with the given application (see [`crate::apps`]). - pub fn new(name: String, app_logic: A) -> Self { +impl Service +where + DS: CustomServerSession, +{ + /// Create a new [`Service`] with a concrete custom downstream session type. + pub fn new_with_custom_session(name: String, app_logic: A) -> Self { Service { name, listeners: Listeners::new(), app_logic: Some(app_logic), + _custom_session: PhantomData, threads: None, runtime_opts_override: None, #[cfg(feature = "connection_filter")] @@ -69,13 +79,17 @@ impl Service { } } - /// Create a new [`Service`] with the given application (see [`crate::apps`]) and the given - /// [`Listeners`]. - pub fn with_listeners(name: String, listeners: Listeners, app_logic: A) -> Self { + /// Create a new [`Service`] with listeners and a concrete custom downstream session type. + pub fn with_listeners_and_custom_session( + name: String, + listeners: Listeners, + app_logic: A, + ) -> Self { Service { name, listeners, app_logic: Some(app_logic), + _custom_session: PhantomData, threads: None, runtime_opts_override: None, #[cfg(feature = "connection_filter")] @@ -187,7 +201,24 @@ impl Service { } } -impl Service { +impl Service { + /// Create a new [`Service`] with the given application (see [`crate::apps`]). + pub fn new(name: String, app_logic: A) -> Self { + Self::new_with_custom_session(name, app_logic) + } + + /// Create a new [`Service`] with the given application (see [`crate::apps`]) and the given + /// [`Listeners`]. + pub fn with_listeners(name: String, listeners: Listeners, app_logic: A) -> Self { + Self::with_listeners_and_custom_session(name, listeners, app_logic) + } +} + +impl Service +where + A: ServerApp + Send + Sync + 'static, + DS: CustomServerSession, +{ pub async fn handle_event(event: Stream, app_logic: Arc, shutdown: ShutdownWatch) { debug!("new event!"); let mut reuse_event = app_logic.process_new(event, &shutdown).await; @@ -275,7 +306,11 @@ impl Service { } #[async_trait] -impl ServiceTrait for Service { +impl ServiceTrait for Service +where + A: ServerApp + Send + Sync + 'static, + DS: CustomServerSession, +{ async fn start_service( &mut self, #[cfg(unix)] fds: Option, diff --git a/pingora-proxy/src/lib.rs b/pingora-proxy/src/lib.rs index 261674dc3..2ce2351b1 100644 --- a/pingora-proxy/src/lib.rs +++ b/pingora-proxy/src/lib.rs @@ -64,6 +64,7 @@ use pingora_core::connectors::{http::Connector, ConnectorOptions}; use pingora_core::modules::http::compression::ResponseCompressionBuilder; use pingora_core::modules::http::{HttpModuleCtx, HttpModules}; use pingora_core::protocols::http::client::HttpSession as ClientSession; +use pingora_core::protocols::http::custom::server::Session as DownstreamSession; use pingora_core::protocols::http::custom::CustomMessageWrite; use pingora_core::protocols::http::subrequest::server::SubrequestHandle; use pingora_core::protocols::http::v1::client::HttpSession as HttpSessionV1; @@ -105,8 +106,12 @@ pub mod prelude { pub use crate::{http_proxy, http_proxy_service, ProxyHttp, ProxyWarnLogContext, Session}; } -pub type ProcessCustomSession = Arc< - dyn Fn(Arc>, Stream, &ShutdownWatch) -> BoxFuture<'static, Option> +/// Type-erased custom-connection callback. +/// +/// The callback remains boxed once per accepted custom connection; the +/// downstream session operations themselves use concrete futures. +pub type ProcessCustomSession = Arc< + dyn Fn(Arc>, Stream, &ShutdownWatch) -> BoxFuture<'static, Option> + Send + Sync + Unpin @@ -178,9 +183,10 @@ impl ShardedNotify { /// The concrete type that holds the user defined HTTP proxy. /// /// Users don't need to interact with this object directly. -pub struct HttpProxy +pub struct HttpProxy where C: custom::Connector, // Upstream custom connector + DS: DownstreamSession, { inner: SV, // TODO: name it better than inner client_upstream: Connector, @@ -192,10 +198,10 @@ where #[cfg(feature = "upstream_modules")] pub upstream_modules: HttpModules, max_retries: usize, - process_custom_session: Option>, + process_custom_session: Option>, } -impl HttpProxy { +impl HttpProxy { /// Create a new [`HttpProxy`] with the given [`ProxyHttp`] implementation and [`ServerConf`]. /// /// After creating an `HttpProxy`, you should call [`HttpProxy::handle_init_modules()`] to @@ -233,21 +239,22 @@ impl HttpProxy { } } -impl HttpProxy +impl HttpProxy where C: custom::Connector, + DS: DownstreamSession, { fn new_custom( inner: SV, conf: Arc, connector: C, - on_custom: Option>, + on_custom: Option>, server_options: Option, client_options: Option, ) -> Self where - SV: ProxyHttp + Send + Sync + 'static, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync + 'static, + >::CTX: Send + Sync, { let client_options = client_options.unwrap_or_else(|| ConnectorOptions::from_server_conf(&conf)); @@ -289,7 +296,7 @@ where /// this method is called automatically. pub fn handle_init_modules(&mut self) where - SV: ProxyHttp, + SV: ProxyHttp, { self.inner .init_downstream_modules(&mut self.downstream_modules); @@ -321,11 +328,11 @@ where async fn handle_new_request( &self, - mut downstream_session: Box, - ) -> Option> + mut downstream_session: Box>, + ) -> Option>> where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { // phase 1 read request header @@ -400,12 +407,12 @@ where // return bool: server_session can be reused, and error if any async fn proxy_to_upstream( &self, - session: &mut Session, - ctx: &mut SV::CTX, + session: &mut Session, + ctx: &mut >::CTX, ) -> (bool, Option>) where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { let peer = match self.inner.upstream_peer(session, ctx).await { Ok(p) => p, @@ -486,13 +493,13 @@ where async fn upstream_filter( &self, - session: &mut Session, + session: &mut Session, task: &mut HttpTask, - ctx: &mut SV::CTX, + ctx: &mut >::CTX, ) -> Result> where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { let duration = match task { HttpTask::Header(header, _eos) => { @@ -523,14 +530,14 @@ where async fn finish( &self, - mut session: Session, - ctx: &mut SV::CTX, + mut session: Session, + ctx: &mut >::CTX, reuse: bool, error: Option>, ) -> Option where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { self.inner .logging(&mut session, error.as_deref(), ctx) @@ -558,7 +565,7 @@ where } } - fn cleanup_sub_req(&self, session: &mut Session) { + fn cleanup_sub_req(&self, session: &mut Session) { if let Some(ctx) = session.subrequest_ctx.as_mut() { ctx.release_write_lock(); } @@ -572,9 +579,12 @@ use pingora_core::protocols::http::compression::ResponseCompressionCtx; /// /// This object is what users interact with in order to access the request itself or change the proxy /// behavior. -pub struct Session { +pub struct Session +where + DS: DownstreamSession, +{ /// the HTTP session to downstream (the client) - pub downstream_session: Box, + pub downstream_session: Box>, /// The interface to control HTTP caching pub cache: HttpCache, /// (de)compress responses coming into the proxy (from upstream) @@ -588,7 +598,7 @@ pub struct Session { /// The context from parent request, if this is a subrequest. pub subrequest_ctx: Option>, /// Handle to allow spawning subrequests, assigned by the `Subrequest` app logic. - pub subrequest_spawner: Option, + pub subrequest_spawner: Option>, // Downstream filter modules pub downstream_modules_ctx: HttpModuleCtx, /// Upstream filter modules. These run before `upstream_compression` and see the raw @@ -612,9 +622,12 @@ pub struct Session { shutdown_flag: Arc, } -impl Session { +impl Session +where + DS: DownstreamSession, +{ fn new( - downstream_session: impl Into>, + downstream_session: impl Into>>, downstream_modules: &HttpModules, #[cfg(feature = "upstream_modules")] upstream_modules: &HttpModules, shutdown_flag: Arc, @@ -640,35 +653,6 @@ impl Session { } } - /// Create a new [Session] from the given [Stream] - /// - /// This function is mostly used for testing and mocking, given the downstream modules and - /// shutdown flags will never be set. - pub fn new_h1(stream: Stream) -> Self { - let modules = HttpModules::new(); - Self::new( - Box::new(HttpSession::new_http1(stream)), - &modules, - #[cfg(feature = "upstream_modules")] - &HttpModules::new(), - Arc::new(AtomicBool::new(false)), - ) - } - - /// Create a new [Session] from the given [Stream] with modules - /// - /// This function is mostly used for testing and mocking, given the shutdown flag will never be - /// set. - pub fn new_h1_with_modules(stream: Stream, downstream_modules: &HttpModules) -> Self { - Self::new( - Box::new(HttpSession::new_http1(stream)), - downstream_modules, - #[cfg(feature = "upstream_modules")] - &HttpModules::new(), - Arc::new(AtomicBool::new(false)), - ) - } - /// Run upstream module filters on the given [`HttpTask`]. /// /// Upstream modules process each task **before** `upstream_compression` and @@ -704,11 +688,11 @@ impl Session { Ok(()) } - pub fn as_downstream_mut(&mut self) -> &mut HttpSession { + pub fn as_downstream_mut(&mut self) -> &mut HttpSession { &mut self.downstream_session } - pub fn as_downstream(&self) -> &HttpSession { + pub fn as_downstream(&self) -> &HttpSession { &self.downstream_session } @@ -1012,6 +996,37 @@ impl Session { } } +impl Session<()> { + /// Create a new [Session] from the given [Stream] + /// + /// This function is mostly used for testing and mocking, given the downstream modules and + /// shutdown flags will never be set. + pub fn new_h1(stream: Stream) -> Self { + let modules = HttpModules::new(); + Self::new( + Box::new(HttpSession::new_http1(stream)), + &modules, + #[cfg(feature = "upstream_modules")] + &HttpModules::new(), + Arc::new(AtomicBool::new(false)), + ) + } + + /// Create a new [Session] from the given [Stream] with modules + /// + /// This function is mostly used for testing and mocking, given the shutdown flag will never be + /// set. + pub fn new_h1_with_modules(stream: Stream, downstream_modules: &HttpModules) -> Self { + Self::new( + Box::new(HttpSession::new_http1(stream)), + downstream_modules, + #[cfg(feature = "upstream_modules")] + &HttpModules::new(), + Arc::new(AtomicBool::new(false)), + ) + } +} + #[derive(Clone, Copy, Debug, Default)] struct H1UpgradeRequestStatus { upstream: Option, @@ -1037,11 +1052,14 @@ impl H1UpgradeRequestSnapshot { /// 101 can establish a tunnel. Otherwise one side changes protocol while the /// other stays in HTTP handling, allowing tunneled traffic to bypass request /// processing or corrupt the connection state. -fn reject_mismatched_h1_upgrade_101( - session: &Session, +fn reject_mismatched_h1_upgrade_101( + session: &Session, header: &ResponseHeader, stage: &'static str, -) -> Result<()> { +) -> Result<()> +where + DS: DownstreamSession, +{ if header.status != http::StatusCode::SWITCHING_PROTOCOLS { return Ok(()); } @@ -1065,11 +1083,14 @@ fn reject_mismatched_h1_upgrade_101( Ok(()) } -fn reject_unexpected_task_after_h1_upgrade( - session: &Session, +fn reject_unexpected_task_after_h1_upgrade( + session: &Session, task: &'static str, task_filter_seen_upgraded: bool, -) -> Result<()> { +) -> Result<()> +where + DS: DownstreamSession, +{ let status = session.h1_upgrade_request_snapshot(); Error::e_explain( InvalidHTTPHeader, @@ -1085,10 +1106,13 @@ fn reject_unexpected_task_after_h1_upgrade( .map_err(|e| e.into_in()) } -fn reject_unexpected_upgraded_body_before_h1_upgrade( - session: &Session, +fn reject_unexpected_upgraded_body_before_h1_upgrade( + session: &Session, task_filter_seen_upgraded: bool, -) -> Result<()> { +) -> Result<()> +where + DS: DownstreamSession, +{ let status = session.h1_upgrade_request_snapshot(); Error::e_explain( InvalidHTTPHeader, @@ -1104,29 +1128,41 @@ fn reject_unexpected_upgraded_body_before_h1_upgrade( .map_err(|e| e.into_in()) } -impl AsRef for Session { - fn as_ref(&self) -> &HttpSession { +impl AsRef> for Session +where + DS: DownstreamSession, +{ + fn as_ref(&self) -> &HttpSession { &self.downstream_session } } -impl AsMut for Session { - fn as_mut(&mut self) -> &mut HttpSession { +impl AsMut> for Session +where + DS: DownstreamSession, +{ + fn as_mut(&mut self) -> &mut HttpSession { &mut self.downstream_session } } use std::ops::{Deref, DerefMut}; -impl Deref for Session { - type Target = HttpSession; +impl Deref for Session +where + DS: DownstreamSession, +{ + type Target = HttpSession; fn deref(&self) -> &Self::Target { &self.downstream_session } } -impl DerefMut for Session { +impl DerefMut for Session +where + DS: DownstreamSession, +{ fn deref_mut(&mut self) -> &mut Self::Target { &mut self.downstream_session } @@ -1144,18 +1180,19 @@ static BAD_GATEWAY: Lazy = Lazy::new(|| { resp }); -impl HttpProxy +impl HttpProxy where C: custom::Connector, + DS: DownstreamSession, { async fn process_request( self: &Arc, - mut session: Session, - mut ctx: ::CTX, + mut session: Session, + mut ctx: >::CTX, ) -> Option where - SV: ProxyHttp + Send + Sync + 'static, - ::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync + 'static, + >::CTX: Send + Sync, { if let Err(e) = self .inner @@ -1369,14 +1406,14 @@ where async fn handle_error( &self, - mut session: Session, - ctx: &mut ::CTX, + mut session: Session, + ctx: &mut >::CTX, e: Box, context: &str, ) -> Option where - SV: ProxyHttp + Send + Sync + 'static, - ::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync + 'static, + >::CTX: Send + Sync, { let res = self.inner.fail_to_proxy(&mut session, &e, ctx).await; if !self.inner.suppress_error_log(&session, ctx, &e) { @@ -1420,24 +1457,28 @@ error[E0391]: cycle detected when computing type of `proxy_cache:: +where + DS: DownstreamSession, +{ async fn process_subrequest( self: Arc, - session: Box, + session: Box>, sub_req_ctx: Box, ); } #[async_trait] -impl Subrequest for HttpProxy +impl Subrequest for HttpProxy where - SV: ProxyHttp + Send + Sync + 'static, - ::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync + 'static, + >::CTX: Send + Sync, C: custom::Connector, + DS: DownstreamSession, { async fn process_subrequest( self: Arc, - session: Box, + session: Box>, sub_req_ctx: Box, ) { debug!("starting subrequest"); @@ -1466,36 +1507,48 @@ where } /// A handle to the underlying HTTP proxy app that allows spawning subrequests. -pub struct SubrequestSpawner { - app: Arc, +pub struct SubrequestSpawner +where + DS: DownstreamSession, +{ + app: Arc + Send + Sync>, } /// A [`PreparedSubrequest`] that is ready to run. -pub struct PreparedSubrequest { - app: Arc, - session: Box, +pub struct PreparedSubrequest +where + DS: DownstreamSession, +{ + app: Arc + Send + Sync>, + session: Box>, sub_req_ctx: Box, } -impl PreparedSubrequest { +impl PreparedSubrequest +where + DS: DownstreamSession, +{ pub async fn run(self) { self.app .process_subrequest(self.session, self.sub_req_ctx) .await } - pub fn session(&self) -> &HttpSession { + pub fn session(&self) -> &HttpSession { self.session.as_ref() } - pub fn session_mut(&mut self) -> &mut HttpSession { + pub fn session_mut(&mut self) -> &mut HttpSession { self.session.deref_mut() } } -impl SubrequestSpawner { +impl SubrequestSpawner +where + DS: DownstreamSession, +{ /// Create a new [`SubrequestSpawner`]. - pub fn new(app: Arc) -> SubrequestSpawner { + pub fn new(app: Arc + Send + Sync>) -> SubrequestSpawner { SubrequestSpawner { app } } @@ -1503,7 +1556,7 @@ impl SubrequestSpawner { // TODO: allow configuring the subrequest session before use pub fn spawn_background_subrequest( &self, - session: &HttpSession, + session: &HttpSession, ctx: SubrequestCtx, ) -> tokio::task::JoinHandle<()> { let new_app = self.app.clone(); // Clone the Arc @@ -1530,9 +1583,9 @@ impl SubrequestSpawner { // TODO: allow configuring the subrequest session before use pub fn create_subrequest( &self, - session: &HttpSession, + session: &HttpSession, ctx: SubrequestCtx, - ) -> (PreparedSubrequest, SubrequestHandle) { + ) -> (PreparedSubrequest, SubrequestHandle) { let new_app = self.app.clone(); // Clone the Arc let (mut session, handle) = subrequest::create_session(session); if ctx.body_mode() == BodyMode::NoBody { @@ -1554,15 +1607,16 @@ impl SubrequestSpawner { } #[async_trait] -impl HttpServerApp for HttpProxy +impl HttpServerApp for HttpProxy where - SV: ProxyHttp + Send + Sync + 'static, - ::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync + 'static, + >::CTX: Send + Sync, C: custom::Connector, + DS: DownstreamSession, { async fn process_new_http( self: &Arc, - mut session: HttpSession, + mut session: HttpSession, shutdown: &ShutdownWatch, ) -> Option { // Extract user context from the previous request before the session is moved into the Box @@ -1696,26 +1750,47 @@ where Service::new(name.to_string(), proxy) } -/// Create a [Service] from the user implemented [ProxyHttp]. +/// Create a [`Service`] with a custom upstream connector and standard HTTP downstream sessions. /// -/// The returned [Service] can be hosted by a [pingora_core::server::Server] directly. -pub fn http_proxy_service_with_name_custom( +/// The returned [`Service`] can be hosted by a [`pingora_core::server::Server`] directly. +pub fn http_proxy_service_with_name_custom_connector( conf: &Arc, inner: SV, name: &str, connector: C, - on_custom: ProcessCustomSession, ) -> Service> where SV: ProxyHttp + Send + Sync + 'static, SV::CTX: Send + Sync + 'static, C: custom::Connector, +{ + let mut proxy = HttpProxy::new_custom(inner, conf.clone(), connector, None, None, None); + proxy.handle_init_modules(); + + Service::new(name.to_string(), proxy) +} + +/// Create a [Service] from the user implemented [ProxyHttp]. +/// +/// The returned [Service] can be hosted by a [pingora_core::server::Server] directly. +pub fn http_proxy_service_with_name_custom( + conf: &Arc, + inner: SV, + name: &str, + connector: C, + on_custom: ProcessCustomSession, +) -> Service, DS> +where + SV: ProxyHttp + Send + Sync + 'static, + >::CTX: Send + Sync + 'static, + C: custom::Connector, + DS: DownstreamSession, { let mut proxy = HttpProxy::new_custom(inner, conf.clone(), connector, Some(on_custom), None, None); proxy.handle_init_modules(); - Service::new(name.to_string(), proxy) + Service::<_, DS>::new_with_custom_session(name.to_string(), proxy) } /// A builder for a [Service] that can be used to create a [HttpProxy] instance @@ -1723,27 +1798,22 @@ where /// The [ProxyServiceBuilder] can be used to construct a [HttpProxy] service with a custom name, /// connector, and custom session handler. /// -pub struct ProxyServiceBuilder +pub struct ProxyServiceBuilder where - SV: ProxyHttp + Send + Sync + 'static, - SV::CTX: Send + Sync + 'static, C: custom::Connector, + DS: DownstreamSession, { conf: Arc, inner: SV, name: String, connector: C, - custom: Option>, + custom: Option>, server_options: Option, client_options: Option, runtime_opts_override: Option, } -impl ProxyServiceBuilder -where - SV: ProxyHttp + Send + Sync + 'static, - SV::CTX: Send + Sync + 'static, -{ +impl ProxyServiceBuilder { /// Create a new [ProxyServiceBuilder] with the given [ServerConf] and [ProxyHttp] /// implementation. /// @@ -1767,11 +1837,10 @@ where } } -impl ProxyServiceBuilder +impl ProxyServiceBuilder where - SV: ProxyHttp + Send + Sync + 'static, - SV::CTX: Send + Sync + 'static, C: custom::Connector, + DS: DownstreamSession, { /// Sets the name of the [HttpProxy] service. pub fn name(mut self, name: impl AsRef) -> Self { @@ -1787,11 +1856,15 @@ where /// between the proxy and the upstream server. /// /// Returns a new [ProxyServiceBuilder] with the custom connector and session handler. - pub fn custom( + pub fn custom( self, connector: C2, - on_custom: ProcessCustomSession, - ) -> ProxyServiceBuilder { + on_custom: ProcessCustomSession, + ) -> ProxyServiceBuilder + where + C2: custom::Connector, + DS2: DownstreamSession, + { let Self { conf, inner, @@ -1846,7 +1919,11 @@ where /// a fully initialized [HttpProxy]. /// /// The returned [Service] is ready to be used by a [pingora_core::server::Server]. - pub fn build(self) -> Service> { + pub fn build(self) -> Service, DS> + where + SV: ProxyHttp + Send + Sync + 'static, + >::CTX: Send + Sync + 'static, + { let Self { conf, inner, @@ -1868,7 +1945,7 @@ where ); proxy.handle_init_modules(); - let mut service = Service::new(name, proxy); + let mut service = Service::<_, DS>::new_with_custom_session(name, proxy); if let Some(runtime_opts_override) = runtime_opts_override { service.set_runtime_opts_override(runtime_opts_override); } diff --git a/pingora-proxy/src/proxy_cache.rs b/pingora-proxy/src/proxy_cache.rs index 162dcb794..e68db9f6c 100644 --- a/pingora-proxy/src/proxy_cache.rs +++ b/pingora-proxy/src/proxy_cache.rs @@ -27,20 +27,21 @@ use std::time::SystemTime; const DEFAULT_MAX_CACHE_LOCK_RETRIES: usize = 2; -impl HttpProxy +impl HttpProxy where C: custom::Connector, + DS: DownstreamSession, { // return bool: server_session can be reused, and error if any pub(crate) async fn proxy_cache( self: &Arc, - session: &mut Session, - ctx: &mut SV::CTX, + session: &mut Session, + ctx: &mut >::CTX, ) -> Option<(bool, Option>)> // None: continue to proxy, Some: return where - SV: ProxyHttp + Send + Sync + 'static, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync + 'static, + >::CTX: Send + Sync, { // Cache logic request phase if let Err(e) = self.inner.request_cache_filter(session, ctx) { @@ -291,12 +292,12 @@ where // return bool: server_session can be reused, and error if any pub(crate) async fn proxy_cache_hit( &self, - session: &mut Session, - ctx: &mut SV::CTX, + session: &mut Session, + ctx: &mut >::CTX, ) -> (bool, Option>) where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { use range_filter::*; @@ -535,11 +536,11 @@ where pub(crate) fn downstream_response_conditional_filter( &self, use_cache: &mut ServeFromCache, - session: &Session, + session: &Session, resp: &mut ResponseHeader, - ctx: &mut SV::CTX, + ctx: &mut >::CTX, ) where - SV: ProxyHttp, + SV: ProxyHttp, { // TODO: range let req = session.req_header(); @@ -570,9 +571,12 @@ where // TODO: cache upstream header filter to add/remove headers - async fn finish_miss_handler_best_effort(&self, session: &mut Session, ctx: &SV::CTX) - where - SV: ProxyHttp, + async fn finish_miss_handler_best_effort( + &self, + session: &mut Session, + ctx: &>::CTX, + ) where + SV: ProxyHttp, { if let Err(e) = session.cache.finish_miss_handler().await { warn!( @@ -585,14 +589,14 @@ where pub(crate) async fn cache_http_task( &self, - session: &mut Session, + session: &mut Session, task: &HttpTask, - ctx: &mut SV::CTX, + ctx: &mut >::CTX, serve_from_cache: &mut ServeFromCache, ) -> Result<()> where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { if !session.cache.enabled() && !session.cache.bypassing() { return Ok(()); @@ -774,13 +778,13 @@ where // Return true if local cache should be used, false otherwise pub(crate) async fn revalidate_or_stale( &self, - session: &mut Session, + session: &mut Session, task: &mut HttpTask, - ctx: &mut SV::CTX, + ctx: &mut >::CTX, ) -> bool where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { if !session.cache.enabled() { return false; @@ -909,13 +913,13 @@ where // bool: can the downstream connection be reused pub(crate) async fn handle_stale_if_error( &self, - session: &mut Session, - ctx: &mut SV::CTX, + session: &mut Session, + ctx: &mut >::CTX, error: &Error, ) -> Option<(bool, Option>)> where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { // the caller might already checked this as an optimization if !session.cache.can_serve_stale_error() { @@ -951,12 +955,12 @@ where // helper function to check when to continue to retry lock (true) or give up (false) fn handle_lock_wait_outcome( &self, - session: &mut Session, - ctx: &SV::CTX, + session: &mut Session, + ctx: &>::CTX, outcome: LockWaitOutcome, ) -> bool where - SV: ProxyHttp, + SV: ProxyHttp, { debug!("cache unlocked {outcome:?}"); match outcome { @@ -1004,12 +1008,12 @@ where fn cache_lock_retry_limit_exceeded( &self, - session: &mut Session, - ctx: &SV::CTX, + session: &mut Session, + ctx: &>::CTX, cache_lock_retries: &mut usize, ) -> bool where - SV: ProxyHttp, + SV: ProxyHttp, { *cache_lock_retries += 1; let max_retries = session diff --git a/pingora-proxy/src/proxy_custom.rs b/pingora-proxy/src/proxy_custom.rs index f21974b8c..ae86f9455 100644 --- a/pingora-proxy/src/proxy_custom.rs +++ b/pingora-proxy/src/proxy_custom.rs @@ -30,23 +30,24 @@ use tokio::sync::oneshot; use super::*; -impl HttpProxy +impl HttpProxy where C: custom::Connector, + DS: DownstreamSession, { /// Proxy to a custom protocol upstream. /// Returns (reuse_server, error) pub(crate) async fn proxy_to_custom_upstream( &self, - session: &mut Session, + session: &mut Session, client_session: &mut C::Session, reused: bool, peer: &HttpPeer, - ctx: &mut SV::CTX, + ctx: &mut >::CTX, ) -> (bool, Option>) where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { #[cfg(windows)] let raw = client_session.fd() as std::os::windows::io::RawSocket; @@ -75,14 +76,14 @@ where /// Returns (reuse_server, error) async fn custom_proxy_down_to_up( &self, - session: &mut Session, + session: &mut Session, client_session: &mut C::Session, peer: &HttpPeer, - ctx: &mut SV::CTX, + ctx: &mut >::CTX, ) -> (bool, Option>) where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { client_session.set_read_timeout(peer.options.read_timeout); client_session.set_write_timeout(peer.options.write_timeout); @@ -284,8 +285,8 @@ where #[allow(clippy::too_many_arguments)] async fn process_upstream_tasks_custom( &self, - session: &mut Session, - ctx: &mut SV::CTX, + session: &mut Session, + ctx: &mut >::CTX, initial_task: HttpTask, rx: &mut mpsc::Receiver, serve_from_cache: &mut ServeFromCache, @@ -293,8 +294,8 @@ where response_state: &mut ResponseStateMachine, ) -> Result> where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { if serve_from_cache.should_discard_upstream() { // Serving the cached response and discarding the upstream one; nothing @@ -370,10 +371,10 @@ where #[allow(clippy::too_many_arguments)] async fn custom_bidirection_down_to_up( &self, - session: &mut Session, + session: &mut Session, client_body: &mut Box, mut rx: mpsc::Receiver, - ctx: &mut SV::CTX, + ctx: &mut >::CTX, mut upstream_custom_message_filter_rx: mpsc::Receiver<( Bytes, oneshot::Sender>, @@ -387,8 +388,8 @@ where pipe_state: Arc, ) -> Result where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { let mut cancel_downstream_reader_tx = Some(cancel_downstream_reader_tx); @@ -718,16 +719,16 @@ where async fn custom_response_filter( &self, - session: &mut Session, + session: &mut Session, mut task: HttpTask, - ctx: &mut SV::CTX, + ctx: &mut >::CTX, serve_from_cache: &mut ServeFromCache, range_body_filter: &mut RangeBodyFilter, from_cache: bool, // are the task from cache already ) -> Result where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { if !from_cache { self.upstream_filter(session, &mut task, ctx).await?; @@ -897,15 +898,15 @@ where async fn send_body_to_custom( &self, - session: &mut Session, + session: &mut Session, mut data: Option, end_of_body: bool, client_body: &mut Box, - ctx: &mut SV::CTX, + ctx: &mut >::CTX, ) -> Result where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { session .downstream_modules_ctx diff --git a/pingora-proxy/src/proxy_h1.rs b/pingora-proxy/src/proxy_h1.rs index a332f8287..3328e84a2 100644 --- a/pingora-proxy/src/proxy_h1.rs +++ b/pingora-proxy/src/proxy_h1.rs @@ -25,20 +25,21 @@ use pingora_core::protocols::http::{ v1::common::is_upgrade_req as is_h1_upgrade_req, }; -impl HttpProxy +impl HttpProxy where C: custom::Connector, + DS: DownstreamSession, { pub(crate) async fn proxy_1to1( &self, - session: &mut Session, + session: &mut Session, client_session: &mut HttpSessionV1, peer: &HttpPeer, - ctx: &mut SV::CTX, + ctx: &mut >::CTX, ) -> (bool, bool, Option>) where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { client_session.read_timeout = peer.options.read_timeout; client_session.write_timeout = peer.options.write_timeout; @@ -193,16 +194,16 @@ where pub(crate) async fn proxy_to_h1_upstream( &self, - session: &mut Session, + session: &mut Session, client_session: &mut HttpSessionV1, reused: bool, peer: &HttpPeer, - ctx: &mut SV::CTX, + ctx: &mut >::CTX, ) -> (bool, bool, Option>) // (reuse_server, reuse_client, error) where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { #[cfg(windows)] let raw = client_session.id() as std::os::windows::io::RawSocket; @@ -253,8 +254,8 @@ where pipe_state: Arc, ) -> Result where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { let mut request_done = false; let mut response_done = false; @@ -362,8 +363,8 @@ where #[allow(clippy::too_many_arguments)] async fn process_upstream_tasks( &self, - session: &mut Session, - ctx: &mut SV::CTX, + session: &mut Session, + ctx: &mut >::CTX, initial_task: HttpTask, rx: &mut mpsc::Receiver, serve_from_cache: &mut ServeFromCache, @@ -371,8 +372,8 @@ where response_state: &mut ResponseStateMachine, ) -> Result> where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { if serve_from_cache.should_discard_upstream() { // Serving the cached response and discarding the upstream one; nothing @@ -443,10 +444,10 @@ where #[allow(clippy::too_many_arguments)] async fn proxy_handle_downstream( &self, - session: &mut Session, + session: &mut Session, tx: mpsc::Sender, mut rx: mpsc::Receiver, - ctx: &mut SV::CTX, + ctx: &mut >::CTX, downstream_custom_message_writer: &mut Option>, downstream_custom_message_reader: &mut Option< Box> + Unpin + Send + Sync + 'static>, @@ -454,8 +455,8 @@ where pipe_state: Arc, ) -> Result where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { // setup custom message forwarding, if downstream supports it let ( @@ -853,16 +854,16 @@ where async fn h1_response_filter( &self, - session: &mut Session, + session: &mut Session, mut task: HttpTask, - ctx: &mut SV::CTX, + ctx: &mut >::CTX, serve_from_cache: &mut ServeFromCache, range_body_filter: &mut RangeBodyFilter, from_cache: bool, // are the task from cache already ) -> Result where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { // skip caching if already served from cache if !from_cache { @@ -1021,15 +1022,15 @@ where // TODO:: use this function to replace send_body_to2 async fn send_body_to_pipe( &self, - session: &mut Session, + session: &mut Session, mut data: Option, end_of_body: bool, tx: mpsc::Permit<'_, HttpTask>, - ctx: &mut SV::CTX, + ctx: &mut >::CTX, ) -> Result where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { // None: end of body // this var is to signal if downstream finish sending the body, which shouldn't be diff --git a/pingora-proxy/src/proxy_h2.rs b/pingora-proxy/src/proxy_h2.rs index 0b13d8a44..ed0c6f706 100644 --- a/pingora-proxy/src/proxy_h2.rs +++ b/pingora-proxy/src/proxy_h2.rs @@ -114,21 +114,22 @@ fn update_h2_scheme_authority( } } -impl HttpProxy +impl HttpProxy where C: custom::Connector, + DS: DownstreamSession, { pub(crate) async fn proxy_down_to_up( &self, - session: &mut Session, + session: &mut Session, client_session: &mut Http2Session, peer: &HttpPeer, - ctx: &mut SV::CTX, + ctx: &mut >::CTX, ) -> (bool, Option>) // (reuse_server, error) where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { let mut req = session.req_header().clone(); let authority_policy = AuthorityPolicy::from(session.downstream_session.is_custom()); @@ -363,15 +364,15 @@ where pub(crate) async fn proxy_to_h2_upstream( &self, - session: &mut Session, + session: &mut Session, client_session: &mut Http2Session, reused: bool, peer: &HttpPeer, - ctx: &mut SV::CTX, + ctx: &mut >::CTX, ) -> (bool, Option>) where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { #[cfg(windows)] let raw = client_session.fd() as std::os::windows::io::RawSocket; @@ -402,8 +403,8 @@ where #[allow(clippy::too_many_arguments)] async fn process_upstream_tasks_h2( &self, - session: &mut Session, - ctx: &mut SV::CTX, + session: &mut Session, + ctx: &mut >::CTX, initial_task: HttpTask, rx: &mut mpsc::Receiver, serve_from_cache: &mut ServeFromCache, @@ -411,8 +412,8 @@ where response_state: &mut ResponseStateMachine, ) -> Result> where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { if serve_from_cache.should_discard_upstream() { // Serving the cached response and discarding the upstream one; nothing @@ -488,10 +489,10 @@ where #[allow(clippy::too_many_arguments)] async fn bidirection_down_to_up( &self, - session: &mut Session, + session: &mut Session, client_body: &mut h2::SendStream, mut rx: mpsc::Receiver, - ctx: &mut SV::CTX, + ctx: &mut >::CTX, write_timeout: Option, downstream_custom_message_writer: &mut Option>, downstream_custom_message_reader: &mut Option< @@ -500,8 +501,8 @@ where pipe_state: Arc, ) -> Result where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { // setup custom message forwarding, if downstream supports it let ( @@ -878,16 +879,16 @@ where async fn h2_response_filter( &self, - session: &mut Session, + session: &mut Session, mut task: HttpTask, - ctx: &mut SV::CTX, + ctx: &mut >::CTX, serve_from_cache: &mut ServeFromCache, range_body_filter: &mut RangeBodyFilter, from_cache: bool, // are the task from cache already ) -> Result where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { if !from_cache { if let Some(duration) = self.upstream_filter(session, &mut task, ctx).await? { @@ -1040,16 +1041,16 @@ where async fn send_body_to2( &self, - session: &mut Session, + session: &mut Session, mut data: Option, end_of_body: bool, client_body: &mut h2::SendStream, - ctx: &mut SV::CTX, + ctx: &mut >::CTX, write_timeout: Option, ) -> Result where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { session .downstream_modules_ctx diff --git a/pingora-proxy/src/proxy_purge.rs b/pingora-proxy/src/proxy_purge.rs index 81218a074..719a4c0de 100644 --- a/pingora-proxy/src/proxy_purge.rs +++ b/pingora-proxy/src/proxy_purge.rs @@ -60,18 +60,19 @@ static NOT_PURGEABLE: Lazy = Lazy::new(|| gen_purge_response(405 // on cache storage or proxy error static INTERNAL_ERROR: Lazy = Lazy::new(|| error_resp::gen_error_response(500)); -impl HttpProxy +impl HttpProxy where C: custom::Connector, + DS: DownstreamSession, { pub(crate) async fn proxy_purge( &self, - session: &mut Session, - ctx: &mut SV::CTX, + session: &mut Session, + ctx: &mut >::CTX, ) -> Option<(bool, Option>)> where - SV: ProxyHttp + Send + Sync, - SV::CTX: Send + Sync, + SV: ProxyHttp + Send + Sync, + >::CTX: Send + Sync, { let purge_status = if session.cache.enabled() { let purged = match self.inner.purge_action(session, ctx) { diff --git a/pingora-proxy/src/proxy_trait.rs b/pingora-proxy/src/proxy_trait.rs index 5e62a4ba1..e33572c3d 100644 --- a/pingora-proxy/src/proxy_trait.rs +++ b/pingora-proxy/src/proxy_trait.rs @@ -45,7 +45,10 @@ pub enum ProxyWarnLogContext { /// /// If any of the filters returns [Result::Err], the request will fail, and the error will be logged. #[cfg_attr(not(doc_async_trait), async_trait)] -pub trait ProxyHttp { +pub trait ProxyHttp +where + DS: DownstreamSession, +{ /// The per request object to share state across the different filters type CTX; @@ -58,7 +61,7 @@ pub trait ProxyHttp { /// be forwarded to. async fn upstream_peer( &self, - session: &mut Session, + session: &mut Session, ctx: &mut Self::CTX, ) -> Result>; @@ -99,7 +102,7 @@ pub trait ProxyHttp { /// the proxy would exit. The proxy continues to the next phases when `Ok(false)` is returned. /// /// By default this filter does nothing and returns `Ok(false)`. - async fn request_filter(&self, _session: &mut Session, _ctx: &mut Self::CTX) -> Result + async fn request_filter(&self, _session: &mut Session, _ctx: &mut Self::CTX) -> Result where Self::CTX: Send + Sync, { @@ -115,7 +118,11 @@ pub trait ProxyHttp { /// Note that because this function is executed before any module that might provide access /// control or rate limiting, logic should stay in request_filter() if it can in order to be /// protected by said modules. - async fn early_request_filter(&self, _session: &mut Session, _ctx: &mut Self::CTX) -> Result<()> + async fn early_request_filter( + &self, + _session: &mut Session, + _ctx: &mut Self::CTX, + ) -> Result<()> where Self::CTX: Send + Sync, { @@ -130,7 +137,7 @@ pub trait ProxyHttp { /// /// Note that this doesn't prevent subrequests from being spawned based on the session by proxy /// core functionality, e.g. background cache revalidation requires spawning subrequests. - fn allow_spawning_subrequest(&self, _session: &Session, _ctx: &Self::CTX) -> bool + fn allow_spawning_subrequest(&self, _session: &Session, _ctx: &Self::CTX) -> bool where Self::CTX: Send + Sync, { @@ -147,7 +154,7 @@ pub trait ProxyHttp { /// who process the requests themselves. async fn request_body_filter( &self, - _session: &mut Session, + _session: &mut Session, _body: &mut Option, _end_of_stream: bool, _ctx: &mut Self::CTX, @@ -164,7 +171,7 @@ pub trait ProxyHttp { /// /// By default this filter does nothing which effectively disables caching. // Ideally only session.cache should be modified, TODO: reflect that in this interface - fn request_cache_filter(&self, _session: &mut Session, _ctx: &mut Self::CTX) -> Result<()> + fn request_cache_filter(&self, _session: &mut Session, _ctx: &mut Self::CTX) -> Result<()> where Self::CTX: Send + Sync, { @@ -187,12 +194,12 @@ pub trait ProxyHttp { /// /// The default implementation panics. You **must** override this method when /// caching is enabled. - fn cache_key_callback(&self, _session: &Session, _ctx: &mut Self::CTX) -> Result { + fn cache_key_callback(&self, _session: &Session, _ctx: &mut Self::CTX) -> Result { unimplemented!("cache_key_callback must be implemented when caching is enabled") } /// This callback is invoked when a cacheable response is ready to be admitted to cache. - fn cache_miss(&self, session: &mut Session, _ctx: &mut Self::CTX) { + fn cache_miss(&self, session: &mut Session, _ctx: &mut Self::CTX) { session.cache.cache_miss(); } @@ -207,7 +214,7 @@ pub trait ProxyHttp { /// and which kind. Returning `None` indicates no forced invalidation async fn cache_hit_filter( &self, - _session: &mut Session, + _session: &mut Session, _meta: &CacheMeta, _hit_handler: &mut HitHandler, _is_fresh: bool, @@ -231,7 +238,7 @@ pub trait ProxyHttp { /// caller's responsibility to disable keepalive or drain the request body if needed. async fn proxy_upstream_filter( &self, - _session: &mut Session, + _session: &mut Session, _ctx: &mut Self::CTX, ) -> Result where @@ -243,7 +250,7 @@ pub trait ProxyHttp { /// Decide if the response is cacheable fn response_cache_filter( &self, - _session: &Session, + _session: &Session, _resp: &ResponseHeader, _ctx: &mut Self::CTX, ) -> Result { @@ -274,7 +281,7 @@ pub trait ProxyHttp { /// be sent. fn cache_not_modified_filter( &self, - session: &Session, + session: &Session, resp: &ResponseHeader, _ctx: &mut Self::CTX, ) -> Result { @@ -298,7 +305,7 @@ pub trait ProxyHttp { /// [RFC7232]: https://www.rfc-editor.org/rfc/rfc7232 fn range_header_filter( &self, - session: &mut Session, + session: &mut Session, resp: &mut ResponseHeader, _ctx: &mut Self::CTX, ) -> range_filter::RangeType { @@ -321,7 +328,7 @@ pub trait ProxyHttp { /// `Transfer-Encoding: chunked`. async fn upstream_request_filter( &self, - _session: &mut Session, + _session: &mut Session, _upstream_request: &mut RequestHeader, _ctx: &mut Self::CTX, ) -> Result<()> @@ -353,7 +360,7 @@ pub trait ProxyHttp { #[cfg(feature = "upstream_modules")] async fn adjust_upstream_modules( &self, - _session: &mut Session, + _session: &mut Session, _upstream_response: &ResponseHeader, _end_of_stream: bool, _ctx: &mut Self::CTX, @@ -373,7 +380,7 @@ pub trait ProxyHttp { /// cached header, not served directly to downstream). async fn upstream_response_filter( &self, - _session: &mut Session, + _session: &mut Session, _upstream_response: &mut ResponseHeader, _ctx: &mut Self::CTX, ) -> Result<()> @@ -389,7 +396,7 @@ pub trait ProxyHttp { /// responses served from cache. async fn response_filter( &self, - _session: &mut Session, + _session: &mut Session, _upstream_response: &mut ResponseHeader, _ctx: &mut Self::CTX, ) -> Result<()> @@ -403,7 +410,7 @@ pub trait ProxyHttp { #[doc(hidden)] async fn custom_forwarding( &self, - _session: &mut Session, + _session: &mut Session, _ctx: &mut Self::CTX, _custom_message_to_upstream: Option>, _custom_message_to_downstream: mpsc::Sender, @@ -418,7 +425,7 @@ pub trait ProxyHttp { #[doc(hidden)] async fn downstream_custom_message_proxy_filter( &self, - _session: &mut Session, + _session: &mut Session, custom_message: Bytes, _ctx: &mut Self::CTX, _final_hop: bool, @@ -433,7 +440,7 @@ pub trait ProxyHttp { #[doc(hidden)] async fn upstream_custom_message_proxy_filter( &self, - _session: &mut Session, + _session: &mut Session, custom_message: Bytes, _ctx: &mut Self::CTX, _final_hop: bool, @@ -453,7 +460,7 @@ pub trait ProxyHttp { /// work without blocking the task processing the request. async fn upstream_response_body_filter( &self, - _session: &mut Session, + _session: &mut Session, _body: &mut Option, _end_of_stream: bool, _ctx: &mut Self::CTX, @@ -470,7 +477,7 @@ pub trait ProxyHttp { /// work without blocking the task processing the request. async fn upstream_response_trailer_filter( &self, - _session: &mut Session, + _session: &mut Session, _upstream_trailers: &mut header::HeaderMap, _ctx: &mut Self::CTX, ) -> Result<()> @@ -486,7 +493,7 @@ pub trait ProxyHttp { /// work without blocking the task processing the request. async fn response_body_filter( &self, - _session: &mut Session, + _session: &mut Session, _body: &mut Option, _end_of_stream: bool, _ctx: &mut Self::CTX, @@ -504,7 +511,7 @@ pub trait ProxyHttp { /// TODO: make this interface more intuitive async fn response_trailer_filter( &self, - _session: &mut Session, + _session: &mut Session, _upstream_trailers: &mut header::HeaderMap, _ctx: &mut Self::CTX, ) -> Result> @@ -519,7 +526,7 @@ pub trait ProxyHttp { /// /// An error log is already emitted if there is any error. This phase is used for collecting /// metrics and sending access logs. - async fn logging(&self, _session: &mut Session, _e: Option<&Error>, _ctx: &mut Self::CTX) + async fn logging(&self, _session: &mut Session, _e: Option<&Error>, _ctx: &mut Self::CTX) where Self::CTX: Send + Sync, { @@ -536,7 +543,7 @@ pub trait ProxyHttp { /// The default implementation returns `None` (no context persisted). fn persist_connection_context( &self, - _session: &Session, + _session: &Session, _ctx: &Self::CTX, ) -> Option> { None @@ -552,7 +559,7 @@ pub trait ProxyHttp { /// Use this to transfer state from the previous request into the new request's context. fn on_connection_reuse( &self, - _session: &mut Session, + _session: &mut Session, _ctx: &mut Self::CTX, _prev_ctx: Box, ) { @@ -561,7 +568,7 @@ pub trait ProxyHttp { /// A value of true means that the log message will be suppressed. The default value is false. /// /// See also: [`Self::suppress_proxy_warn_log`]. - fn suppress_error_log(&self, _session: &Session, _ctx: &Self::CTX, _error: &Error) -> bool { + fn suppress_error_log(&self, _session: &Session, _ctx: &Self::CTX, _error: &Error) -> bool { false } @@ -581,7 +588,7 @@ pub trait ProxyHttp { /// Experimental: this API may change or be removed until indicated otherwise. fn suppress_proxy_warn_log( &self, - _session: &Session, + _session: &Session, _ctx: &Self::CTX, _error: &Error, _context: ProxyWarnLogContext, @@ -602,7 +609,7 @@ pub trait ProxyHttp { fn error_while_proxy( &self, peer: &HttpPeer, - session: &mut Session, + session: &mut Session, e: Box, _ctx: &mut Self::CTX, client_reused: bool, @@ -627,7 +634,7 @@ pub trait ProxyHttp { /// available. fn fail_to_connect( &self, - _session: &mut Session, + _session: &mut Session, _peer: &HttpPeer, _ctx: &mut Self::CTX, e: Box, @@ -645,7 +652,7 @@ pub trait ProxyHttp { /// selection, and the keepalive configured on the `Session` itself still takes precedent. async fn fail_to_proxy( &self, - session: &mut Session, + session: &mut Session, e: &Error, _ctx: &mut Self::CTX, ) -> FailToProxy @@ -693,7 +700,7 @@ pub trait ProxyHttp { // 5xx HTTP status will be encoded as ErrorType::HTTPStatus(code) fn should_serve_stale( &self, - _session: &mut Session, + _session: &mut Session, _ctx: &mut Self::CTX, error: Option<&Error>, // None when it is called during stale while revalidate ) -> bool { @@ -709,7 +716,7 @@ pub trait ProxyHttp { /// This filter allows user to log timing and connection related info. async fn connected_to_upstream( &self, - _session: &mut Session, + _session: &mut Session, _reused: bool, _peer: &HttpPeer, #[cfg(unix)] _fd: std::os::unix::io::RawFd, @@ -726,7 +733,7 @@ pub trait ProxyHttp { /// This callback is invoked every time request related error log needs to be generated /// /// Users can define what is important to be written about this request via the returned string. - fn request_summary(&self, session: &Session, _ctx: &Self::CTX) -> String { + fn request_summary(&self, session: &Session, _ctx: &Self::CTX) -> String { session.as_ref().request_summary() } @@ -734,7 +741,7 @@ pub trait ProxyHttp { /// /// - `true`: this request will be used to invalidate the cache. /// - `false`: this request is a treated as a normal request - fn is_purge(&self, _session: &Session, _ctx: &Self::CTX) -> bool { + fn is_purge(&self, _session: &Session, _ctx: &Self::CTX) -> bool { false } @@ -744,7 +751,7 @@ pub trait ProxyHttp { /// Returning [`PurgeAction::Expire`] asks to keep it and mark it stale instead, so it /// revalidates against the origin rather than being refetched in full. Storage that cannot /// mark an entry stale falls back to deleting it. - fn purge_action(&self, _session: &Session, _ctx: &Self::CTX) -> PurgeAction { + fn purge_action(&self, _session: &Session, _ctx: &Self::CTX) -> PurgeAction { PurgeAction::Delete } @@ -756,7 +763,7 @@ pub trait ProxyHttp { /// If the filter returns `Err`, the proxy will instead send a 500 response. fn purge_response_filter( &self, - _session: &Session, + _session: &Session, _ctx: &mut Self::CTX, _purge_status: PurgeStatus, _purge_response: &mut std::borrow::Cow<'static, ResponseHeader>, diff --git a/pingora-proxy/src/subrequest/mod.rs b/pingora-proxy/src/subrequest/mod.rs index 8141f8c45..f987f5e67 100644 --- a/pingora-proxy/src/subrequest/mod.rs +++ b/pingora-proxy/src/subrequest/mod.rs @@ -146,11 +146,19 @@ impl Ctx { } } -use crate::HttpSession; - -pub(crate) fn create_session(parsed_session: &HttpSession) -> (HttpSession, SubrequestHandle) { +use crate::{DownstreamSession, HttpSession}; + +pub(crate) fn create_session( + parsed_session: &HttpSession, +) -> (HttpSession, SubrequestHandle) +where + DS: DownstreamSession, +{ let (session, handle) = SessionSubrequest::new_from_session(parsed_session); - (HttpSession::new_subrequest(session), handle) + ( + HttpSession::::new_subrequest_with_custom_session(session), + handle, + ) } #[tokio::test] diff --git a/pingora-proxy/src/subrequest/pipe.rs b/pingora-proxy/src/subrequest/pipe.rs index 1cc5b1a56..0317792cc 100644 --- a/pingora-proxy/src/subrequest/pipe.rs +++ b/pingora-proxy/src/subrequest/pipe.rs @@ -27,7 +27,7 @@ use crate::proxy_common::{DownstreamStateMachine, ResponseStateMachine}; use crate::subrequest::*; -use crate::{PreparedSubrequest, Session}; +use crate::{DownstreamSession, PreparedSubrequest, Session}; use bytes::Bytes; use futures::FutureExt; use log::{debug, warn}; @@ -178,14 +178,15 @@ impl std::convert::From for InputBody { } } -pub async fn pipe_subrequest( - session: &mut Session, - mut subrequest: PreparedSubrequest, +pub async fn pipe_subrequest( + session: &mut Session, + mut subrequest: PreparedSubrequest, subrequest_handle: SubrequestHandle, mut task_filter: F, input_body: InputBodyType, ) -> std::result::Result where + DS: DownstreamSession, F: FnMut(HttpTask) -> Result>, { let (maybe_preset_body, saved_body) = match input_body { @@ -375,13 +376,16 @@ where } // Mostly the same as proxy_common, but does not run proxy request_body_filter -async fn send_body_to_pipe( - session: &mut Session, +async fn send_body_to_pipe( + session: &mut Session, mut data: Option, end_of_body: bool, saved_body: Option<&mut SavedBody>, tx: mpsc::Permit<'_, HttpTask>, -) -> Result { +) -> Result +where + DS: DownstreamSession, +{ // None: end of body // this var is to signal if downstream finish sending the body, which shouldn't be // affected by the request_body_filter From 9e5a8d9901e34f2cdb7166e8583edcd6b0e6af1d Mon Sep 17 00:00:00 2001 From: ewang Date: Tue, 25 Aug 2026 11:19:15 -0700 Subject: [PATCH 4/5] add owned HTTP test origin --- .bleep | 2 +- .config/nextest.toml | 11 ++ Cargo.toml | 1 + pingora-proxy/Cargo.toml | 1 + pingora-proxy/tests/test_basic.rs | 64 +++++++- pingora-proxy/tests/test_upstream.rs | 147 ++++++++++------- pingora-proxy/tests/utils/server_utils.rs | 20 +++ pingora-test-utils/Cargo.toml | 25 +++ pingora-test-utils/src/http_origin.rs | 186 ++++++++++++++++++++++ pingora-test-utils/src/lib.rs | 18 +++ pingora-test-utils/tests/http_origin.rs | 184 +++++++++++++++++++++ 11 files changed, 592 insertions(+), 67 deletions(-) create mode 100644 .config/nextest.toml create mode 100644 pingora-test-utils/Cargo.toml create mode 100644 pingora-test-utils/src/http_origin.rs create mode 100644 pingora-test-utils/src/lib.rs create mode 100644 pingora-test-utils/tests/http_origin.rs diff --git a/.bleep b/.bleep index aec9211d7..9c5bf4193 100644 --- a/.bleep +++ b/.bleep @@ -1 +1 @@ -75f781c5d07dc5f9e292668db1d8872043d49b83 \ No newline at end of file +958d857aa734c9e76c8a9df4c82707e07236e597 \ No newline at end of file diff --git a/.config/nextest.toml b/.config/nextest.toml new file mode 100644 index 000000000..a16f4dff2 --- /dev/null +++ b/.config/nextest.toml @@ -0,0 +1,11 @@ +# These integration binaries bind fixed ports in each nextest process. +[test-groups] +pingora-fixed-port-integration = { max-threads = 1 } + +[[profile.default.overrides]] +filter = 'package(pingora-core) & binary(test_basic)' +test-group = 'pingora-fixed-port-integration' + +[[profile.default.overrides]] +filter = 'package(pingora-proxy) & (binary(test_basic) | binary(test_upstream))' +test-group = 'pingora-fixed-port-integration' diff --git a/Cargo.toml b/Cargo.toml index f75b6d365..a2a9858cf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,6 +27,7 @@ members = [ "pingora-memory-cache", "pingora-prometheus", "pingora-foundations", + "pingora-test-utils", "tinyufo", ] diff --git a/pingora-proxy/Cargo.toml b/pingora-proxy/Cargo.toml index e3fb582c6..3b33e1d04 100644 --- a/pingora-proxy/Cargo.toml +++ b/pingora-proxy/Cargo.toml @@ -55,6 +55,7 @@ tokio-tungstenite = "0.26" pingora-limits = { version = "0.9.0", path = "../pingora-limits" } pingora-load-balancing = { version = "0.9.0", path = "../pingora-load-balancing", default-features=false } pingora-prometheus = { version = "0.9.0", path = "../pingora-prometheus" } +pingora-test-utils = { version = "0.9.0", path = "../pingora-test-utils" } prometheus = "0" futures-util = "0.3" serde = { version = "1.0", features = ["derive"] } diff --git a/pingora-proxy/tests/test_basic.rs b/pingora-proxy/tests/test_basic.rs index 3cbe9bcdc..0e8e27886 100644 --- a/pingora-proxy/tests/test_basic.rs +++ b/pingora-proxy/tests/test_basic.rs @@ -16,17 +16,18 @@ mod utils; use bytes::Bytes; use h2::client; -use http::Request; +use http::{Request, Response}; use http_body_util::BodyExt; use hyper_util::client::legacy::Client; #[cfg(unix)] use hyperlocal::{UnixClientExt, Uri}; +use pingora_test_utils::http_origin::HttpOrigin; use reqwest::{header, StatusCode}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use utils::server_utils::{ - downstream_cache_warn_log_calls, init, reset_suppress_proxy_warn_log_calls, + downstream_cache_warn_log_calls, init, init_proxy, reset_suppress_proxy_warn_log_calls, suppress_proxy_warn_log_calls, }; @@ -47,8 +48,23 @@ async fn test_origin_alive() { #[tokio::test] async fn test_simple_proxy() { - init(); - let res = reqwest::get("http://127.0.0.1:6147").await.unwrap(); + init_proxy().await; + let origin = HttpOrigin::bind(|_request| async { + Response::builder() + .header(header::CONTENT_LENGTH, "13") + .body(Bytes::from_static(b"Hello World!\n")) + .unwrap() + }) + .await + .unwrap(); + let origin_addr = origin.addr(); + + let res = reqwest::Client::new() + .get("http://127.0.0.1:6147") + .header("x-port", origin_addr.port().to_string()) + .send() + .await + .unwrap(); assert_eq!(res.status(), StatusCode::OK); let headers = res.headers(); @@ -62,7 +78,10 @@ async fn test_simple_proxy() { assert_eq!(sockaddr.ip().to_string(), "127.0.0.1"); assert!(is_specified_port(sockaddr.port())); - assert_eq!(headers["x-upstream-server-addr"], "127.0.0.1:8000"); + assert_eq!( + headers["x-upstream-server-addr"].to_str().unwrap(), + origin_addr.to_string() + ); let sockaddr = headers["x-upstream-client-addr"] .to_str() .unwrap() @@ -73,6 +92,7 @@ async fn test_simple_proxy() { let body = res.text().await.unwrap(); assert_eq!(body, "Hello World!\n"); + origin.shutdown().await; } #[tokio::test] @@ -693,13 +713,23 @@ async fn test_upstream_compression() { #[tokio::test] async fn test_downstream_compression() { - init(); + init_proxy().await; + let origin = HttpOrigin::bind(|_request| async { + Response::builder() + .header(header::CONTENT_TYPE, "text/plain") + .body(Bytes::from_static(b"BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB")) + .unwrap() + }) + .await + .unwrap(); + let origin_port = origin.addr().port().to_string(); // 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:6147/no_compression") + .header("x-port", &origin_port) // tell the test proxy to use downstream compression module instead of upstream .header("x-downstream-compression", "1") .header("accept-encoding", "gzip") @@ -715,6 +745,7 @@ async fn test_downstream_compression() { let client = reqwest::ClientBuilder::new().gzip(true).build().unwrap(); let res = client .get("http://127.0.0.1:6147/no_compression") + .header("x-port", &origin_port) .header("accept-encoding", "gzip") .send() .await @@ -722,15 +753,30 @@ async fn test_downstream_compression() { assert_eq!(res.status(), StatusCode::OK); let body = res.bytes().await.unwrap(); assert_eq!(body.as_ref(), &[b'B'; 32]); + origin.shutdown().await; } #[tokio::test] async fn test_connect_close() { - init(); + init_proxy().await; + let origin = HttpOrigin::bind(|_request| async { + Response::builder() + .header(header::CONTENT_LENGTH, "13") + .body(Bytes::from_static(b"Hello World!\n")) + .unwrap() + }) + .await + .unwrap(); + let origin_port = origin.addr().port().to_string(); // default keep-alive let client = reqwest::ClientBuilder::new().build().unwrap(); - let res = client.get("http://127.0.0.1:6147").send().await.unwrap(); + let res = client + .get("http://127.0.0.1:6147") + .header("x-port", &origin_port) + .send() + .await + .unwrap(); assert_eq!(res.status(), StatusCode::OK); let headers = res.headers(); assert_eq!(headers[header::CONTENT_LENGTH], "13"); @@ -742,6 +788,7 @@ async fn test_connect_close() { let client = reqwest::ClientBuilder::new().build().unwrap(); let res = client .get("http://127.0.0.1:6147") + .header("x-port", &origin_port) .header("connection", "close") .send() .await @@ -752,6 +799,7 @@ async fn test_connect_close() { assert_eq!(headers[header::CONNECTION], "close"); let body = res.text().await.unwrap(); assert_eq!(body, "Hello World!\n"); + origin.shutdown().await; } #[tokio::test] diff --git a/pingora-proxy/tests/test_upstream.rs b/pingora-proxy/tests/test_upstream.rs index 19443b995..17f7cd26c 100644 --- a/pingora-proxy/tests/test_upstream.rs +++ b/pingora-proxy/tests/test_upstream.rs @@ -14,7 +14,7 @@ mod utils; -use utils::server_utils::init; +use utils::server_utils::{init, init_proxy}; use utils::websocket::{WS_ECHO, WS_ECHO_RAW}; use bytes::Bytes; @@ -22,6 +22,7 @@ use futures::{SinkExt, StreamExt}; use http::header::{HeaderName, HeaderValue}; use http_body_util::BodyExt; use pingora_http::ResponseHeader; +use pingora_test_utils::http_origin::HttpOrigin; use reqwest::{StatusCode, Version}; use std::sync::{ atomic::{AtomicUsize, Ordering}, @@ -30,7 +31,7 @@ use std::sync::{ use std::time::{Duration, Instant}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; -use tokio::sync::oneshot; +use tokio::sync::{mpsc, oneshot}; use tokio::time::timeout; use tokio_tungstenite::tungstenite::{client::IntoClientRequest, Message}; @@ -386,6 +387,30 @@ async fn capture_upstream_request( (port, rx) } +async fn capture_http_origin() -> (HttpOrigin, mpsc::UnboundedReceiver>) { + let (tx, rx) = mpsc::unbounded_channel(); + let origin = HttpOrigin::bind(move |request| { + let tx = tx.clone(); + async move { + let _ = tx.send(request); + http::Response::new(Bytes::new()) + } + }) + .await + .unwrap(); + + (origin, rx) +} + +async fn receive_http_origin_request( + received: &mut mpsc::UnboundedReceiver>, +) -> http::Request { + timeout(Duration::from_secs(5), received.recv()) + .await + .expect("upstream was never contacted") + .expect("origin stopped before receiving the request") +} + // Send h2c with independently controlled `:authority` and `Host`. // A direct h2 client is required because higher-level clients make them agree. async fn send_h2c_authority_request( @@ -722,7 +747,8 @@ async fn test_h1_absolute_form_host_override_rewrites_target_authority() { #[tokio::test] async fn test_h1_absolute_form_host_is_restored_after_filter_deletion() { - init(); + init_proxy().await; + let (origin, mut received) = capture_http_origin().await; let request = concat!( "GET http://client.example/test HTTP/1.1\r\n", @@ -731,13 +757,13 @@ async fn test_h1_absolute_form_host_is_restored_after_filter_deletion() { "x-port: {port}\r\n", "\r\n", ); - let (port, received) = capture_h1_upstream().await; + let port = origin.addr().port(); assert!(send_h1_raw_request(port, request).await.contains("200 OK")); - let received = String::from_utf8(received.await.unwrap()) - .unwrap() - .to_ascii_lowercase(); - assert!(received.starts_with("get http://client.example/test http/1.1\r\n")); - assert!(received.contains("\r\nhost: client.example\r\n")); + let upstream = receive_http_origin_request(&mut received).await; + origin.shutdown().await; + + assert_eq!(upstream.uri().to_string(), "http://client.example/test"); + assert_eq!(upstream.headers()[http::header::HOST], "client.example"); } #[tokio::test] @@ -827,17 +853,18 @@ async fn test_hostless_h1_absolute_form_uses_target_authority() { #[tokio::test] async fn test_hostless_http10_request_remains_hostless_upstream() { - init(); + init_proxy().await; + let (origin, mut received) = capture_http_origin().await; - let (port, received) = capture_h1_upstream().await; + let port = origin.addr().port(); let request = "GET /test HTTP/1.0\r\nx-port: {port}\r\n\r\n"; assert!(send_h1_raw_request(port, request).await.contains("200 OK")); + let upstream = receive_http_origin_request(&mut received).await; + origin.shutdown().await; - let received = String::from_utf8(received.await.unwrap()) - .unwrap() - .to_ascii_lowercase(); - assert!(received.starts_with("get /test http/1.1\r\n")); - assert!(!received.contains("\r\nhost:")); + assert_eq!(upstream.uri(), "/test"); + assert_eq!(upstream.version(), http::Version::HTTP_11); + assert!(!upstream.headers().contains_key(http::header::HOST)); } // Ambiguous H1 authority must not reach upstream. @@ -952,10 +979,9 @@ async fn send_raw_request_to_test_proxy(request: String) -> ResponseHeader { #[tokio::test] async fn test_h1_upstream_strips_hop_by_hop_and_connection_nominated_headers() { - init(); - let (port, received) = - capture_upstream_request(b"\r\n\r\n", b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n") - .await; + init_proxy().await; + let (origin, mut received) = capture_http_origin().await; + let port = origin.addr().port(); let req = format!( concat!( @@ -977,23 +1003,30 @@ async fn test_h1_upstream_strips_hop_by_hop_and_connection_nominated_headers() { ); assert_eq!(send_raw_request_to_test_proxy(req).await.status, 200); - let upstream = String::from_utf8(received.await.unwrap()) - .unwrap() - .to_ascii_lowercase(); - assert!(!upstream.contains("\r\nconnection:")); - assert!(!upstream.contains("\r\nkeep-alive:")); - assert!(!upstream.contains("\r\nproxy-connection:")); - assert!(!upstream.contains("\r\nproxy-authenticate:")); - assert!(!upstream.contains("\r\nproxy-authorization:")); - assert!(!upstream.contains("\r\nte:")); - assert!(!upstream.contains("\r\ntrailer:")); - assert!(!upstream.contains("\r\nx-private-hop:")); - assert!(upstream.contains("\r\nx-regular: keep\r\n")); + let upstream = receive_http_origin_request(&mut received).await; + origin.shutdown().await; + + for removed in [ + "connection", + "keep-alive", + "proxy-connection", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "x-private-hop", + ] { + assert!( + !upstream.headers().contains_key(removed), + "{removed} reached the upstream" + ); + } + assert_eq!(upstream.headers()["x-regular"], "keep"); } #[tokio::test] async fn test_h1_upstream_rejects_sensitive_fields_nominated_by_connection() { - init(); + init_proxy().await; for nominated in [ "Host", "X-Forwarded-For", @@ -1038,7 +1071,7 @@ async fn test_h1_upstream_rejects_sensitive_fields_nominated_by_connection() { #[tokio::test] async fn test_h1_upstream_rejects_excessive_connection_nominations() { - init(); + init_proxy().await; let (port, received) = capture_upstream_request(b"\r\n\r\n", b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n") .await; @@ -1503,10 +1536,9 @@ async fn test_h1_upstream_finishes_chunked_body_when_filter_discards_body() { #[tokio::test] async fn test_h1_upstream_can_retain_connection_nominated_fields_separately() { - init(); - let (port, received) = - capture_upstream_request(b"\r\n\r\n", b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n") - .await; + init_proxy().await; + let (origin, mut received) = capture_http_origin().await; + let port = origin.addr().port(); let req = format!( concat!( @@ -1522,19 +1554,18 @@ async fn test_h1_upstream_can_retain_connection_nominated_fields_separately() { ); assert_eq!(send_raw_request_to_test_proxy(req).await.status, 200); - let upstream = String::from_utf8(received.await.unwrap()) - .unwrap() - .to_ascii_lowercase(); - assert!(!upstream.contains("\r\nconnection:")); - assert!(upstream.contains("\r\nx-private-hop: retained\r\n")); + let upstream = receive_http_origin_request(&mut received).await; + origin.shutdown().await; + + assert!(!upstream.headers().contains_key(http::header::CONNECTION)); + assert_eq!(upstream.headers()["x-private-hop"], "retained"); } #[tokio::test] async fn test_h1_upstream_does_not_validate_nominations_when_removal_disabled() { - init(); - let (port, received) = - capture_upstream_request(b"\r\n\r\n", b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n") - .await; + init_proxy().await; + let (origin, mut received) = capture_http_origin().await; + let port = origin.addr().port(); let req = format!( concat!( @@ -1549,18 +1580,17 @@ async fn test_h1_upstream_does_not_validate_nominations_when_removal_disabled() ); assert_eq!(send_raw_request_to_test_proxy(req).await.status, 200); - let upstream = String::from_utf8(received.await.unwrap()) - .unwrap() - .to_ascii_lowercase(); - assert!(upstream.contains("\r\nhost: intended.example\r\n")); + let upstream = receive_http_origin_request(&mut received).await; + origin.shutdown().await; + + assert_eq!(upstream.headers()[http::header::HOST], "intended.example"); } #[tokio::test] async fn test_h1_upstream_always_sends_http11_request_version() { - init(); - let (port, received) = - capture_upstream_request(b"\r\n\r\n", b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n") - .await; + init_proxy().await; + let (origin, mut received) = capture_http_origin().await; + let port = origin.addr().port(); let req = format!( concat!( @@ -1573,9 +1603,10 @@ async fn test_h1_upstream_always_sends_http11_request_version() { ); assert_eq!(send_raw_request_to_test_proxy(req).await.status, 200); - assert!(String::from_utf8(received.await.unwrap()) - .unwrap() - .starts_with("GET / HTTP/1.1\r\n")); + let upstream = receive_http_origin_request(&mut received).await; + origin.shutdown().await; + + assert_eq!(upstream.version(), http::Version::HTTP_11); } #[tokio::test] diff --git a/pingora-proxy/tests/utils/server_utils.rs b/pingora-proxy/tests/utils/server_utils.rs index 67ca1f4a5..0d00805f3 100644 --- a/pingora-proxy/tests/utils/server_utils.rs +++ b/pingora-proxy/tests/utils/server_utils.rs @@ -1127,6 +1127,18 @@ impl Server { let server_handle = thread::spawn(|| { test_main(); }); + let addr = "127.0.0.1:6147".parse().unwrap(); + let deadline = std::time::Instant::now() + Duration::from_secs(10); + loop { + if std::net::TcpStream::connect_timeout(&addr, Duration::from_millis(100)).is_ok() { + break; + } + assert!( + std::time::Instant::now() < deadline, + "test proxy failed to start within 10s" + ); + thread::sleep(Duration::from_millis(50)); + } Server { handle: server_handle, } @@ -1221,6 +1233,14 @@ pub static TEST_SERVER: Lazy = Lazy::new(Server::start); pub static TEST_PSK_TLS_SERVER: Lazy = Lazy::new(PskTlsServer::start); use super::mock_origin::MOCK_ORIGIN; +pub async fn init_proxy() { + tokio::task::spawn_blocking(|| { + let _ = *TEST_SERVER; + }) + .await + .expect("test proxy startup task panicked"); +} + pub fn init() { let _ = *TEST_SERVER; let _ = *MOCK_ORIGIN; diff --git a/pingora-test-utils/Cargo.toml b/pingora-test-utils/Cargo.toml new file mode 100644 index 000000000..6beb84428 --- /dev/null +++ b/pingora-test-utils/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "pingora-test-utils" +version = "0.9.0" +authors = ["Pingora Team at Cloudflare "] +license = "Apache-2.0" +edition = "2021" +repository = "https://github.com/cloudflare/pingora" +description = """ +Test utilities for Pingora services and libraries. +""" + +[lints] +workspace = true + +[dependencies] +bytes = { workspace = true } +http = { workspace = true } +http-body-util = "0.1" +hyper = { version = "1", features = ["http1", "server"] } +hyper-util = { version = "0.1", features = ["tokio"] } +tokio = { workspace = true, features = ["macros", "net", "rt", "sync"] } + +[dev-dependencies] +reqwest = { version = "0.12", features = ["rustls-tls"], default-features = false } +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time"] } diff --git a/pingora-test-utils/src/http_origin.rs b/pingora-test-utils/src/http_origin.rs new file mode 100644 index 000000000..08ec30c68 --- /dev/null +++ b/pingora-test-utils/src/http_origin.rs @@ -0,0 +1,186 @@ +// 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. + +//! An owned HTTP/1 origin for integration tests. + +use bytes::Bytes; +use http::{Request, Response}; +use http_body_util::{BodyExt, Full}; +use hyper::body::Incoming; +use hyper::server::conn::http1; +use hyper::service::service_fn; +use hyper_util::rt::TokioIo; +use std::future::Future; +use std::io; +use std::net::{Ipv4Addr, SocketAddr}; +use std::panic; +use std::sync::Arc; +use tokio::net::{TcpListener, TcpStream}; +use tokio::sync::oneshot; +use tokio::task::{JoinHandle, JoinSet}; + +/// An HTTP/1 origin whose listener and connection tasks are owned by this value. +/// +/// The origin binds to an ephemeral localhost port by default. Dropping it +/// immediately stops the listener and all active connections. Use +/// [`HttpOrigin::shutdown`] when tests need to wait for cleanup or surface +/// handler and listener failures; [`Drop`] cannot await the server task. +pub struct HttpOrigin { + addr: SocketAddr, + shutdown_tx: Option>, + task: Option>, +} + +impl HttpOrigin { + /// Bind an origin to an ephemeral IPv4 localhost port. + pub async fn bind(handler: H) -> io::Result + where + H: Fn(Request) -> F + Send + Sync + 'static, + F: Future> + Send + 'static, + { + Self::bind_to(SocketAddr::from((Ipv4Addr::LOCALHOST, 0)), handler).await + } + + /// Bind an origin to `addr`. + pub async fn bind_to(addr: SocketAddr, handler: H) -> io::Result + where + H: Fn(Request) -> F + Send + Sync + 'static, + F: Future> + Send + 'static, + { + let listener = TcpListener::bind(addr).await?; + let addr = listener.local_addr()?; + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + let task = tokio::spawn(run(listener, Arc::new(handler), shutdown_rx)); + + Ok(Self { + addr, + shutdown_tx: Some(shutdown_tx), + task: Some(task), + }) + } + + /// Return the address on which the origin is listening. + pub fn addr(&self) -> SocketAddr { + self.addr + } + + /// Return the origin's base URL without a trailing slash. + pub fn url(&self) -> String { + format!("http://{}", self.addr) + } + + /// Stop the listener and active connections, then wait for cleanup. + /// + /// This method consumes the origin. If the returned future is cancelled, + /// dropping it still aborts the server task. Panics from request handlers + /// and listener accept failures are rethrown here. + pub async fn shutdown(mut self) { + self.signal_shutdown(); + if let Some(task) = self.task.as_mut() { + if let Err(error) = task.await { + if error.is_panic() { + panic::resume_unwind(error.into_panic()); + } + } + } + self.task.take(); + } + + fn signal_shutdown(&mut self) { + if let Some(shutdown_tx) = self.shutdown_tx.take() { + let _ = shutdown_tx.send(()); + } + } +} + +impl Drop for HttpOrigin { + fn drop(&mut self) { + self.signal_shutdown(); + if let Some(task) = self.task.take() { + task.abort(); + } + } +} + +async fn run(listener: TcpListener, handler: Arc, mut shutdown_rx: oneshot::Receiver<()>) +where + H: Fn(Request) -> F + Send + Sync + 'static, + F: Future> + Send + 'static, +{ + let mut connections = JoinSet::new(); + let mut connection_panic = None; + let mut accept_error = None; + + loop { + tokio::select! { + _ = &mut shutdown_rx => break, + accepted = listener.accept() => { + let (stream, _) = match accepted { + Ok(accepted) => accepted, + Err(error) => { + accept_error = Some(error); + break; + } + }; + let handler = Arc::clone(&handler); + connections.spawn(serve_connection(stream, handler)); + } + result = connections.join_next(), if !connections.is_empty() => { + if let Some(Err(error)) = result { + if error.is_panic() { + connection_panic = Some(error.into_panic()); + break; + } + } + } + } + } + + connections.abort_all(); + while let Some(result) = connections.join_next().await { + if let Err(error) = result { + if error.is_panic() && connection_panic.is_none() { + connection_panic = Some(error.into_panic()); + } + } + } + + if let Some(connection_panic) = connection_panic { + panic::resume_unwind(connection_panic); + } + if let Some(accept_error) = accept_error { + panic!("HTTP test origin failed to accept a connection: {accept_error}"); + } +} + +async fn serve_connection(stream: TcpStream, handler: Arc) +where + H: Fn(Request) -> F + Send + Sync + 'static, + F: Future> + Send + 'static, +{ + let service = service_fn(move |request: Request| { + let handler = Arc::clone(&handler); + async move { + let (parts, body) = request.into_parts(); + let body = body.collect().await?.to_bytes(); + let response = handler(Request::from_parts(parts, body)).await; + let (parts, body) = response.into_parts(); + Ok::<_, hyper::Error>(Response::from_parts(parts, Full::new(body))) + } + }); + + let _ = http1::Builder::new() + .serve_connection(TokioIo::new(stream), service) + .await; +} diff --git a/pingora-test-utils/src/lib.rs b/pingora-test-utils/src/lib.rs new file mode 100644 index 000000000..18c78c4f6 --- /dev/null +++ b/pingora-test-utils/src/lib.rs @@ -0,0 +1,18 @@ +// 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. + +#![warn(clippy::all)] +//! Test utilities for Pingora services and libraries. + +pub mod http_origin; diff --git a/pingora-test-utils/tests/http_origin.rs b/pingora-test-utils/tests/http_origin.rs new file mode 100644 index 000000000..559348199 --- /dev/null +++ b/pingora-test-utils/tests/http_origin.rs @@ -0,0 +1,184 @@ +// 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. + +use bytes::Bytes; +use http::{Request, Response, StatusCode}; +use pingora_test_utils::http_origin::HttpOrigin; +use std::future::{poll_fn, Future}; +use std::sync::Arc; +use std::task::Poll; +use std::time::Duration; +use tokio::net::TcpStream; +use tokio::sync::{Mutex, Notify}; + +#[tokio::test] +async fn serves_buffered_requests_with_shared_state() { + let requests = Arc::new(Mutex::new(Vec::new())); + let origin = HttpOrigin::bind({ + let requests = Arc::clone(&requests); + move |request: Request| { + let requests = Arc::clone(&requests); + async move { + requests.lock().await.push(( + request.method().clone(), + request.uri().clone(), + request.body().clone(), + )); + Response::builder() + .status(StatusCode::CREATED) + .header("x-origin", "rust") + .body(Bytes::from_static(b"response body")) + .unwrap() + } + } + }) + .await + .unwrap(); + + assert!(origin.addr().ip().is_loopback()); + assert_ne!(origin.addr().port(), 0); + + let response = reqwest::Client::new() + .post(format!("{}/resource?query=value", origin.url())) + .body("request body") + .send() + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::CREATED); + assert_eq!(response.headers()["x-origin"], "rust"); + assert_eq!(response.bytes().await.unwrap(), "response body"); + + let requests = requests.lock().await; + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].0, http::Method::POST); + assert_eq!(requests[0].1, "/resource?query=value"); + assert_eq!(requests[0].2, "request body"); +} + +#[tokio::test] +async fn explicit_shutdown_stops_the_listener() { + let origin = HttpOrigin::bind(ok_handler).await.unwrap(); + let addr = origin.addr(); + + TcpStream::connect(addr).await.unwrap(); + origin.shutdown().await; + + assert!(TcpStream::connect(addr).await.is_err()); +} + +#[tokio::test] +async fn shutdown_aborts_active_requests() { + let started = Arc::new(Notify::new()); + let blocked = Arc::new(Notify::new()); + let origin = HttpOrigin::bind({ + let started = Arc::clone(&started); + let blocked = Arc::clone(&blocked); + move |_request| { + let started = Arc::clone(&started); + let blocked = Arc::clone(&blocked); + async move { + started.notify_one(); + blocked.notified().await; + Response::new(Bytes::new()) + } + } + }) + .await + .unwrap(); + + let request = tokio::spawn(reqwest::get(origin.url())); + started.notified().await; + origin.shutdown().await; + + assert!(request.await.unwrap().is_err()); +} + +#[tokio::test] +async fn drop_stops_the_listener() { + let origin = HttpOrigin::bind(ok_handler).await.unwrap(); + let addr = origin.addr(); + + TcpStream::connect(addr).await.unwrap(); + drop(origin); + + wait_until_refused(addr).await; +} + +#[tokio::test] +async fn cancelling_shutdown_stops_the_listener_and_active_requests() { + let started = Arc::new(Notify::new()); + let blocked = Arc::new(Notify::new()); + let origin = HttpOrigin::bind({ + let started = Arc::clone(&started); + let blocked = Arc::clone(&blocked); + move |_request| { + let started = Arc::clone(&started); + let blocked = Arc::clone(&blocked); + async move { + started.notify_one(); + blocked.notified().await; + Response::new(Bytes::new()) + } + } + }) + .await + .unwrap(); + let addr = origin.addr(); + + let request = tokio::spawn(reqwest::get(origin.url())); + started.notified().await; + + let mut shutdown = Box::pin(origin.shutdown()); + poll_fn(|context| { + assert!(shutdown.as_mut().poll(context).is_pending()); + Poll::Ready(()) + }) + .await; + drop(shutdown); + + wait_until_refused(addr).await; + assert!(request.await.unwrap().is_err()); +} + +#[tokio::test] +#[should_panic(expected = "handler panic")] +async fn shutdown_propagates_handler_panics() { + let origin = HttpOrigin::bind(panic_handler).await.unwrap(); + + assert!(reqwest::get(origin.url()).await.is_err()); + origin.shutdown().await; +} + +async fn ok_handler(_request: Request) -> Response { + Response::new(Bytes::from_static(b"ok")) +} + +async fn panic_handler(_request: Request) -> Response { + panic!("handler panic") +} + +async fn wait_until_refused(addr: std::net::SocketAddr) { + tokio::time::timeout(Duration::from_secs(1), async { + loop { + match TcpStream::connect(addr).await { + Ok(stream) => drop(stream), + Err(_) => break, + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); +} From 64a3767475dd32d5af2022d559f7fccb479dc075 Mon Sep 17 00:00:00 2001 From: Jeremy Brinegar Date: Thu, 3 Sep 2026 08:49:14 -0500 Subject: [PATCH 5/5] Abort tls offload tasks when dropped --- .bleep | 2 +- Cargo.toml | 1 + pingora-core/Cargo.toml | 3 +- .../listeners/tls/boringssl_openssl/mod.rs | 48 +++++++++++++---- pingora-core/src/listeners/tls/rustls/mod.rs | 20 +++---- pingora-core/src/listeners/tls/s2n/mod.rs | 6 ++- pingora-core/src/offload.rs | 52 +++++++++++++++++++ 7 files changed, 108 insertions(+), 24 deletions(-) diff --git a/.bleep b/.bleep index 9c5bf4193..468a71a29 100644 --- a/.bleep +++ b/.bleep @@ -1 +1 @@ -958d857aa734c9e76c8a9df4c82707e07236e597 \ No newline at end of file +ae908d928b483c33b231680ec6060191d1e0415a \ No newline at end of file diff --git a/Cargo.toml b/Cargo.toml index a2a9858cf..37b384320 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,6 +35,7 @@ members = [ bstr = "1.12.0" tokio = "1" tokio-stream = { version = "0.1" } +tokio-util = { version = "0.7.12", features = ["rt"] } async-trait = "0.1.42" httparse = "1" bytes = "1.0" diff --git a/pingora-core/Cargo.toml b/pingora-core/Cargo.toml index 13fba01fc..519a84471 100644 --- a/pingora-core/Cargo.toml +++ b/pingora-core/Cargo.toml @@ -31,6 +31,7 @@ pingora-s2n = { version = "0.9.0", path = "../pingora-s2n", optional = true } bstr = { workspace = true } tokio = { workspace = true, features = ["net", "rt-multi-thread", "signal"] } tokio-stream = { workspace = true } +tokio-util = { workspace = true, optional = true } futures = { workspace = true } async-trait = { workspace = true } httparse = { workspace = true } @@ -123,6 +124,6 @@ rustls = ["pingora-rustls", "any_tls", "dep:x509-parser", "ouroboros"] s2n = ["pingora-s2n", "any_tls", "dep:x509-parser", "ouroboros", "lru"] patched_http1 = ["pingora-http/patched_http1"] openssl_derived = ["any_tls"] -any_tls = [] +any_tls = ["dep:tokio-util"] sentry = ["dep:sentry"] connection_filter = [] diff --git a/pingora-core/src/listeners/tls/boringssl_openssl/mod.rs b/pingora-core/src/listeners/tls/boringssl_openssl/mod.rs index f5a6ff36d..627591c86 100644 --- a/pingora-core/src/listeners/tls/boringssl_openssl/mod.rs +++ b/pingora-core/src/listeners/tls/boringssl_openssl/mod.rs @@ -185,16 +185,16 @@ impl Acceptor { if let Some(offload) = self.offload.as_ref() { let ssl_acceptor = self.ssl_acceptor.clone(); let callbacks = self.callbacks.clone(); - let rt = offload.get_runtime(stream.id() as u64); - rt.spawn(async move { - if let Some(cb) = callbacks.as_ref() { - handshake_with_callback(&ssl_acceptor, stream, cb.as_ref()).await - } else { - handshake(&ssl_acceptor, stream).await - } - }) - .await - .or_err(InternalError, "TLS offload runtime failure")? + offload + .spawn_abort_on_drop(stream.id() as u64, async move { + if let Some(cb) = callbacks.as_ref() { + handshake_with_callback(&ssl_acceptor, stream, cb.as_ref()).await + } else { + handshake(&ssl_acceptor, stream).await + } + }) + .await + .or_err(InternalError, "TLS offload runtime failure")? } else if let Some(cb) = self.callbacks.as_ref() { handshake_with_callback(&self.ssl_acceptor, stream, cb.as_ref()).await } else { @@ -268,3 +268,31 @@ mod alpn { } } } + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + use tokio::io::AsyncReadExt; + + #[tokio::test] + async fn canceled_offloaded_handshake_drops_stream() { + let cert_path = format!("{}/tests/keys/server.crt", env!("CARGO_MANIFEST_DIR")); + let key_path = format!("{}/tests/keys/key.pem", env!("CARGO_MANIFEST_DIR")); + let mut settings = TlsSettings::intermediate(&cert_path, &key_path).unwrap(); + settings.set_offload_threadpool(1, 1); + let acceptor = settings.build(); + let (mut client, server) = tokio::io::duplex(1024); + + let mut handshake = Box::pin(acceptor.tls_handshake(server)); + assert!(futures::poll!(handshake.as_mut()).is_pending()); + drop(handshake); + + let mut buf = [0]; + let bytes_read = tokio::time::timeout(Duration::from_secs(1), client.read(&mut buf)) + .await + .unwrap() + .unwrap(); + assert_eq!(bytes_read, 0); + } +} diff --git a/pingora-core/src/listeners/tls/rustls/mod.rs b/pingora-core/src/listeners/tls/rustls/mod.rs index bf838cde0..3b3deb778 100644 --- a/pingora-core/src/listeners/tls/rustls/mod.rs +++ b/pingora-core/src/listeners/tls/rustls/mod.rs @@ -251,16 +251,16 @@ impl Acceptor { offload: None, }; let callbacks = self.callbacks.clone(); - let rt = offload.get_runtime(stream.id() as u64); - rt.spawn(async move { - if let Some(cb) = callbacks.as_ref() { - handshake_with_callback(&acceptor, stream, cb.as_ref()).await - } else { - handshake(&acceptor, stream).await - } - }) - .await - .or_err(InternalError, "TLS offload runtime failure")? + offload + .spawn_abort_on_drop(stream.id() as u64, async move { + if let Some(cb) = callbacks.as_ref() { + handshake_with_callback(&acceptor, stream, cb.as_ref()).await + } else { + handshake(&acceptor, stream).await + } + }) + .await + .or_err(InternalError, "TLS offload runtime failure")? } else if let Some(cb) = self.callbacks.as_ref() { handshake_with_callback(self, stream, cb.as_ref()).await } else { diff --git a/pingora-core/src/listeners/tls/s2n/mod.rs b/pingora-core/src/listeners/tls/s2n/mod.rs index 8cc2b1ce4..cf72ea323 100644 --- a/pingora-core/src/listeners/tls/s2n/mod.rs +++ b/pingora-core/src/listeners/tls/s2n/mod.rs @@ -232,8 +232,10 @@ impl Acceptor { acceptor: self.acceptor.clone(), offload: None, }; - let rt = offload.get_runtime(stream.id() as u64); - rt.spawn(async move { handshake(&acceptor, stream).await }) + offload + .spawn_abort_on_drop(stream.id() as u64, async move { + handshake(&acceptor, stream).await + }) .await .or_err(InternalError, "TLS offload runtime failure")? } else { diff --git a/pingora-core/src/offload.rs b/pingora-core/src/offload.rs index 0db7ed9d3..90b31b17c 100644 --- a/pingora-core/src/offload.rs +++ b/pingora-core/src/offload.rs @@ -15,8 +15,12 @@ use log::debug; use once_cell::sync::OnceCell; use rand::Rng; +#[cfg(feature = "any_tls")] +use std::future::Future; use tokio::runtime::{Builder, Handle}; use tokio::sync::oneshot::{channel, Sender}; +#[cfg(feature = "any_tls")] +use tokio_util::task::AbortOnDropHandle; // NOTE: use dedicated current-thread runtimes until pingora-runtime can preserve // the lazy-after-daemonize initialization behavior below. @@ -104,4 +108,52 @@ impl OffloadRuntime { let pools = self.pools.get_or_init(|| self.init_pools()); &pools[shard * self.thread_per_shard + thread_in_shard].0 } + + /// Spawn work that is aborted when its awaiting task is canceled. + #[cfg(feature = "any_tls")] + pub fn spawn_abort_on_drop(&self, hash: u64, future: F) -> AbortOnDropHandle + where + F: Future + Send + 'static, + F::Output: Send + 'static, + { + AbortOnDropHandle::new(self.get_runtime(hash).spawn(future)) + } +} + +#[cfg(all(test, feature = "any_tls"))] +mod tests { + use super::*; + use std::{future::pending, time::Duration}; + + struct DropSignal(Option>); + + impl Drop for DropSignal { + fn drop(&mut self) { + if let Some(sender) = self.0.take() { + let _ = sender.send(()); + } + } + } + + #[tokio::test] + async fn spawned_task_is_aborted_when_handle_is_dropped() { + let offload = OffloadRuntime::new("test offload", 1, 1); + let (started_sender, started_receiver) = channel(); + let (dropped_sender, dropped_receiver) = channel(); + let task = offload.spawn_abort_on_drop(0, async move { + let _drop_signal = DropSignal(Some(dropped_sender)); + started_sender.send(()).unwrap(); + pending::<()>().await; + }); + + tokio::time::timeout(Duration::from_secs(1), started_receiver) + .await + .unwrap() + .unwrap(); + drop(task); + tokio::time::timeout(Duration::from_secs(1), dropped_receiver) + .await + .unwrap() + .unwrap(); + } }