Allow reusing session on errors prior to proxy upstream

Allow session reuse if configured in `fail_to_proxy`. Applies only to
errors prior to proxying upstream, for now.

Previously any short-circuiting error during request processing would
result in closing the connection to downstream. This is unnecessary if
there was an internal error that simply resulted in an error response
being written downstream. Note that any `set_keepalive` on the Session
itself is still respected, and this API currently does not overwrite
reuse decisions made while proxying upstream.

The user is responsible for implications of reusing the session in this
manner, such as draining any remaining request body.

Author: drcaramelsyrup

.bleep

@@ -1 +1 @@
-98d239a60a476873e1de444763a555802aa41f6c
+e6597ea67a5a6c4fcdd79d63de1e3d04a46d7a35

pingora-proxy/src/lib.rs

@@ -79,7 +79,7 @@ use subrequest::Ctx as SubReqCtx;
 
 pub use proxy_cache::range_filter::{range_header_filter, RangeType};
 pub use proxy_purge::PurgeStatus;
-pub use proxy_trait::ProxyHttp;
+pub use proxy_trait::{FailToProxy, ProxyHttp};
 
 pub mod prelude {
     pub use crate::{http_proxy_service, ProxyHttp, Session};
@@ -493,10 +493,9 @@ impl<SV> HttpProxy<SV> {
             .early_request_filter(&mut session, &mut ctx)
             .await
         {
-            self.handle_error(&mut session, &mut ctx, e, "Fail to early filter request:")
+            return self
+                .handle_error(session, &mut ctx, e, "Fail to early filter request:")
                 .await;
-            self.cleanup_sub_req(&mut session);
-            return None;
         }
 
         let req = session.downstream_session.req_header_mut();
@@ -507,15 +506,14 @@ impl<SV> HttpProxy<SV> {
             .request_header_filter(req)
             .await
         {
-            self.handle_error(
-                &mut session,
-                &mut ctx,
-                e,
-                "Failed in downstream modules request filter:",
-            )
-            .await;
-            self.cleanup_sub_req(&mut session);
-            return None;
+            return self
+                .handle_error(
+                    session,
+                    &mut ctx,
+                    e,
+                    "Failed in downstream modules request filter:",
+                )
+                .await;
         }
 
         match self.inner.request_filter(&mut session, &mut ctx).await {
@@ -529,10 +527,9 @@ impl<SV> HttpProxy<SV> {
                 /* else continue */
             }
             Err(e) => {
-                self.handle_error(&mut session, &mut ctx, e, "Fail to filter request:")
+                return self
+                    .handle_error(session, &mut ctx, e, "Fail to filter request:")
                     .await;
-                self.cleanup_sub_req(&mut session);
-                return None;
             }
         }
 
@@ -563,15 +560,14 @@ impl<SV> HttpProxy<SV> {
                         match session.write_response_header_ref(&BAD_GATEWAY).await {
                             Ok(()) => {}
                             Err(e) => {
-                                self.handle_error(
-                                    &mut session,
-                                    &mut ctx,
-                                    e,
-                                    "Error responding with Bad Gateway:",
-                                )
-                                .await;
-
-                                return None;
+                                return self
+                                    .handle_error(
+                                        session,
+                                        &mut ctx,
+                                        e,
+                                        "Error responding with Bad Gateway:",
+                                    )
+                                    .await;
                             }
                         }
                     }
@@ -585,14 +581,14 @@ impl<SV> HttpProxy<SV> {
                     session.cache.disable(NoCacheReason::InternalError);
                 }
 
-                self.handle_error(
-                    &mut session,
-                    &mut ctx,
-                    e,
-                    "Error deciding if we should proxy to upstream:",
-                )
-                .await;
-                return None;
+                return self
+                    .handle_error(
+                        session,
+                        &mut ctx,
+                        e,
+                        "Error deciding if we should proxy to upstream:",
+                    )
+                    .await;
             }
         }
 
@@ -657,14 +653,14 @@ impl<SV> HttpProxy<SV> {
                 };
                 session.cache.disable(reason);
             }
-            let status = self.inner.fail_to_proxy(&mut session, e, &mut ctx).await;
+            let res = self.inner.fail_to_proxy(&mut session, e, &mut ctx).await;
 
             // final error will have > 0 status unless downstream connection is dead
             if !self.inner.suppress_error_log(&session, &ctx, e) {
                 error!(
                     "Fail to proxy: {}, status: {}, tries: {}, retry: {}, {}",
                     final_error.as_ref().unwrap(),
-                    status,
+                    res.error_code,
                     retries,
                     false, // we never retry here
                     self.inner.request_summary(&session, &ctx)
@@ -679,23 +675,32 @@ impl<SV> HttpProxy<SV> {
 
     async fn handle_error(
         &self,
-        session: &mut Session,
+        mut session: Session,
         ctx: &mut <SV as ProxyHttp>::CTX,
         e: Box<Error>,
         context: &str,
-    ) where
+    ) -> Option<Stream>
+    where
         SV: ProxyHttp + Send + Sync + 'static,
         <SV as ProxyHttp>::CTX: Send + Sync,
     {
-        if !self.inner.suppress_error_log(session, ctx, &e) {
+        let res = self.inner.fail_to_proxy(&mut session, &e, ctx).await;
+        if !self.inner.suppress_error_log(&session, ctx, &e) {
             error!(
-                "{context} {}, {}",
+                "{context} {}, status: {}, {}",
                 e,
-                self.inner.request_summary(session, ctx)
+                res.error_code,
+                self.inner.request_summary(&session, ctx)
             );
         }
-        self.inner.fail_to_proxy(session, &e, ctx).await;
-        self.inner.logging(session, Some(&e), ctx).await;
+        self.inner.logging(&mut session, Some(&e), ctx).await;
+        self.cleanup_sub_req(&mut session);
+
+        if res.can_reuse_downstream {
+            session.downstream_session.finish().await.ok().flatten()
+        } else {
+            None
+        }
     }
 }
 

pingora-proxy/src/proxy_trait.rs

@@ -391,8 +391,16 @@ pub trait ProxyHttp {
     ///
     /// Users may write an error response to the downstream if the downstream is still writable.
     ///
-    /// The response status code of the error response maybe returned for logging purpose.
-    async fn fail_to_proxy(&self, session: &mut Session, e: &Error, _ctx: &mut Self::CTX) -> u16
+    /// The response status code of the error response may be returned for logging purposes.
+    /// Additionally, the user can return whether this session may be reused in spite of the error.
+    /// Today this reuse status is only respected for errors that occur prior to upstream peer
+    /// selection, and the keepalive configured on the `Session` itself still takes precedent.
+    async fn fail_to_proxy(
+        &self,
+        session: &mut Session,
+        e: &Error,
+        _ctx: &mut Self::CTX,
+    ) -> FailToProxy
     where
         Self::CTX: Send + Sync,
     {
@@ -419,7 +427,12 @@ pub trait ProxyHttp {
                 error!("failed to send error response to downstream: {e}");
             });
         }
-        code
+
+        FailToProxy {
+            error_code: code,
+            // default to no reuse, which is safest
+            can_reuse_downstream: false,
+        }
     }
 
     /// Decide whether should serve stale when encountering an error or during revalidation
@@ -493,3 +506,9 @@ pub trait ProxyHttp {
         Ok(())
     }
 }
+
+/// Context struct returned by `fail_to_proxy`.
+pub struct FailToProxy {
+    pub error_code: u16,
+    pub can_reuse_downstream: bool,
+}