Add a write timeout to write body buf and an option to set a minimum send rate

3 files changed, 82 insertions, 0 deletions

diff --git a/pingora-proxy/tests/test_upstream.rs b/pingora-proxy/tests/test_upstream.rs
index ad7aaaf..2600832 100644
--- a/pingora-proxy/tests/test_upstream.rs
+++ b/pingora-proxy/tests/test_upstream.rs
@@ -133,6 +133,60 @@ async fn test_ws_server_ends_conn() {
     assert!(ws_stream.next().await.is_none());
 }
 
+#[tokio::test]
+async fn test_download_timeout() {
+    init();
+    use hyper::body::HttpBody;
+    use tokio::time::sleep;
+
+    let client = hyper::Client::new();
+    let uri: hyper::Uri = "http://127.0.0.1:6147/download/".parse().unwrap();
+    let req = hyper::Request::builder()
+        .uri(uri)
+        .header("x-write-timeout", "1")
+        .body(hyper::Body::empty())
+        .unwrap();
+    let mut res = client.request(req).await.unwrap();
+    assert_eq!(res.status(), StatusCode::OK);
+
+    let mut err = false;
+    sleep(Duration::from_secs(2)).await;
+    while let Some(chunk) = res.body_mut().data().await {
+        if chunk.is_err() {
+            err = true;
+        }
+    }
+    assert!(err);
+}
+
+#[tokio::test]
+async fn test_download_timeout_min_rate() {
+    init();
+    use hyper::body::HttpBody;
+    use tokio::time::sleep;
+
+    let client = hyper::Client::new();
+    let uri: hyper::Uri = "http://127.0.0.1:6147/download/".parse().unwrap();
+    let req = hyper::Request::builder()
+        .uri(uri)
+        .header("x-write-timeout", "1")
+        .header("x-min-rate", "10000")
+        .body(hyper::Body::empty())
+        .unwrap();
+    let mut res = client.request(req).await.unwrap();
+    assert_eq!(res.status(), StatusCode::OK);
+
+    let mut err = false;
+    sleep(Duration::from_secs(2)).await;
+    while let Some(chunk) = res.body_mut().data().await {
+        if chunk.is_err() {
+            err = true;
+        }
+    }
+    // no error as write timeout is overridden by min rate
+    assert!(!err);
+}
+
 mod test_cache {
     use super::*;
     use std::str::FromStr;
diff --git a/pingora-proxy/tests/utils/conf/origin/conf/nginx.conf b/pingora-proxy/tests/utils/conf/origin/conf/nginx.conf
index 07e57e8..a41a743 100644
--- a/pingora-proxy/tests/utils/conf/origin/conf/nginx.conf
+++ b/pingora-proxy/tests/utils/conf/origin/conf/nginx.conf
@@ -296,6 +296,15 @@ http {
             }
         }
 
+        location /download/ {
+            content_by_lua_block {
+                ngx.req.read_body()
+                local body = string.rep("A", 4194304)
+                ngx.header["Content-Length"] = #body
+                ngx.print(body)
+            }
+        }
+
         location /tls_verify {
             keepalive_timeout 0;
             return 200;
diff --git a/pingora-proxy/tests/utils/server_utils.rs b/pingora-proxy/tests/utils/server_utils.rs
index 7b20630..ec1a962 100644
--- a/pingora-proxy/tests/utils/server_utils.rs
+++ b/pingora-proxy/tests/utils/server_utils.rs
@@ -39,6 +39,7 @@ use pingora_proxy::{ProxyHttp, Session};
 use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 use std::thread;
+use std::time::Duration;
 
 pub struct ExampleProxyHttps {}
 
@@ -230,6 +231,17 @@ impl ProxyHttp for ExampleProxyHttp {
 
     async fn request_filter(&self, session: &mut Session, _ctx: &mut Self::CTX) -> Result<bool> {
         let req = session.req_header();
+
+        let write_timeout = req
+            .headers
+            .get("x-write-timeout")
+            .and_then(|v| v.to_str().ok().and_then(|v| v.parse().ok()));
+
+        let min_rate = req
+            .headers
+            .get("x-min-rate")
+            .and_then(|v| v.to_str().ok().and_then(|v| v.parse().ok()));
+
         let downstream_compression = req.headers.get("x-downstream-compression").is_some();
         if !downstream_compression {
             // enable upstream compression for all requests by default
@@ -242,6 +254,13 @@ impl ProxyHttp for ExampleProxyHttp {
                 .adjust_level(0);
         }
 
+        if let Some(min_rate) = min_rate {
+            session.set_min_send_rate(min_rate);
+        }
+        if let Some(write_timeout) = write_timeout {
+            session.set_write_timeout(Duration::from_secs(write_timeout));
+        }
+
         Ok(false)
     }