Progress towards allowing large writes

src/proto/streams/prioritize.rs

@@ -22,6 +22,17 @@ pub(super) struct Prioritize<B> {
     conn_task: Option<task::Task>,
 }
 
+#[derive(Debug)]
+pub(crate) struct Prioritized<B> {
+    // The buffer
+    inner: B,
+
+    // The stream that this is associated with
+    stream: store::Key,
+}
+
+// ===== impl Prioritize =====
+
 impl<B> Prioritize<B>
     where B: Buf,
 {
@@ -61,40 +72,56 @@ impl<B> Prioritize<B>
     pub fn queue_frame(&mut self,
                        frame: Frame<B>,
                        stream: &mut store::Ptr<B>)
+    {
+        if self.queue_frame2(frame, stream) {
+            // Notification required
+            if let Some(ref task) = self.conn_task {
+                task.notify();
+            }
+        }
+    }
+
+    /// Queue frame without actually notifying. Returns ture if the queue was
+    /// succesfful.
+    fn queue_frame2(&mut self, frame: Frame<B>, stream: &mut store::Ptr<B>)
+        -> bool
     {
         self.buffered_data += frame.flow_len();
 
         // queue the frame in the buffer
         stream.pending_send.push_back(&mut self.buffer, frame);
 
-        if stream.is_pending_send {
-            debug_assert!(!self.pending_send.is_empty());
-
-            // Already queued to have frame processed.
-            return;
-        }
-
         // Queue the stream
-        push_sender(&mut self.pending_send, stream);
+        !push_sender(&mut self.pending_send, stream)
+    }
 
-        if let Some(ref task) = self.conn_task {
-            task.notify();
-        }
+    /// Push the frame to the front of the stream's deque, scheduling the
+    /// steream if needed.
+    fn push_back_frame(&mut self, frame: Frame<B>, stream: &mut store::Ptr<B>) {
+        // Push the frame to the front of the stream's deque
+        stream.pending_send.push_front(&mut self.buffer, frame);
+
+        // If needed, schedule the sender
+        push_sender(&mut self.pending_capacity, stream);
     }
 
     pub fn poll_complete<T>(&mut self,
                             store: &mut Store<B>,
-                            dst: &mut Codec<T, B>)
+                            dst: &mut Codec<T, Prioritized<B>>)
         -> Poll<(), ConnectionError>
         where T: AsyncWrite,
     {
+        // Track the task
         self.conn_task = Some(task::current());
 
+        // Ensure codec is ready
+        try_ready!(dst.poll_ready());
+
+        // Reclaim any frame that has previously been written
+        self.reclaim_frame(store, dst);
+
         trace!("poll_complete");
         loop {
-            // Ensure codec is ready
-            try_ready!(dst.poll_ready());
-
             match self.pop_frame(store) {
                 Some(frame) => {
                     trace!("writing frame={:?}", frame);
@@ -106,15 +133,31 @@ impl<B> Prioritize<B>
                     // We already verified that `dst` is ready to accept the
                     // write
                     assert!(res.is_ready());
+
+                    // Ensure the codec is ready to try the loop again.
+                    try_ready!(dst.poll_ready());
+
+                    // Because, always try to reclaim...
+                    self.reclaim_frame(store, dst);
+
+                }
+                None => {
+                    // Try to flush the codec.
+                    try_ready!(dst.poll_complete());
+
+                    // This might release a data frame...
+                    if !self.reclaim_frame(store, dst) {
+                        return Ok(().into());
+                    }
+
+                    // No need to poll ready as poll_complete() does this for
+                    // us...
                 }
-                None => break,
             }
         }
-
-        Ok(().into())
     }
 
-    fn pop_frame(&mut self, store: &mut Store<B>) -> Option<Frame<B>> {
+    fn pop_frame(&mut self, store: &mut Store<B>) -> Option<Frame<Prioritized<B>>> {
         loop {
             match self.pop_sender(store) {
                 Some(mut stream) => {
@@ -124,11 +167,7 @@ impl<B> Prioritize<B>
 
                             if len > self.flow_control.effective_window_size() as usize {
                                 // TODO: This could be smarter...
-                                stream.pending_send.push_front(&mut self.buffer, frame.into());
-
-                                // Push the stream onto the list of streams
-                                // waiting for connection capacity
-                                push_sender(&mut self.pending_capacity, &mut stream);
+                                self.push_back_frame(frame.into(), &mut stream);
 
                                 // Try again w/ the next stream
                                 continue;
@@ -143,6 +182,14 @@ impl<B> Prioritize<B>
                         push_sender(&mut self.pending_send, &mut stream);
                     }
 
+                    // Add prioritization logic
+                    let frame = frame.map(|buf| {
+                        Prioritized {
+                            inner: buf,
+                            stream: stream.key(),
+                        }
+                    });
+
                     return Some(frame);
                 }
                 None => return None,
@@ -174,10 +221,58 @@ impl<B> Prioritize<B>
             }
         }
     }
+
+    fn reclaim_frame<T>(&mut self,
+                        store: &mut Store<B>,
+                        dst: &mut Codec<T, Prioritized<B>>) -> bool
+    {
+        // First check if there are any data chunks to take back
+        if let Some(frame) = dst.take_last_data_frame() {
+            let mut stream = store.resolve(frame.payload().stream);
+
+            let frame = frame.map(|prioritized| {
+                // TODO: Ensure fully written
+                prioritized.inner
+            });
+
+            self.push_back_frame(frame.into(), &mut stream);
+
+            true
+        } else {
+            false
+        }
+    }
 }
 
-fn push_sender<B>(list: &mut store::List<B>, stream: &mut store::Ptr<B>) {
-    debug_assert!(!stream.is_pending_send);
+/// Push the stream onto the `pending_send` list. Returns true if the sender was
+/// not already queued.
+fn push_sender<B>(list: &mut store::List<B>, stream: &mut store::Ptr<B>)
+    -> bool
+{
+    if stream.is_pending_send {
+        return false;
+    }
+
     list.push::<stream::Next>(stream);
     stream.is_pending_send = true;
+
+    true
+}
+
+// ===== impl Prioritized =====
+
+impl<B> Buf for Prioritized<B>
+    where B: Buf,
+{
+    fn remaining(&self) -> usize {
+        self.inner.remaining()
+    }
+
+    fn bytes(&self) -> &[u8] {
+        self.inner.bytes()
+    }
+
+    fn advance(&mut self, cnt: usize) {
+        self.inner.advance(cnt)
+    }
 }