From 74ff39ec6d3ed1c0dfd5fefc9f9e2e32c2191477 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Tue, 22 Apr 2025 11:35:45 +1000
Subject: [PATCH] fix(connlib): shortcut `DatagramSegmentIter` on `len` (#8877)

When the `recv` syscall completes, `quinn-udp` tells us how many batches
we have read. On Windows, this is always 1 because Windows doesn't have
an APIs to read more than a single GRO batch. The `DatagramSegmentIter`
already has a way of detecting this, however it currently needs to
iterator through all batches (10) and check that their `meta.length ==
0` before realising this.

We can shortcut the iterator early which might improve download
performance on Windows.

I can't measure a direct improvement here but I believe that is because
we are currently limited by the buffer size on Windows. Regardless, this
feels like the right thing to do.
---
 rust/socket-factory/src/lib.rs | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/rust/socket-factory/src/lib.rs b/rust/socket-factory/src/lib.rs
index 4404cce9b..e1b704a85 100644
--- a/rust/socket-factory/src/lib.rs
+++ b/rust/socket-factory/src/lib.rs
@@ -267,9 +267,7 @@ impl UdpSocket {
                     ],
                 );
 
-                // Note: We don't need to use `len` here because the iterator will stop once it encounteres `meta.len == 0`.
-
-                return Poll::Ready(Ok(DatagramSegmentIter::new(bufs, meta, *port)));
+                return Poll::Ready(Ok(DatagramSegmentIter::new(bufs, meta, *port, len)));
             }
         }
     }
@@ -476,6 +474,7 @@ pub struct DatagramSegmentIter<
     #[debug(skip)]
     buffers: [B; N],
     metas: [quinn_udp::RecvMeta; N],
+    len: usize,
 
     port: u16,
 
@@ -487,7 +486,7 @@ pub struct DatagramSegmentIter<
 }
 
 impl<B, const N: usize> DatagramSegmentIter<N, B> {
-    fn new(buffers: [B; N], metas: [quinn_udp::RecvMeta; N], port: u16) -> Self {
+    fn new(buffers: [B; N], metas: [quinn_udp::RecvMeta; N], port: u16, len: usize) -> Self {
         let total_bytes = metas.iter().map(|m| m.len).sum::<usize>();
         let num_packets = metas
             .iter()
@@ -503,6 +502,7 @@ impl<B, const N: usize> DatagramSegmentIter<N, B> {
         Self {
             buffers,
             metas,
+            len,
             port,
             buf_index: 0,
             segment_index: 0,
@@ -520,7 +520,7 @@ where
 
     fn next(&mut self) -> Option<Self::Item<'_>> {
         loop {
-            if self.buf_index >= N {
+            if self.buf_index >= N || self.buf_index >= self.len {
                 return None;
             }
 
@@ -621,6 +621,7 @@ mod tests {
                 quinn_udp::RecvMeta::default(),
             ],
             0,
+            3,
         );
 
         assert_eq!(iter.next().unwrap().packet, b"foobar1");