From 13b0594348916b901ad7e1c838b9d90298db6af4 Mon Sep 17 00:00:00 2001
From: Slaveanu Ionut <ionut.slaveanu@adswizz.com>
Date: Fri, 25 Oct 2024 18:38:17 +0300
Subject: [PATCH] fix(http2): improve graceful shutdown during handshake
 (#3729)

Before, if a graceful shutdown was triggered while the handshake was in-progress, the connection would just be dropped instantly. However, if some requests were in-flight as part of the handshake, they'd get dropped along with it.

Now, if handshake is in-progress, we record that a close is desired, and as soon as the handshake has completed, the real graceful shutdown process starts, which should send the proper signals back about what happened with the in-flight requests.
---
 src/proto/h2/server.rs | 19 ++++++++-----------
 1 file changed, 8 insertions(+), 11 deletions(-)

diff --git a/src/proto/h2/server.rs b/src/proto/h2/server.rs
index be897a54cf..a434ba9e6c 100644
--- a/src/proto/h2/server.rs
+++ b/src/proto/h2/server.rs
@@ -89,6 +89,7 @@ pin_project! {
         service: S,
         state: State<T, B>,
         date_header: bool,
+        close_pending: bool
     }
 }
 
@@ -101,7 +102,6 @@ where
         hs: Handshake<Compat<T>, SendBuf<B::Data>>,
     },
     Serving(Serving<T, B>),
-    Closed,
 }
 
 struct Serving<T, B>
@@ -172,6 +172,7 @@ where
             },
             service,
             date_header: config.date_header,
+            close_pending: false,
         }
     }
 
@@ -179,7 +180,8 @@ where
         trace!("graceful_shutdown");
         match self.state {
             State::Handshaking { .. } => {
-                // fall-through, to replace state with Closed
+                self.close_pending = true;
+                return;
             }
             State::Serving(ref mut srv) => {
                 if srv.closing.is_none() {
@@ -187,11 +189,7 @@ where
                 }
                 return;
             }
-            State::Closed => {
-                return;
-            }
         }
-        self.state = State::Closed;
     }
 }
 
@@ -228,12 +226,11 @@ where
                     })
                 }
                 State::Serving(ref mut srv) => {
-                    ready!(srv.poll_server(cx, &mut me.service, &mut me.exec))?;
-                    return Poll::Ready(Ok(Dispatched::Shutdown));
-                }
-                State::Closed => {
                     // graceful_shutdown was called before handshaking finished,
-                    // nothing to do here...
+                    if true == me.close_pending && srv.closing.is_none() {
+                        srv.conn.graceful_shutdown();
+                    }
+                    ready!(srv.poll_server(cx, &mut me.service, &mut me.exec))?;
                     return Poll::Ready(Ok(Dispatched::Shutdown));
                 }
             };