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

@@ -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,26 +172,24 @@ where
             },
             service,
             date_header: config.date_header,
+            close_pending: false,
         }
     }
 
     pub(crate) fn graceful_shutdown(&mut self) {
         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() {
                     srv.conn.graceful_shutdown();
                 }
                 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));
                 }
             };