Now graceful_shutdown close the connection after handshake is finished

Author: ionut-slaveanu

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));
                 }
             };