WIP: send flow control

src/proto/streams/prioritize.rs

@@ -2,8 +2,18 @@ use super::*;
 
 #[derive(Debug)]
 pub(super) struct Prioritize<B> {
+    /// Streams that have pending frames
     pending_send: store::List<B>,
 
+    /// Streams that are waiting for connection level flow control capacity
+    pending_capacity: store::List<B>,
+
+    /// Connection level flow control governing sent data
+    flow_control: FlowControl,
+
+    /// Total amount of buffered data in data frames
+    buffered_data: usize,
+
     /// Holds frames that are waiting to be written to the socket
     buffer: Buffer<B>,
 }
@@ -11,17 +21,44 @@ pub(super) struct Prioritize<B> {
 impl<B> Prioritize<B>
     where B: Buf,
 {
-    pub fn new() -> Prioritize<B> {
+    pub fn new(config: &Config) -> Prioritize<B> {
         Prioritize {
             pending_send: store::List::new(),
+            pending_capacity: store::List::new(),
+            flow_control: FlowControl::new(config.init_local_window_sz),
+            buffered_data: 0,
             buffer: Buffer::new(),
         }
     }
 
+    pub fn available_window(&self) -> WindowSize {
+        let win = self.flow_control.window_size();
+
+        if self.buffered_data >= win as usize {
+            0
+        } else {
+            win - self.buffered_data as WindowSize
+        }
+    }
+
+    pub fn recv_window_update(&mut self, frame: frame::WindowUpdate)
+        -> Result<(), ConnectionError>
+    {
+        // Expand the window
+        self.flow_control.expand_window(frame.size_increment())?;
+
+        // Imediately apply the update
+        self.flow.apply_window_update();
+
+        Ok(())
+    }
+
     pub fn queue_frame(&mut self,
                        frame: Frame<B>,
                        stream: &mut store::Ptr<B>)
     {
+        self.buffered_data += frame.flow_len();
+
         // queue the frame in the buffer
         stream.pending_send.push_back(&mut self.buffer, frame);
 
@@ -33,7 +70,7 @@ impl<B> Prioritize<B>
         }
 
         // Queue the stream
-        self.push_sender(stream);
+        push_sender(&mut self.pending_send, stream);
     }
 
     pub fn poll_complete<T>(&mut self,
@@ -48,7 +85,9 @@ impl<B> Prioritize<B>
 
             match self.pop_frame(store) {
                 Some(frame) => {
-                    // TODO: data frames should be handled specially...
+                    // Subtract the data size
+                    self.buffered_data -= frame.flow_len();
+
                     let res = dst.start_send(frame)?;
 
                     // We already verified that `dst` is ready to accept the
@@ -63,35 +102,63 @@ impl<B> Prioritize<B>
     }
 
     fn pop_frame(&mut self, store: &mut Store<B>) -> Option<Frame<B>> {
-        match self.pop_sender(store) {
-            Some(mut stream) => {
-                let frame = stream.pending_send.pop_front(&mut self.buffer).unwrap();
+        loop {
+            match self.pop_sender(store) {
+                Some(mut stream) => {
+                    let frame = match stream.pending_send.pop_front(&mut self.buffer).unwrap() {
+                        Frame::Data(frame) => {
+                            let len = frame.payload().remaining();
 
-                if !stream.pending_send.is_empty() {
-                    self.push_sender(&mut stream);
+                            if len > self.flow_control.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);
+
+                                // Try again w/ the next stream
+                                continue;
+                            }
+
+                            frame.into()
+                        }
+                        frame => frame,
+                    };
+
+                    if !stream.pending_send.is_empty() {
+                        push_sender(&mut self.pending_send, &mut stream);
+                    }
+
+                    return Some(frame);
                 }
-
-                Some(frame)
+                None => return None,
             }
-            None => None,
         }
     }
 
-    fn push_sender(&mut self, stream: &mut store::Ptr<B>) {
-        debug_assert!(!stream.is_pending_send);
-
-        self.pending_send.push(stream);
-
-        stream.is_pending_send = true;
-    }
-
     fn pop_sender<'a>(&mut self, store: &'a mut Store<B>) -> Option<store::Ptr<'a, B>> {
-        match self.pending_send.pop(store) {
-            Some(mut stream) => {
-                stream.is_pending_send = false;
-                Some(stream)
+        // If the connection level window has capacity, pop off of the pending
+        // capacity list first.
+
+        if self.flow_control.has_capacity() && !self.pending_capacity.is_empty() {
+            let mut stream = self.pending_capacity.pop(store).unwrap();
+            stream.is_pending_send = false;
+            Some(stream)
+        } else {
+            match self.pending_send.pop(store) {
+                Some(mut stream) => {
+                    stream.is_pending_send = false;
+                    Some(stream)
+                }
+                None => None,
             }
-            None => None,
         }
     }
 }
+
+fn push_sender<B>(list: &mut store::List<B>, stream: &mut store::Ptr<B>) {
+    debug_assert!(!stream.is_pending_send);
+    list.push(stream);
+    stream.is_pending_send = true;
+}