Sync internal, 2026-09-11

.bleep

@@ -1 +1 @@
-3615b9f6e16a0a0c2d7523c8e8b812d1681dcf10
+ae908d928b483c33b231680ec6060191d1e0415a

.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'

Cargo.toml

@@ -27,13 +27,15 @@ members = [
     "pingora-memory-cache",
     "pingora-prometheus",
     "pingora-foundations",
+    "pingora-test-utils",
     "tinyufo",
 ]
 
 [workspace.dependencies]
 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"

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<C: CachePut> CachePutCtx<C> {
     /// 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<DS: DownstreamSession>(
         &mut self,
-        session: &mut ServerSession,
+        session: &mut ServerSession<DS>,
     ) -> Result<Option<NoCacheReason>> {
         let mut no_cache_reason = None;
         while let Some(data) = session.read_request_body().await? {

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 = []

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<DS = ()>
+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<CS>(session: &ServerSession<CS>) -> 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<CS>(self, session: &mut ServerSession<CS>)
+    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<DS = ()>
+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<Self>,
-        mut session: ServerSession,
+        mut session: ServerSession<DS>,
         // TODO: make this ShutdownWatch so that all task can await on this event
         shutdown: &ShutdownWatch,
     ) -> Option<ReusedHttpStream>;
@@ -258,9 +271,10 @@ pub trait HttpServerApp {
 }
 
 #[async_trait]
-impl<T> ServerApp for T
+impl<T, DS> ServerApp<DS> for T
 where
-    T: HttpServerApp + Send + Sync + 'static,
+    T: HttpServerApp<DS> + Send + Sync + 'static,
+    DS: CustomServerSession,
 {
     async fn process_new(
         self: &Arc<Self>,
@@ -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::<DS>::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::<DS>::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::<DS>::new_http1_with_custom_session(stream);
                 if let Some(persistent_settings) = persistent_settings {
                     persistent_settings.apply_to_session(&mut session);
                 }

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);
+    }
+}

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 {

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 {

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<F>(&self, hash: u64, future: F) -> AbortOnDropHandle<F::Output>
+    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<Sender<()>>);
+
+    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();
+    }
 }