From 03a5b99f291aa7e96bd59f6083280d1aebff79b4 Mon Sep 17 00:00:00 2001
From: Jordan Wu <101218661+jordan-definitive@users.noreply.github.com>
Date: Tue, 17 Sep 2024 19:08:07 -0700
Subject: [PATCH] shutdown h2 connection gracefully with GOAWAYs

---
 pingora-core/src/apps/mod.rs | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/pingora-core/src/apps/mod.rs b/pingora-core/src/apps/mod.rs
index 32fd82f4..a3470e48 100644
--- a/pingora-core/src/apps/mod.rs
+++ b/pingora-core/src/apps/mod.rs
@@ -126,12 +126,20 @@ where
                 }
                 Ok(c) => c,
             };
+            let mut shutdown = shutdown.clone();
 
             loop {
                 // this loop ends when the client decides to close the h2 conn
                 // TODO: add a timeout?
-                let h2_stream =
-                    server::HttpSession::from_h2_conn(&mut h2_conn, digest.clone()).await;
+                let h2_stream = tokio::select! {
+                    _ = shutdown.changed() => {
+                        h2_conn.graceful_shutdown();
+                        // continue to the next loop to wait for client disconnect before returning.
+                        // the client will disconnect when all in flight streams are closed.
+                        continue;
+                    }
+                    h2_stream = server::HttpSession::from_h2_conn(&mut h2_conn, digest.clone()) => h2_stream
+                };
                 let h2_stream = match h2_stream {
                     Err(e) => {
                         // It is common for the client to just disconnect TCP without properly