Wire in trailers (#34)

Add send and receive trailer support.

src/client.rs

@@ -1,4 +1,4 @@
-use {frame, ConnectionError};
+use {frame, HeaderMap, ConnectionError};
 use Body;
 use frame::StreamId;
 use proto::{self, Connection, WindowSize};
@@ -173,10 +173,10 @@ impl<B: IntoBuf> Stream<B> {
     }
 
     /// Send trailers
-    pub fn send_trailers(&mut self, _trailers: ())
+    pub fn send_trailers(&mut self, trailers: HeaderMap)
         -> Result<(), ConnectionError>
     {
-        unimplemented!();
+        self.inner.send_trailers::<Peer>(trailers)
     }
 }
 

src/frame/headers.rs

@@ -98,6 +98,7 @@ const ALL: u8 = END_STREAM
 // ===== impl Headers =====
 
 impl Headers {
+    /// Create a new HEADERS frame
     pub fn new(stream_id: StreamId, pseudo: Pseudo, fields: HeaderMap) -> Self {
         Headers {
             stream_id: stream_id,
@@ -108,6 +109,19 @@ impl Headers {
         }
     }
 
+    pub fn trailers(stream_id: StreamId, fields: HeaderMap) -> Self {
+        let mut flags = HeadersFlag::default();
+        flags.set_end_stream();
+
+        Headers {
+            stream_id,
+            stream_dep: None,
+            fields: fields,
+            pseudo: Pseudo::default(),
+            flags: flags,
+        }
+    }
+
     /// Loads the header frame but doesn't actually do HPACK decoding.
     ///
     /// HPACK decoding is done in the `load_hpack` step.
@@ -290,6 +304,10 @@ impl Headers {
         Ok(request)
     }
 
+    pub fn into_fields(self) -> HeaderMap {
+        self.fields
+    }
+
     pub fn encode(self, encoder: &mut hpack::Encoder, dst: &mut BytesMut)
         -> Option<Continuation>
     {

src/lib.rs

@@ -61,6 +61,13 @@ impl<B: IntoBuf> Body<B> {
     pub fn release_capacity(&mut self, sz: usize) -> Result<(), ConnectionError> {
         self.inner.release_capacity(sz as proto::WindowSize)
     }
+
+    /// Poll trailers
+    ///
+    /// This function **must** not be called until `Body::poll` returns `None`.
+    pub fn poll_trailers(&mut self) -> Poll<Option<HeaderMap>, ConnectionError> {
+        self.inner.poll_trailers()
+    }
 }
 
 impl<B: IntoBuf> futures::Stream for Body<B> {

src/proto/streams/prioritize.rs

@@ -43,6 +43,8 @@ impl<B> Prioritize<B>
 
         flow.assign_capacity(config.init_local_window_sz);
 
+        trace!("Prioritize::new; flow={:?}", flow);
+
         Prioritize {
             pending_send: store::Queue::new(),
             pending_capacity: store::Queue::new(),

src/proto/streams/recv.rs

@@ -1,4 +1,4 @@
-use {client, server, frame, ConnectionError};
+use {client, server, frame, HeaderMap, ConnectionError};
 use proto::*;
 use super::*;
 
@@ -525,19 +525,15 @@ impl<B> Recv<B> where B: Buf {
         -> Poll<Option<Bytes>, ConnectionError>
     {
         match stream.pending_recv.pop_front(&mut self.buffer) {
+            Some(Frame::Data(frame)) => {
+                Ok(Some(frame.into_payload()).into())
+            }
             Some(frame) => {
-                match frame {
-                    Frame::Data(frame) => {
-                        Ok(Some(frame.into_payload()).into())
-                    }
-                    frame => {
-                        // Frame is trailer
-                        stream.pending_recv.push_front(&mut self.buffer, frame);
+                // Frame is trailer
+                stream.pending_recv.push_front(&mut self.buffer, frame);
 
-                        // No more data frames
-                        Ok(None.into())
-                    }
-                }
+                // No more data frames
+                Ok(None.into())
             }
             None => {
                 if stream.state.is_recv_closed() {
@@ -552,6 +548,32 @@ impl<B> Recv<B> where B: Buf {
         }
     }
 
+    pub fn poll_trailers(&mut self, stream: &mut Stream<B>)
+        -> Poll<Option<HeaderMap>, ConnectionError>
+    {
+        match stream.pending_recv.pop_front(&mut self.buffer) {
+            Some(Frame::Headers(frame)) => {
+                Ok(Some(frame.into_fields()).into())
+            }
+            Some(_) => {
+                // TODO: This is a user error. `poll_trailers` was called before
+                // the entire set of data frames have been consumed. What should
+                // we do?
+                unimplemented!();
+            }
+            None => {
+                if stream.state.is_recv_closed() {
+                    // There will be no trailer frame
+                    Ok(None.into())
+                } else {
+                    // Request to get notified once another frame arrives
+                    stream.recv_task = Some(task::current());
+                    Ok(Async::NotReady)
+                }
+            }
+        }
+    }
+
     fn reset(&mut self, _stream_id: StreamId, _reason: Reason) {
         unimplemented!();
     }

src/proto/streams/send.rs

@@ -133,6 +133,25 @@ impl<B> Send<B> where B: Buf {
         self.prioritize.send_data(frame, stream, task)
     }
 
+    pub fn send_trailers(&mut self,
+                         frame: frame::Headers,
+                         stream: &mut store::Ptr<B>,
+                         task: &mut Option<Task>)
+        -> Result<(), ConnectionError>
+    {
+        // TODO: Should this logic be moved into state.rs?
+        if !stream.state.is_send_streaming() {
+            return Err(UnexpectedFrameType.into());
+        }
+
+        stream.state.send_close()?;
+
+        trace!("send_trailers -- queuing; frame={:?}", frame);
+        self.prioritize.queue_frame(frame.into(), stream, task);
+
+        Ok(())
+    }
+
     pub fn poll_complete<T>(&mut self,
                             store: &mut Store<B>,
                             dst: &mut Codec<T, Prioritized<B>>)

src/proto/streams/streams.rs

@@ -1,4 +1,4 @@
-use {client, server};
+use {client, server, HeaderMap};
 use proto::*;
 use super::*;
 
@@ -335,6 +335,22 @@ impl<B> StreamRef<B>
         })
     }
 
+    pub fn send_trailers<P: Peer>(&mut self, trailers: HeaderMap) -> Result<(), ConnectionError>
+    {
+        let mut me = self.inner.lock().unwrap();
+        let me = &mut *me;
+
+        let stream = me.store.resolve(self.key);
+
+        // Create the trailers frame
+        let frame = frame::Headers::trailers(stream.id, trailers);
+
+        me.actions.transition::<P, _, _>(stream, |actions, stream| {
+            // Send the trailers frame
+            actions.send.send_trailers(frame, stream, &mut actions.task)
+        })
+    }
+
     /// Called by the server after the stream is accepted. Given that clients
     /// initialize streams by sending HEADERS, the request will always be
     /// available.
@@ -403,6 +419,15 @@ impl<B> StreamRef<B>
         me.actions.recv.poll_data(&mut stream)
     }
 
+    pub fn poll_trailers(&mut self) -> Poll<Option<HeaderMap>, ConnectionError> {
+        let mut me = self.inner.lock().unwrap();
+        let me = &mut *me;
+
+        let mut stream = me.store.resolve(self.key);
+
+        me.actions.recv.poll_trailers(&mut stream)
+    }
+
     /// Releases recv capacity back to the peer. This will result in sending
     /// WINDOW_UPDATE frames on both the stream and connection.
     pub fn release_capacity(&mut self, capacity: WindowSize)

src/server.rs

@@ -1,4 +1,4 @@
-use {Body, ConnectionError};
+use {Body, HeaderMap, ConnectionError};
 use frame::{self, StreamId};
 use proto::{self, Connection, WindowSize};
 use error::Reason;
@@ -185,10 +185,10 @@ impl<B: IntoBuf> Stream<B> {
     }
 
     /// Send trailers
-    pub fn send_trailers(&mut self, _trailers: ())
+    pub fn send_trailers(&mut self, trailers: HeaderMap)
         -> Result<(), ConnectionError>
     {
-        unimplemented!();
+        self.inner.send_trailers::<Peer>(trailers)
     }
 
     pub fn send_reset(mut self, reason: Reason) -> Result<(), ConnectionError> {