StreamState: allow reusing Recv instances

alessandrod · djc · commit 41850c8a304f · 2024-10-14T08:07:06.000Z
Recv is used to receive data during the lifetime of a stream.  When a
stream is closed, instead of throwing away the corresponding Recv, put
it in a free list so they can be reused.

This reduces allocations by avoiding repeated allocations of BinaryHeaps
within Assemblers, and allocations of Box(es) for the actual Recvs. For
servers that serve millions of streams, this drastically reduces
allocations, effectively allocating a max_(uni|bi)_streams window worth
of Recvs and then reusing them instead of allocating for each new
stream.
diff --git a/quinn-proto/src/connection/assembler.rs b/quinn-proto/src/connection/assembler.rs
@@ -29,6 +29,14 @@ impl Assembler {
         Self::default()
     }
 
+    /// Reset to the initial state
+    pub(super) fn reinit(&mut self) {
+        let old_data = mem::take(&mut self.data);
+        *self = Self::default();
+        self.data = old_data;
+        self.data.clear();
+    }
+
     pub(super) fn ensure_ordering(&mut self, ordered: bool) -> Result<(), IllegalOrderedRead> {
         if ordered && !self.state.is_ordered() {
             return Err(IllegalOrderedRead);
diff --git a/quinn-proto/src/connection/streams/mod.rs b/quinn-proto/src/connection/streams/mod.rs
@@ -7,10 +7,11 @@ use bytes::Bytes;
 use thiserror::Error;
 use tracing::trace;
 
-use self::state::get_or_insert_recv;
-
 use super::spaces::{Retransmits, ThinRetransmits};
-use crate::{connection::streams::state::get_or_insert_send, frame, Dir, StreamId, VarInt};
+use crate::{
+    connection::streams::state::{get_or_insert_recv, get_or_insert_send},
+    frame, Dir, StreamId, VarInt,
+};
 
 mod recv;
 use recv::Recv;
@@ -149,8 +150,8 @@ impl<'a> RecvStream<'a> {
         // connection-level flow control to account for discarded data. Otherwise, we can discard
         // state immediately.
         if !stream.final_offset_unknown() {
-            entry.remove();
-            self.state.stream_freed(self.id, StreamHalf::Recv);
+            let recv = entry.remove().expect("must have recv when stopping");
+            self.state.stream_recv_freed(self.id, recv);
         }
 
         if self.state.add_read_credits(read_credits).should_transmit() {
@@ -168,7 +169,7 @@ impl<'a> RecvStream<'a> {
         let hash_map::Entry::Occupied(entry) = self.state.recv.entry(self.id) else {
             return Err(ClosedStream { _private: () });
         };
-        let Some(s) = entry.get().as_ref() else {
+        let Some(s) = entry.get().as_ref().and_then(|s| s.as_open_recv()) else {
             return Ok(None);
         };
         if s.stopped {
@@ -180,8 +181,9 @@ impl<'a> RecvStream<'a> {
 
         // Clean up state after application observes the reset, since there's no reason for the
         // application to attempt to read or stop the stream once it knows it's reset
-        entry.remove_entry();
-        self.state.stream_freed(self.id, StreamHalf::Recv);
+        let (_, recv) = entry.remove_entry();
+        self.state
+            .stream_recv_freed(self.id, recv.expect("must have recv on reset"));
         self.state.queue_max_stream_id(self.pending);
 
         Ok(Some(code))
diff --git a/quinn-proto/src/connection/streams/recv.rs b/quinn-proto/src/connection/streams/recv.rs
@@ -5,12 +5,14 @@ use thiserror::Error;
 use tracing::debug;
 
 use super::state::get_or_insert_recv;
-use super::{ClosedStream, Retransmits, ShouldTransmit, StreamHalf, StreamId, StreamsState};
+use super::{ClosedStream, Retransmits, ShouldTransmit, StreamId, StreamsState};
 use crate::connection::assembler::{Assembler, Chunk, IllegalOrderedRead};
+use crate::connection::streams::state::StreamRecv;
 use crate::{frame, TransportError, VarInt};
 
 #[derive(Debug, Default)]
 pub(super) struct Recv {
+    // NB: when adding or removing fields, remember to update `reinit`.
     state: RecvState,
     pub(super) assembler: Assembler,
     sent_max_stream_data: u64,
@@ -29,6 +31,15 @@ impl Recv {
         })
     }
 
+    /// Reset to the initial state
+    pub(super) fn reinit(&mut self, initial_max_data: u64) {
+        self.state = RecvState::default();
+        self.assembler.reinit();
+        self.sent_max_stream_data = initial_max_data;
+        self.end = 0;
+        self.stopped = false;
+    }
+
     /// Process a STREAM frame
     ///
     /// Return value is `(number_of_new_bytes_ingested, stream_is_closed)`
@@ -255,7 +266,7 @@ impl<'a> Chunks<'a> {
         let mut recv =
             match get_or_insert_recv(streams.stream_receive_window)(entry.get_mut()).stopped {
                 true => return Err(ReadableError::ClosedStream),
-                false => entry.remove().unwrap(), // this can't fail due to the previous get_or_insert_with
+                false => entry.remove().unwrap().into_inner(), // this can't fail due to the previous get_or_insert_with
             };
 
         recv.assembler.ensure_ordering(ordered)?;
@@ -292,14 +303,24 @@ impl<'a> Chunks<'a> {
         match rs.state {
             RecvState::ResetRecvd { error_code, .. } => {
                 debug_assert_eq!(self.read, 0, "reset streams have empty buffers");
-                self.streams.stream_freed(self.id, StreamHalf::Recv);
-                self.state = ChunksState::Reset(error_code);
+                let state = mem::replace(&mut self.state, ChunksState::Reset(error_code));
+                // At this point if we have `rs` self.state must be `ChunksState::Readable`
+                let recv = match state {
+                    ChunksState::Readable(recv) => StreamRecv::Open(recv),
+                    _ => unreachable!("state must be ChunkState::Readable"),
+                };
+                self.streams.stream_recv_freed(self.id, recv);
                 Err(ReadError::Reset(error_code))
             }
             RecvState::Recv { size } => {
                 if size == Some(rs.end) && rs.assembler.bytes_read() == rs.end {
-                    self.streams.stream_freed(self.id, StreamHalf::Recv);
-                    self.state = ChunksState::Finished;
+                    let state = mem::replace(&mut self.state, ChunksState::Finished);
+                    // At this point if we have `rs` self.state must be `ChunksState::Readable`
+                    let recv = match state {
+                        ChunksState::Readable(recv) => StreamRecv::Open(recv),
+                        _ => unreachable!("state must be ChunkState::Readable"),
+                    };
+                    self.streams.stream_recv_freed(self.id, recv);
                     Ok(None)
                 } else {
                     // We don't need a distinct `ChunksState` variant for a blocked stream because
@@ -345,7 +366,9 @@ impl<'a> Chunks<'a> {
                 self.pending.max_stream_data.insert(self.id);
             }
             // Return the stream to storage for future use
-            self.streams.recv.insert(self.id, Some(rs));
+            self.streams
+                .recv
+                .insert(self.id, Some(StreamRecv::Open(rs)));
         }
 
         // Issue connection-level flow control credit for any data we read regardless of state
diff --git a/quinn-proto/src/connection/streams/state.rs b/quinn-proto/src/connection/streams/state.rs
@@ -20,12 +20,58 @@ use crate::{
     Dir, Side, StreamId, TransportError, VarInt, MAX_STREAM_COUNT,
 };
 
+/// Wrapper around `Recv` that facilitates reusing `Recv` instances
+#[derive(Debug)]
+pub(super) enum StreamRecv {
+    /// A `Recv` that is ready to be opened
+    Free(Box<Recv>),
+    /// A `Recv` that has been opened
+    Open(Box<Recv>),
+}
+
+impl StreamRecv {
+    /// Returns a reference to the inner `Recv` if the stream is open
+    pub(super) fn as_open_recv(&self) -> Option<&Recv> {
+        match self {
+            Self::Open(r) => Some(r),
+            _ => None,
+        }
+    }
+
+    // Returns a mutable reference to the inner `Recv` if the stream is open
+    pub(super) fn as_open_recv_mut(&mut self) -> Option<&mut Recv> {
+        match self {
+            Self::Open(r) => Some(r),
+            _ => None,
+        }
+    }
+
+    // Returns the inner `Recv`
+    pub(super) fn into_inner(self) -> Box<Recv> {
+        match self {
+            Self::Free(r) | Self::Open(r) => r,
+        }
+    }
+
+    // Reinitialize the stream so the inner `Recv` can be reused
+    pub(super) fn free(self, initial_max_data: u64) -> Self {
+        match self {
+            Self::Free(_) => unreachable!("Self::Free on reinit()"),
+            Self::Open(mut recv) => {
+                recv.reinit(initial_max_data);
+                Self::Free(recv)
+            }
+        }
+    }
+}
+
 #[allow(unreachable_pub)] // fuzzing only
 pub struct StreamsState {
     pub(super) side: Side,
     // Set of streams that are currently open, or could be immediately opened by the peer
     pub(super) send: FxHashMap<StreamId, Option<Box<Send>>>,
-    pub(super) recv: FxHashMap<StreamId, Option<Box<Recv>>>,
+    pub(super) recv: FxHashMap<StreamId, Option<StreamRecv>>,
+    pub(super) free_recv: Vec<StreamRecv>,
     pub(super) next: [u64; 2],
     /// Maximum number of locally-initiated streams that may be opened over the lifetime of the
     /// connection so far, per direction
@@ -105,6 +151,7 @@ impl StreamsState {
             side,
             send: FxHashMap::default(),
             recv: FxHashMap::default(),
+            free_recv: Vec::new(),
             next: [0, 0],
             max: [0, 0],
             max_remote: [max_remote_bi.into(), max_remote_uni.into()],
@@ -240,8 +287,8 @@ impl StreamsState {
 
         // Stopped streams become closed instantly on FIN, so check whether we need to clean up
         if closed {
-            self.recv.remove(&id);
-            self.stream_freed(id, StreamHalf::Recv);
+            let rs = self.recv.remove(&id).flatten().unwrap();
+            self.stream_recv_freed(id, rs);
         }
 
         // We don't buffer data on stopped streams, so issue flow control credit immediately
@@ -293,8 +340,8 @@ impl StreamsState {
         let end = rs.end;
         if stopped {
             // Stopped streams should be disposed immediately on reset
-            self.recv.remove(&id);
-            self.stream_freed(id, StreamHalf::Recv);
+            let rs = self.recv.remove(&id).flatten().unwrap();
+            self.stream_recv_freed(id, rs);
         }
         self.on_stream_frame(!stopped, id);
 
@@ -358,6 +405,7 @@ impl StreamsState {
         self.recv
             .get(&id)
             .and_then(|s| s.as_ref())
+            .and_then(|s| s.as_open_recv())
             .map_or(false, |s| s.can_send_flow_control())
     }
 
@@ -442,7 +490,12 @@ impl StreamsState {
                 None => break,
             };
             pending.max_stream_data.remove(&id);
-            let rs = match self.recv.get_mut(&id).and_then(|s| s.as_mut()) {
+            let rs = match self
+                .recv
+                .get_mut(&id)
+                .and_then(|s| s.as_mut())
+                .and_then(|s| s.as_open_recv_mut())
+            {
                 Some(x) => x,
                 None => continue,
             };
@@ -831,7 +884,8 @@ impl StreamsState {
         }
         // bidirectional OR (unidirectional AND remote)
         if bi || remote {
-            assert!(self.recv.insert(id, None).is_none());
+            let recv = self.free_recv.pop();
+            assert!(self.recv.insert(id, recv).is_none());
         }
     }
 
@@ -883,6 +937,11 @@ impl StreamsState {
         }
     }
 
+    pub(super) fn stream_recv_freed(&mut self, id: StreamId, recv: StreamRecv) {
+        self.free_recv.push(recv.free(self.stream_receive_window));
+        self.stream_freed(id, StreamHalf::Recv);
+    }
+
     pub(super) fn max_send_data(&self, id: StreamId) -> VarInt {
         let remote = self.side != id.initiator();
         match id.dir() {
@@ -905,8 +964,16 @@ pub(super) fn get_or_insert_send(
 #[inline]
 pub(super) fn get_or_insert_recv(
     initial_max_data: u64,
-) -> impl Fn(&mut Option<Box<Recv>>) -> &mut Box<Recv> {
-    move |opt| opt.get_or_insert_with(|| Recv::new(initial_max_data))
+) -> impl FnMut(&mut Option<StreamRecv>) -> &mut Recv {
+    move |opt| {
+        *opt = opt.take().map(|s| match s {
+            StreamRecv::Free(recv) => StreamRecv::Open(recv),
+            s => s,
+        });
+        opt.get_or_insert_with(|| StreamRecv::Open(Recv::new(initial_max_data)))
+            .as_open_recv_mut()
+            .unwrap()
+    }
 }
 
 #[cfg(test)]
