Add Graceful Shutdown support

If graceful shutdown is initiated, a GOAWAY of the max stream ID - 1 is
sent, followed by a PING frame, to measure RTT. When the PING is ACKed,
the connection sends a new GOAWAY with the proper last processed stream
ID. From there, once all active streams have completely, the connection
will finally close.

src/proto/ping_pong.rs

@@ -10,25 +10,67 @@ use tokio_io::AsyncWrite;
 /// Acknowledges ping requests from the remote.
 #[derive(Debug)]
 pub struct PingPong {
-    sending_pong: Option<PingPayload>,
+    pending_ping: Option<PendingPing>,
+    pending_pong: Option<PingPayload>,
+}
+
+#[derive(Debug)]
+struct PendingPing {
+    payload: PingPayload,
+    sent: bool,
+}
+
+/// Status returned from `PingPong::recv_ping`.
+#[derive(Debug)]
+pub(crate) enum ReceivedPing {
+    MustAck,
+    Unknown,
+    Shutdown,
 }
 
 impl PingPong {
     pub fn new() -> Self {
         PingPong {
-            sending_pong: None,
+            pending_ping: None,
+            pending_pong: None,
         }
     }
 
+    pub fn ping_shutdown(&mut self) {
+        assert!(self.pending_ping.is_none());
+
+        self.pending_ping = Some(PendingPing {
+            payload: Ping::SHUTDOWN,
+            sent: false,
+        });
+    }
+
     /// Process a ping
-    pub fn recv_ping(&mut self, ping: Ping) {
+    pub(crate) fn recv_ping(&mut self, ping: Ping) -> ReceivedPing {
         // The caller should always check that `send_pongs` returns ready before
         // calling `recv_ping`.
-        assert!(self.sending_pong.is_none());
+        assert!(self.pending_pong.is_none());
 
-        if !ping.is_ack() {
+        if ping.is_ack() {
+            if let Some(pending) = self.pending_ping.take() {
+                if &pending.payload == ping.payload() {
+                    trace!("recv PING ack");
+                    return ReceivedPing::Shutdown;
+                }
+
+                // if not the payload we expected, put it back.
+                self.pending_ping = Some(pending);
+            }
+
+            // else we were acked a ping we didn't send?
+            // The spec doesn't require us to do anything about this,
+            // so for resiliency, just ignore it for now.
+            warn!("recv PING ack that we never sent: {:?}", ping);
+            ReceivedPing::Unknown
+        } else {
             // Save the ping's payload to be sent as an acknowledgement.
-            self.sending_pong = Some(ping.into_payload());
+            self.pending_pong = Some(ping.into_payload());
+            ReceivedPing::MustAck
         }
     }
 
@@ -38,15 +80,46 @@ impl PingPong {
         T: AsyncWrite,
         B: Buf,
     {
-        if let Some(pong) = self.sending_pong.take() {
+        if let Some(pong) = self.pending_pong.take() {
             if !dst.poll_ready()?.is_ready() {
-                self.sending_pong = Some(pong);
+                self.pending_pong = Some(pong);
                 return Ok(Async::NotReady);
             }
 
-            dst.buffer(Ping::pong(pong).into()).ok().expect("invalid pong frame");
+            dst.buffer(Ping::pong(pong).into())
+                .expect("invalid pong frame");
+        }
+
+        Ok(Async::Ready(()))
+    }
+
+    /// Send any pending pings.
+    pub fn send_pending_ping<T, B>(&mut self, dst: &mut Codec<T, B>) -> Poll<(), io::Error>
+    where
+        T: AsyncWrite,
+        B: Buf,
+    {
+        if let Some(ref mut ping) = self.pending_ping {
+            if !ping.sent {
+                if !dst.poll_ready()?.is_ready() {
+                    return Ok(Async::NotReady);
+                }
+
+                dst.buffer(Ping::new(ping.payload).into())
+                    .expect("invalid ping frame");
+                ping.sent = true;
+            }
         }
 
         Ok(Async::Ready(()))
     }
 }
+
+impl ReceivedPing {
+    pub fn is_shutdown(&self) -> bool {
+        match *self {
+            ReceivedPing::Shutdown => true,
+            _ => false,
+        }
+    }
+}