Implement stream prioritization (fix #165)

Author: geieredgar

quinn-proto/src/connection/mod.rs

@@ -1653,6 +1653,37 @@ where
         Ok(())
     }
 
+    /// Set the priority of a stream
+    ///
+    /// # Panics
+    /// - when applied to a receive stream
+    pub fn set_priority(
+        &mut self,
+        stream_id: StreamId,
+        priority: i32,
+    ) -> Result<(), UnknownStream> {
+        assert!(
+            stream_id.dir() == Dir::Bi || stream_id.initiator() == self.side,
+            "only streams supporting outgoing data may change priority"
+        );
+
+        self.streams.set_priority(stream_id, priority)?;
+        Ok(())
+    }
+
+    /// Get the priority of a stream
+    ///
+    /// # Panics
+    /// - when applied to a receive stream
+    pub fn priority(&mut self, stream_id: StreamId) -> Result<i32, UnknownStream> {
+        assert!(
+            stream_id.dir() == Dir::Bi || stream_id.initiator() == self.side,
+            "only streams supporting outgoing data have a priority"
+        );
+
+        Ok(self.streams.priority(stream_id)?)
+    }
+
     /// Handle the already-decrypted first packet from the client
     ///
     /// Decrypting the first packet in the `Endpoint` allows stateless packet handling to be more

quinn-proto/src/connection/streams.rs

@@ -1,5 +1,6 @@
 use std::{
-    collections::{hash_map, HashMap, VecDeque},
+    cell::RefCell,
+    collections::{binary_heap::PeekMut, hash_map, BinaryHeap, HashMap, VecDeque},
     convert::TryFrom,
     mem,
 };
@@ -50,7 +51,7 @@ pub struct Streams {
     /// permitted to open but which have not yet been opened.
     send_streams: usize,
     /// Streams with outgoing data queued
-    pending: VecDeque<StreamId>,
+    pending: BinaryHeap<PendingLevel>,
 
     events: VecDeque<StreamEvent>,
     /// Streams blocked on connection-level flow control or stream window space
@@ -105,7 +106,7 @@ impl Streams {
             opened: [false, false],
             next_reported_remote: [0, 0],
             send_streams: 0,
-            pending: VecDeque::new(),
+            pending: BinaryHeap::new(),
             events: VecDeque::new(),
             connection_blocked: Vec::new(),
             max_data: 0,
@@ -272,7 +273,7 @@ impl Streams {
         self.unacked_data += len as u64;
         trace!(stream = %id, "wrote {} bytes", len);
         if !was_pending {
-            self.pending.push_back(id);
+            push_pending(&mut self.pending, id, stream.priority);
         }
         Ok(len)
     }
@@ -393,7 +394,7 @@ impl Streams {
         let was_pending = stream.is_pending();
         stream.finish()?;
         if !was_pending {
-            self.pending.push_back(id);
+            push_pending(&mut self.pending, id, stream.priority);
         }
         Ok(())
     }
@@ -423,6 +424,25 @@ impl Streams {
         Ok(())
     }
 
+    pub fn set_priority(&mut self, id: StreamId, priority: i32) -> Result<(), UnknownStream> {
+        let stream = match self.send.get_mut(&id) {
+            Some(ss) => ss,
+            None => return Err(UnknownStream { _private: () }),
+        };
+
+        stream.priority = priority;
+        Ok(())
+    }
+
+    pub fn priority(&mut self, id: StreamId) -> Result<i32, UnknownStream> {
+        let stream = match self.send.get_mut(&id) {
+            Some(ss) => ss,
+            None => return Err(UnknownStream { _private: () }),
+        };
+
+        Ok(stream.priority)
+    }
+
     pub fn reset_acked(&mut self, id: StreamId) {
         match self.send.entry(id) {
             hash_map::Entry::Vacant(_) => {}
@@ -594,14 +614,21 @@ impl Streams {
                 Some(x) => x,
                 None => break,
             };
+            let mut level = match self.pending.peek_mut() {
+                Some(x) => x,
+                None => break,
+            };
             // Poppping data from the front of the queue, storing as much data
             // as possible in a single frame, and enqueing sending further
             // remaining data at the end of the queue helps with fairness.
             // Other streams will have a chance to write data before we touch
             // this stream again.
-            let id = match self.pending.pop_front() {
+            let id = match level.queue.get_mut().pop_front() {
                 Some(x) => x,
-                None => break,
+                None => {
+                    PeekMut::pop(level);
+                    continue;
+                }
             };
             let stream = match self.send.get_mut(&id) {
                 Some(s) => s,
@@ -622,7 +649,12 @@ impl Streams {
                 stream.fin_pending = false;
             }
             if stream.is_pending() {
-                self.pending.push_back(id);
+                if level.priority == stream.priority {
+                    level.queue.get_mut().push_back(id);
+                } else {
+                    drop(level);
+                    push_pending(&mut self.pending, id, stream.priority);
+                }
             }
 
             let meta = frame::StreamMeta { id, offsets, fin };
@@ -691,7 +723,7 @@ impl Streams {
             Some(x) => x,
         };
         if !stream.is_pending() {
-            self.pending.push_back(frame.id);
+            push_pending(&mut self.pending, frame.id, stream.priority);
         }
         stream.fin_pending |= frame.fin;
         stream.pending.retransmit(frame.offsets);
@@ -708,7 +740,7 @@ impl Streams {
                     continue;
                 }
                 if !stream.is_pending() {
-                    self.pending.push_back(id);
+                    push_pending(&mut self.pending, id, stream.priority);
                 }
                 stream.pending.retransmit_all_for_0rtt();
             }
@@ -916,6 +948,47 @@ impl Streams {
     }
 }
 
+fn push_pending(pending: &mut BinaryHeap<PendingLevel>, id: StreamId, priority: i32) {
+    for level in pending.iter() {
+        if priority == level.priority {
+            level.queue.borrow_mut().push_back(id);
+            return;
+        }
+    }
+    let mut queue = VecDeque::new();
+    queue.push_back(id);
+    pending.push(PendingLevel {
+        queue: RefCell::new(queue),
+        priority,
+    });
+}
+
+struct PendingLevel {
+    // RefCell is needed because BinaryHeap doesn't have an iter_mut()
+    queue: RefCell<VecDeque<StreamId>>,
+    priority: i32,
+}
+
+impl PartialEq for PendingLevel {
+    fn eq(&self, other: &Self) -> bool {
+        self.priority.eq(&other.priority)
+    }
+}
+
+impl PartialOrd for PendingLevel {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Eq for PendingLevel {}
+
+impl Ord for PendingLevel {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.priority.cmp(&other.priority)
+    }
+}
+
 /// Application events about streams
 #[derive(Debug)]
 pub enum StreamEvent {
@@ -1209,4 +1282,28 @@ mod tests {
             TransportErrorCode::FLOW_CONTROL_ERROR
         );
     }
+
+    #[test]
+    fn stream_priority() {
+        let mut server = make(Side::Server);
+        server.set_params(&TransportParameters {
+            initial_max_streams_bidi: 3u32.into(),
+            initial_max_data: 10u32.into(),
+            initial_max_stream_data_bidi_remote: 10u32.into(),
+            ..Default::default()
+        });
+        let id_high = server.open(Dir::Bi).unwrap();
+        let id_mid = server.open(Dir::Bi).unwrap();
+        let id_low = server.open(Dir::Bi).unwrap();
+        server.set_priority(id_low, -1).unwrap();
+        server.set_priority(id_high, 1).unwrap();
+        server.write(id_mid, b"mid").unwrap();
+        server.write(id_low, b"low").unwrap();
+        server.write(id_high, b"high").unwrap();
+        let mut buf = Vec::with_capacity(40);
+        let meta = server.write_stream_frames(&mut buf, 40);
+        assert_eq!(meta[0].id, id_high);
+        assert_eq!(meta[1].id, id_mid);
+        assert_eq!(meta[2].id, id_low);
+    }
 }

quinn-proto/src/connection/streams/send.rs

@@ -9,6 +9,7 @@ pub(super) struct Send {
     pub(super) max_data: u64,
     pub(super) state: SendState,
     pub(super) pending: SendBuffer,
+    pub(super) priority: i32,
     /// Whether a frame containing a FIN bit must be transmitted, even if we don't have any new data
     pub(super) fin_pending: bool,
     /// Whether this stream is in the `connection_blocked` list of `Streams`
@@ -23,6 +24,7 @@ impl Send {
             max_data: max_data.into(),
             state: SendState::Ready,
             pending: SendBuffer::new(),
+            priority: 0,
             fin_pending: false,
             connection_blocked: false,
             stop_reason: None,

quinn/src/streams.rs

@@ -150,6 +150,25 @@ where
         Ok(())
     }
 
+    /// Set the priority of the send stream
+    ///
+    /// Every send stream has an initial priority of 0. Locally buffered data from streams with
+    /// higher priority will be transmitted before data from streams with lower priority. Changing
+    /// the priority of a stream with pending data may only take effect after that data has been
+    /// transmitted. Using many different priority levels per connection may have a negative
+    /// impact on performance.
+    pub fn set_priority(&self, priority: i32) -> Result<(), UnknownStream> {
+        let mut conn = self.conn.lock().unwrap();
+        conn.inner.set_priority(self.stream, priority)?;
+        Ok(())
+    }
+
+    /// Get the priority of the send stream
+    pub fn priority(&self) -> Result<i32, UnknownStream> {
+        let mut conn = self.conn.lock().unwrap();
+        Ok(conn.inner.priority(self.stream)?)
+    }
+
     /// Completes if/when the peer stops the stream, yielding the error code
     pub fn stopped(&mut self) -> Stopped<'_, S> {
         Stopped { stream: self }