Flush already received data if upstream write errors

Author: andrewhavck

.bleep

@@ -1 +1 @@
-fea06c09da4b7504dddd3273d88e14339455628c
+da81db0dcc6bb8aad2d01f0b6449accbeaef7013

pingora-proxy/src/proxy_h1.rs

@@ -154,6 +154,7 @@ impl<SV> HttpProxy<SV> {
     {
         let mut request_done = false;
         let mut response_done = false;
+        let mut send_error = None;
 
         /* duplex mode, wait for either to complete */
         while !request_done || !response_done {
@@ -178,18 +179,28 @@ impl<SV> HttpProxy<SV> {
                         Err(e) => {
                             // Push the error to downstream and then quit
                             // Don't care if send fails: downstream already gone
-                            let _ = tx.send(HttpTask::Failed(e.into_up())).await;
+                            let _ = tx.send(HttpTask::Failed(send_error.unwrap_or(e).into_up())).await;
                             // Downstream should consume all remaining data and handle the error
                             return Ok(())
                         }
                     }
                 },
 
                 body = rx.recv(), if !request_done => {
-                    request_done = send_body_to1(client_session, body).await?;
-                    // An upgraded request is terminated when either side is done
-                    if request_done && client_session.is_upgrade_req() {
-                        response_done = true;
+                    match send_body_to1(client_session, body).await {
+                        Ok(send_done) => {
+                            request_done = send_done;
+                            // An upgraded request is terminated when either side is done
+                            if request_done && client_session.is_upgrade_req() {
+                                response_done = true;
+                            }
+                        },
+                        Err(e) => {
+                           debug!("send error, draining read buf: {e}");
+                           request_done = true;
+                           send_error = Some(e);
+                           continue
+                        }
                     }
                 },
 

pingora-proxy/src/proxy_h2.rs

@@ -275,10 +275,18 @@ impl<SV> HttpProxy<SV> {
                         }
                     };
                     let is_body_done = session.is_body_done();
-                    let request_done =
-                        self.send_body_to2(session, body, is_body_done, client_body, ctx)
-                        .await?;
-                    downstream_state.maybe_finished(request_done);
+                    match self.send_body_to2(session, body, is_body_done, client_body, ctx).await {
+                        Ok(request_done) =>  {
+                            downstream_state.maybe_finished(request_done);
+                        },
+                        Err(e) => {
+                            // mark request done, attempt to drain receive
+                            warn!("Upstream h2 body send error: {e}");
+                            // upstream is what actually errored but we don't want to continue
+                            // polling the downstream body
+                            downstream_state.to_errored();
+                        }
+                    };
                 },
 
                 task = rx.recv(), if !response_state.upstream_done() => {

pingora-proxy/tests/test_upstream.rs

@@ -64,6 +64,36 @@ async fn test_connection_die() {
     assert!(body.is_err());
 }
 
+#[tokio::test]
+async fn test_upload_connection_die() {
+    init();
+    let client = reqwest::Client::new();
+    let res = client
+        .post("http://127.0.0.1:6147/upload_connection_die/")
+        .body("b".repeat(15 * 1024 * 1024)) // 15 MB upload
+        .timeout(Duration::from_secs(5))
+        .send()
+        .await
+        .unwrap();
+    // should get 200 status before connection dies
+    assert_eq!(res.status(), StatusCode::OK);
+    let _ = res.text().await;
+
+    // try h2
+    let client = reqwest::Client::new();
+    let res = client
+        .post("http://127.0.0.1:6147/upload_connection_die/")
+        .body("b".repeat(15 * 1024 * 1024)) // 15 MB upload
+        .timeout(Duration::from_secs(5))
+        .header("x-h2", "true")
+        .send()
+        .await
+        .unwrap();
+    // should get 200 status before connection dies
+    assert_eq!(res.status(), StatusCode::OK);
+    let _ = res.text().await;
+}
+
 #[tokio::test]
 async fn test_upload() {
     init();

pingora-proxy/tests/utils/conf/origin/conf/nginx.conf

@@ -299,6 +299,17 @@ http {
             }
         }
 
+        location /upload_connection_die/ {
+            content_by_lua_block {
+                ngx.status = ngx.HTTP_OK
+                ngx.print("")
+                ngx.flush(true)
+
+                time.sleep(1)
+                ngx.exit(444)
+            }
+        }
+
         location /download/ {
             content_by_lua_block {
                 ngx.req.read_body()