add SendResponse::poll_reset and SendStream::poll_reset to listen for reset streams (#279)

src/proto/mod.rs

@@ -10,7 +10,7 @@ pub(crate) use self::connection::{Config, Connection};
 pub(crate) use self::error::Error;
 pub(crate) use self::peer::{Peer, Dyn as DynPeer};
 pub(crate) use self::streams::{Key as StreamKey, StreamRef, OpaqueStreamRef, Streams};
-pub(crate) use self::streams::{Prioritized, Open};
+pub(crate) use self::streams::{PollReset, Prioritized, Open};
 
 use codec::Codec;
 

src/proto/streams/mod.rs

@@ -11,6 +11,7 @@ mod streams;
 
 pub(crate) use self::prioritize::Prioritized;
 pub(crate) use self::recv::Open;
+pub(crate) use self::send::PollReset;
 pub(crate) use self::store::Key;
 pub(crate) use self::streams::{StreamRef, OpaqueStreamRef, Streams};
 

src/proto/streams/send.rs

@@ -1,10 +1,15 @@
 use codec::{RecvError, UserError};
-use codec::UserError::*;
 use frame::{self, Reason};
-use http;
-use super::*;
+use super::{
+    store, Buffer, Codec, Config, Counts, Frame, Prioritize,
+    Prioritized, Store, Stream, StreamId, StreamIdOverflow, WindowSize,
+};
 
 use bytes::Buf;
+use http;
+use futures::{Async, Poll};
+use futures::task::Task;
+use tokio_io::AsyncWrite;
 
 use std::{cmp, io};
 
@@ -21,6 +26,13 @@ pub(super) struct Send {
     prioritize: Prioritize,
 }
 
+/// A value to detect which public API has called `poll_reset`.
+#[derive(Debug)]
+pub(crate) enum PollReset {
+    AwaitingHeaders,
+    Streaming,
+}
+
 impl Send {
     /// Create a new `Send`
     pub fn new(config: &Config) -> Self {
@@ -196,7 +208,7 @@ impl Send {
     ) -> Result<(), UserError> {
         // TODO: Should this logic be moved into state.rs?
         if !stream.state.is_send_streaming() {
-            return Err(UnexpectedFrameType.into());
+            return Err(UserError::UnexpectedFrameType);
         }
 
         stream.state.send_close();
@@ -263,6 +275,20 @@ impl Send {
         }
     }
 
+    pub fn poll_reset(
+        &self,
+        stream: &mut Stream,
+        mode: PollReset,
+    ) -> Poll<Reason, ::Error> {
+        match stream.state.ensure_reason(mode)? {
+            Some(reason) => Ok(reason.into()),
+            None => {
+                stream.wait_send();
+                Ok(Async::NotReady)
+            },
+        }
+    }
+
     pub fn recv_connection_window_update(
         &mut self,
         frame: frame::WindowUpdate,
@@ -405,6 +431,6 @@ impl Send {
     }
 
     pub fn ensure_next_stream_id(&self) -> Result<StreamId, UserError> {
-        self.next_stream_id.map_err(|_| OverflowedStreamId)
+        self.next_stream_id.map_err(|_| UserError::OverflowedStreamId)
     }
 }

src/proto/streams/state.rs

@@ -1,7 +1,9 @@
+use std::io;
+
 use codec::{RecvError, UserError};
 use codec::UserError::*;
 use frame::Reason;
-use proto;
+use proto::{self, PollReset};
 
 use self::Inner::*;
 use self::Peer::*;
@@ -399,8 +401,6 @@ impl State {
     }
 
     pub fn ensure_recv_open(&self) -> Result<bool, proto::Error> {
-        use std::io;
-
         // TODO: Is this correct?
         match self.inner {
             Closed(Cause::Proto(reason)) |
@@ -412,6 +412,24 @@ impl State {
             _ => Ok(true),
         }
     }
+
+    /// Returns a reason if the stream has been reset.
+    pub(super) fn ensure_reason(&self, mode: PollReset) -> Result<Option<Reason>, ::Error> {
+        match self.inner {
+            Closed(Cause::Proto(reason)) |
+            Closed(Cause::LocallyReset(reason)) |
+            Closed(Cause::Scheduled(reason)) => Ok(Some(reason)),
+            Closed(Cause::Io) => Err(proto::Error::Io(io::ErrorKind::BrokenPipe.into()).into()),
+            Open { local: Streaming, .. } |
+            HalfClosedRemote(Streaming) => match mode {
+                PollReset::AwaitingHeaders => {
+                    Err(UserError::PollResetAfterSendResponse.into())
+                },
+                PollReset::Streaming => Ok(None),
+            },
+            _ => Ok(None),
+        }
+    }
 }
 
 impl Default for State {

src/proto/streams/stream.rs

@@ -47,7 +47,7 @@ pub(super) struct Stream {
     pub buffered_send_data: WindowSize,
 
     /// Task tracking additional send capacity (i.e. window updates).
-    pub send_task: Option<task::Task>,
+    send_task: Option<task::Task>,
 
     /// Frames pending for this stream being sent to the socket
     pub pending_send: buffer::Deque,

src/proto/streams/streams.rs

@@ -708,7 +708,7 @@ where
             let mut stream = me.store.resolve(*key);
             trace!("poll_pending_open; stream = {:?}", stream.is_pending_open);
             if stream.is_pending_open {
-                stream.send_task = Some(task::current());
+                stream.wait_send();
                 return Ok(Async::NotReady);
             }
         }
@@ -930,6 +930,17 @@ impl<B> StreamRef<B> {
         me.actions.send.poll_capacity(&mut stream)
     }
 
+    /// Request to be notified for if a `RST_STREAM` is received for this stream.
+    pub(crate) fn poll_reset(&mut self, mode: proto::PollReset) -> Poll<Reason, ::Error> {
+        let mut me = self.opaque.inner.lock().unwrap();
+        let me = &mut *me;
+
+        let mut stream = me.store.resolve(self.opaque.key);
+
+        me.actions.send.poll_reset(&mut stream, mode)
+            .map_err(From::from)
+    }
+
     pub(crate) fn key(&self) -> store::Key {
         self.opaque.key
     }