From 130de254acff6640eb1bba2f67c14eceb09fc8a6 Mon Sep 17 00:00:00 2001
From: Benjamin Saunders <ben.e.saunders@gmail.com>
Date: Fri, 21 Jul 2023 15:01:15 -0700
Subject: [PATCH] Fix excessive receive buffer allocation

Allocation logic previously considered a GRO batch to be one datagram,
rather than (when supported) 64. This led to an unintended ~128MiB
allocation, now scaled back down to ~4MiB.
---
 quinn/src/endpoint.rs | 34 +++++++++++++++++++++++-----------
 1 file changed, 23 insertions(+), 11 deletions(-)

diff --git a/quinn/src/endpoint.rs b/quinn/src/endpoint.rs
index d5e1e2406..965f97ce5 100644
--- a/quinn/src/endpoint.rs
+++ b/quinn/src/endpoint.rs
@@ -384,20 +384,20 @@ pub(crate) struct Shared {
 impl State {
     fn drive_recv<'a>(&'a mut self, cx: &mut Context, now: Instant) -> Result<bool, io::Error> {
         self.recv_limiter.start_cycle();
+        let recv_count = recv_datagram_count(&self.udp_state) / self.udp_state.gro_segments();
         let mut metas = [RecvMeta::default(); BATCH_SIZE];
-        let mut iovs = MaybeUninit::<[IoSliceMut<'a>; BATCH_SIZE]>::uninit();
+        let mut iovs = std::array::from_fn::<MaybeUninit<IoSliceMut<'a>>, BATCH_SIZE, _>(|_| {
+            MaybeUninit::uninit()
+        });
         self.recv_buf
-            .chunks_mut(self.recv_buf.len() / BATCH_SIZE)
+            .chunks_mut(self.recv_buf.len() / recv_count)
             .enumerate()
-            .for_each(|(i, buf)| unsafe {
-                iovs.as_mut_ptr()
-                    .cast::<IoSliceMut>()
-                    .add(i)
-                    .write(IoSliceMut::<'a>::new(buf));
+            .for_each(|(i, buf)| {
+                iovs[i].write(IoSliceMut::<'a>::new(buf));
             });
-        let mut iovs = unsafe { iovs.assume_init() };
+        let iovs = unsafe { slice_assume_init_mut(&mut iovs[..recv_count]) };
         loop {
-            match self.socket.poll_recv(cx, &mut iovs, &mut metas) {
+            match self.socket.poll_recv(cx, iovs, &mut metas) {
                 Poll::Ready(Ok(msgs)) => {
                     self.recv_limiter.record_work(msgs);
                     for (meta, buf) in metas.iter().zip(iovs.iter()).take(msgs) {
@@ -552,6 +552,12 @@ impl State {
     }
 }
 
+// TODO: Remove once feature maybe_uninit_slice (https://github.com/rust-lang/rust/issues/63569)
+// stabilizes
+unsafe fn slice_assume_init_mut<T>(x: &mut [MaybeUninit<T>]) -> &mut [T] {
+    std::slice::from_raw_parts_mut(x.as_mut_ptr().cast(), x.len())
+}
+
 #[inline]
 fn udp_transmit(t: proto::Transmit) -> udp::Transmit {
     udp::Transmit {
@@ -673,8 +679,7 @@ impl EndpointRef {
         let recv_buf = vec![
             0;
             inner.config().get_max_udp_payload_size().min(64 * 1024) as usize
-                * udp_state.gro_segments()
-                * BATCH_SIZE
+                * recv_datagram_count(&udp_state)
         ];
         let (sender, events) = mpsc::unbounded_channel();
         Self(Arc::new(EndpointInner {
@@ -737,3 +742,10 @@ impl std::ops::Deref for EndpointRef {
         &self.0
     }
 }
+
+/// Number of datagrams to allocate space for per recv
+fn recv_datagram_count(udp: &UdpState) -> usize {
+    // We must receive in multiples of a GRO batch to avoid datagram truncation, and we want to
+    // accept at least `BATCH_SIZE` packets per receive.
+    usize::max(1, BATCH_SIZE / udp.gro_segments()) * udp.gro_segments()
+}