From 8c89bfb751e3ec584be548cea1488e00f328090d Mon Sep 17 00:00:00 2001
From: Matthias Einwag <matthias.einwag@live.com>
Date: Sat, 22 May 2021 15:44:54 -0700
Subject: [PATCH] Limit the amount of datagrams produced per connection driver
 iteration
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Creating datagrams using `poll_transmit` is rather CPU intensive,
and while doing that no other events can be handled (e.g. processing new
incoming ACKs). This change places an upper bound on the amount of
datagrams produced.

This significantly reduces RTT for a loopback connection since the connection
is no longer busy producing packets.

Raw throughput can be a bit slower since a lot more ACK packets are now
processed in the loopback test. However RTT is halfed. The difference
in ACK packets (7200 vs 83000) and RTT (800us vs 400us) is very
visible in the benchmark:

**Before:**

```
Overall stats:

Sent 1073741824 bytes on 1 streams in 1.70s (603.88 MiB/s)

Server connection stats:
ConnectionStats {
    frame_tx: FrameStats {
        ACK: 7199,
    },
    frame_rx: FrameStats {
        ACK: 903129,
    },
    path: PathStats {
        rtt: 459.755µs,
        cwnd: 12320,
    },
}

Client connection stats:
ConnectionStats {
    frame_tx: FrameStats {
        ACK: 903129,
    },
    frame_rx: FrameStats {
        ACK: 7199,
    },
    path: PathStats {
        rtt: 769.84µs,
        cwnd: 1094264975,
    },
}
```

**After:**
```
Sent 5389680640 bytes on 1 streams in 8.51s (603.90 MiB/s)

Server connection stats:
ConnectionStats {
    frame_tx: FrameStats {
        ACK: 83555,
    },
    frame_rx: FrameStats {
        ACK: 4563602,
        STREAM: 4562497,
    },
    path: PathStats {
        rtt: 205.582µs,
        cwnd: 12320,
    },
}

Client connection stats:
ConnectionStats {
    frame_tx: FrameStats {
        ACK: 4563602,
        STREAM: 4562497,
    },
    frame_rx: FrameStats {
        ACK: 83555,
    },
    path: PathStats {
        rtt: 393.448µs,
        cwnd: 495821204,
    },
}
```
---
 quinn/src/connection.rs | 54 ++++++++++++++++++++++++++++-------------
 1 file changed, 37 insertions(+), 17 deletions(-)

diff --git a/quinn/src/connection.rs b/quinn/src/connection.rs
index a6cd73f22..fabea1d6f 100644
--- a/quinn/src/connection.rs
+++ b/quinn/src/connection.rs
@@ -287,25 +287,24 @@ where
         let span = info_span!("drive", id = conn.handle.0);
         let _guard = span.enter();
 
-        loop {
-            let mut keep_going = false;
-            if let Err(e) = conn.process_conn_events(cx) {
-                conn.terminate(e);
-                return Poll::Ready(());
-            }
-            conn.drive_transmit();
-            // If a timer expires, there might be more to transmit. When we transmit something, we
-            // might need to reset a timer. Hence, we must loop until neither happens.
-            keep_going |= conn.drive_timer(cx);
-            conn.forward_endpoint_events();
-            conn.forward_app_events();
-            if !keep_going || conn.inner.is_drained() {
-                break;
-            }
+        if let Err(e) = conn.process_conn_events(cx) {
+            conn.terminate(e);
+            return Poll::Ready(());
         }
+        let mut keep_going = conn.drive_transmit();
+        // If a timer expires, there might be more to transmit. When we transmit something, we
+        // might need to reset a timer. Hence, we must loop until neither happens.
+        keep_going |= conn.drive_timer(cx);
+        conn.forward_endpoint_events();
+        conn.forward_app_events();
 
         if !conn.inner.is_drained() {
-            conn.driver = Some(cx.waker().clone());
+            if keep_going {
+                // If the connection hasn't processed all tasks, schedule it again
+                cx.waker().wake_by_ref();
+            } else {
+                conn.driver = Some(cx.waker().clone());
+            }
             return Poll::Pending;
         }
         if conn.error.is_none() {
@@ -799,17 +798,32 @@ impl<S> ConnectionInner<S>
 where
     S: proto::crypto::Session,
 {
-    fn drive_transmit(&mut self) {
+    fn drive_transmit(&mut self) -> bool {
         let now = Instant::now();
+        let mut transmits = 0;
 
         let max_datagrams = caps().max_gso_segments;
 
         while let Some(t) = self.inner.poll_transmit(now, max_datagrams) {
+            transmits += match t.segment_size {
+                None => 1,
+                Some(s) => (t.contents.len() + s - 1) / s, // round up
+            };
             // If the endpoint driver is gone, noop.
             let _ = self
                 .endpoint_events
                 .unbounded_send((self.handle, EndpointEvent::Transmit(t)));
+
+            if transmits >= MAX_TRANSMIT_DATAGRAMS {
+                // TODO: What isn't ideal here yet is that if we don't poll all
+                // datagrams that could be sent we don't go into the `app_limited`
+                // state and CWND continues to grow until we get here the next time.
+                // See https://github.com/quinn-rs/quinn/issues/1126
+                return true;
+            }
         }
+
+        false
     }
 
     fn forward_endpoint_events(&mut self) {
@@ -1073,3 +1087,9 @@ pub enum SendDatagramError {
     #[error("connection closed: {0}")]
     ConnectionClosed(#[source] ConnectionError),
 }
+
+/// The maximum amount of datagrams which will be produced in a single `drive_transmit` call
+///
+/// This limits the amount of CPU resources consumed by datagram generation,
+/// and allows other tasks (like receiving ACKs) to run in between.
+const MAX_TRANSMIT_DATAGRAMS: usize = 20;