Add method to disable keepalive if downstream is unfinished

This change adds a method for H1 servers to disable downstream keepalive
upon sending the response if the downstream is unfinished e.g. to avoid
needing to drain the body.

Author: drcaramelsyrup

.bleep

@@ -1 +1 @@
-bb2739d21ee0ff8c345a0c03c39f888cb087f311
+9c4fb3cdd246164ff27d0d1b8c4f5775cef8f1fc

pingora-core/src/protocols/http/server.rs

@@ -273,6 +273,17 @@ impl Session {
         }
     }
 
+    /// Sets whether keepalive should be disabled if response is written prior to
+    /// downstream body finishing.
+    ///
+    /// This is a noop for h2.
+    pub fn set_close_on_response_before_downstream_finish(&mut self, close: bool) {
+        match self {
+            Self::H1(s) => s.set_close_on_response_before_downstream_finish(close),
+            Self::H2(_) => {} // always ignored
+        }
+    }
+
     /// Return a digest of the request including the method, path and Host header
     // TODO: make this use a `Formatter`
     pub fn request_summary(&self) -> String {

pingora-core/src/protocols/http/v1/server.rs

@@ -79,6 +79,8 @@ pub struct HttpSession {
     min_send_rate: Option<usize>,
     /// When this is enabled informational response headers will not be proxied downstream
     ignore_info_resp: bool,
+    /// Disable keepalive if response is sent before downstream body is finished
+    close_on_response_before_downstream_finish: bool,
 }
 
 impl HttpSession {
@@ -116,6 +118,7 @@ impl HttpSession {
             digest,
             min_send_rate: None,
             ignore_info_resp: false,
+            close_on_response_before_downstream_finish: false,
         }
     }
 
@@ -469,6 +472,11 @@ impl HttpSession {
             }
         }
 
+        if self.close_on_response_before_downstream_finish && !self.is_body_done() {
+            debug!("set connection close before downstream finish");
+            self.set_keepalive(None);
+        }
+
         // no need to add these headers to 1xx responses
         if !header.status.is_informational() && self.update_resp_headers {
             /* update headers */
@@ -946,6 +954,14 @@ impl HttpSession {
         self.ignore_info_resp = ignore;
     }
 
+    /// Sets whether keepalive should be disabled if response is written prior to
+    /// downstream body finishing.
+    ///
+    /// This may be set to avoid draining downstream if the body is no longer necessary.
+    pub fn set_close_on_response_before_downstream_finish(&mut self, close: bool) {
+        self.close_on_response_before_downstream_finish = close;
+    }
+
     /// Return the [Digest] of the connection.
     pub fn digest(&self) -> &Digest {
         &self.digest

pingora-proxy/tests/test_upstream.rs

@@ -47,6 +47,8 @@ async fn test_duplex() {
         .send()
         .await
         .unwrap();
+    let headers = res.headers();
+    assert_eq!(headers["Connection"], "keep-alive");
     assert_eq!(res.status(), StatusCode::OK);
     let body = res.text().await.unwrap();
     assert_eq!(body.len(), 64 * 5);
@@ -110,6 +112,25 @@ async fn test_upload() {
     assert_eq!(body.len(), 64 * 5);
 }
 
+#[tokio::test]
+async fn test_close_on_response_before_downstream_finish() {
+    init();
+    let client = reqwest::Client::new();
+    let res = client
+        .post("http://127.0.0.1:6147/test2")
+        .header("x-close-on-response-before-downstream-finish", "1")
+        .body("b".repeat(15 * 1024 * 1024)) // 15 MB upload
+        .timeout(Duration::from_secs(5))
+        .send()
+        .await
+        .unwrap();
+    assert_eq!(res.status(), StatusCode::OK);
+    let headers = res.headers();
+    assert_eq!(headers["Connection"], "close");
+    let body = res.text().await.unwrap();
+    assert_eq!(body.len(), 11);
+}
+
 #[tokio::test]
 async fn test_ws_server_ends_conn() {
     init();

pingora-proxy/tests/utils/server_utils.rs

@@ -249,6 +249,11 @@ impl ProxyHttp for ExampleProxyHttp {
             .get("x-min-rate")
             .and_then(|v| v.to_str().ok().and_then(|v| v.parse().ok()));
 
+        let close_on_response_before_downstream_finish = req
+            .headers
+            .get("x-close-on-response-before-downstream-finish")
+            .is_some();
+
         let downstream_compression = req.headers.get("x-downstream-compression").is_some();
         if !downstream_compression {
             // enable upstream compression for all requests by default
@@ -263,6 +268,9 @@ impl ProxyHttp for ExampleProxyHttp {
 
         session.set_min_send_rate(min_rate);
         session.set_write_timeout(write_timeout.map(Duration::from_secs));
+        session.set_close_on_response_before_downstream_finish(
+            close_on_response_before_downstream_finish,
+        );
 
         Ok(false)
     }