When Streams are dropped, close Connection (#221) (#222)

When all Streams are dropped / finished, the Connection was held
open until the peer hangs up. Instead, the Connection should hang up
once it knows that nothing more will be sent.

To fix this, we notify the Connection when a stream is no longer
referenced. On the Connection poll(), we check that there are no
active, held, reset streams or any references to the Streams
and transition to sending a GOAWAY if that is case.

The specific behavior depends on if running as a client or server.

src/proto/connection.rs

@@ -122,6 +122,27 @@ where
         Ok(().into())
     }
 
+    fn transition_to_go_away(&mut self, id: StreamId, e: Reason) {
+        let goaway = frame::GoAway::new(id, e);
+        self.state = State::GoAway(goaway);
+    }
+
+    /// Closes the connection by transitioning to a GOAWAY state
+    /// iff there are no streams or references
+    pub fn maybe_close_connection_if_no_streams(&mut self) {
+        // If we poll() and realize that there are no streams or references
+        // then we can close the connection by transitioning to GOAWAY
+        if self.streams.num_active_streams() == 0 && !self.streams.has_streams_or_other_references() {
+            self.close_connection();
+        }
+    }
+
+    /// Closes the connection by transitioning to a GOAWAY state
+    pub fn close_connection(&mut self) {
+        let last_processed_id = self.streams.last_processed_id();
+        self.transition_to_go_away(last_processed_id, Reason::NO_ERROR);
+    }
+
     /// Advances the internal state of the connection.
     pub fn poll(&mut self) -> Poll<(), proto::Error> {
         use codec::RecvError::*;
@@ -143,9 +164,8 @@ where
 
                             if self.error.is_some() {
                                 if self.streams.num_active_streams() == 0 {
-                                    let id = self.streams.last_processed_id();
-                                    let goaway = frame::GoAway::new(id, Reason::NO_ERROR);
-                                    self.state = State::GoAway(goaway);
+                                    let last_processed_id = self.streams.last_processed_id();
+                                    self.transition_to_go_away(last_processed_id, Reason::NO_ERROR);
                                     continue;
                                 }
                             }
@@ -160,12 +180,7 @@ where
 
                             // Reset all active streams
                             let last_processed_id = self.streams.recv_err(&e.into());
-
-                            // Create the GO_AWAY frame with the last_processed_id
-                            let frame = frame::GoAway::new(last_processed_id, e);
-
-                            // Transition to the going away state.
-                            self.state = State::GoAway(frame);
+                            self.transition_to_go_away(last_processed_id, e);
                         },
                         // Attempting to read a frame resulted in a stream level error.
                         // This is handled by resetting the frame then trying to read

src/proto/peer.rs

@@ -63,6 +63,7 @@ impl Dyn {
     /// Returns true if the remote peer can initiate a stream with the given ID.
     pub fn ensure_can_open(&self, id: StreamId) -> Result<(), RecvError> {
         if !self.is_server() {
+            trace!("Cannot open stream {:?} - not server, PROTOCOL_ERROR", id);
             // Remote is a server and cannot open streams. PushPromise is
             // registered by reserving, so does not go through this path.
             return Err(RecvError::Connection(Reason::PROTOCOL_ERROR));
@@ -70,6 +71,7 @@ impl Dyn {
 
         // Ensure that the ID is a valid server initiated ID
         if !id.is_client_initiated() {
+            trace!("Cannot open stream {:?} - not client initiated, PROTOCOL_ERROR", id);
             return Err(RecvError::Connection(Reason::PROTOCOL_ERROR));
         }
 

src/proto/streams/counts.rs

@@ -46,6 +46,10 @@ impl Counts {
         self.peer
     }
 
+    pub fn has_streams(&self) -> bool {
+        self.num_send_streams != 0 || self.num_recv_streams != 0
+    }
+
     /// Returns true if the receive stream concurrency can be incremented
     pub fn can_inc_num_recv_streams(&self) -> bool {
         self.max_recv_streams > self.num_recv_streams

src/proto/streams/recv.rs

@@ -119,6 +119,7 @@ impl Recv {
 
         let next_id = self.next_stream_id()?;
         if id < next_id {
+            trace!("id ({:?}) < next_id ({:?}), PROTOCOL_ERROR", id, next_id);
             return Err(RecvError::Connection(Reason::PROTOCOL_ERROR));
         }
 

src/proto/streams/streams.rs

@@ -1,10 +1,10 @@
-use super::{Buffer, Config, Counts, Prioritized, Recv, Send, Stream, StreamId};
-use super::recv::RecvHeaderBlockError;
-use super::store::{self, Entry, Resolve, Store};
 use {client, proto, server};
 use codec::{Codec, RecvError, SendError, UserError};
 use frame::{self, Frame, Reason};
 use proto::{peer, Peer, WindowSize};
+use super::{Buffer, Config, Counts, Prioritized, Recv, Send, Stream, StreamId};
+use super::recv::RecvHeaderBlockError;
+use super::store::{self, Entry, Resolve, Store};
 
 use bytes::{Buf, Bytes};
 use futures::{task, Async, Poll};
@@ -520,8 +520,8 @@ where
         end_of_stream: bool,
         pending: Option<&store::Key>,
     ) -> Result<StreamRef<B>, SendError> {
-        use super::stream::ContentLength;
         use http::Method;
+        use super::stream::ContentLength;
 
         // TODO: There is a hazard with assigning a stream ID before the
         // prioritize layer. If prioritization reorders new streams, this
@@ -661,6 +661,19 @@ where
         me.store.num_active_streams()
     }
 
+    pub fn has_streams_or_other_references(&self) -> bool {
+        if Arc::strong_count(&self.inner) > 1 {
+            return true;
+        }
+
+        if Arc::strong_count(&self.send_buffer) > 1 {
+            return true;
+        }
+
+        let me = self.inner.lock().unwrap();
+        me.counts.has_streams()
+    }
+
     #[cfg(feature = "unstable")]
     pub fn num_wired_streams(&self) -> usize {
         let me = self.inner.lock().unwrap();
@@ -951,6 +964,16 @@ fn drop_stream_ref(inner: &Mutex<Inner>, key: store::Key) {
 
     let actions = &mut me.actions;
 
+    // If the stream is not referenced and it is already
+    // closed (does not have to go through logic below
+    // of canceling the stream), we should notify the task
+    // (connection) so that it can close properly
+    if stream.ref_count == 0 && stream.is_closed() {
+        if let Some(task) = actions.task.take() {
+            task.notify();
+        }
+    }
+
     me.counts.transition(stream, |counts, stream| {
         maybe_cancel(stream, actions, counts);