From 578cbb1eb0dc7f90cbd5d9dec4011441fc85f56c Mon Sep 17 00:00:00 2001
From: Matthias Einwag <matthias.einwag@live.com>
Date: Sat, 18 Sep 2021 17:37:39 -0700
Subject: [PATCH] Use WorkLimiter also for sending data

This adds time-based yielding to the send loop in the same fashion it
had been previously added ot the receive loop.
In my performance testing this didn't show a noticeable difference - likely
because in the current benchmark the client is the bottleneck. But it should
make things more deterministic.
---
 quinn/src/endpoint.rs | 31 ++++++++++++++++++++-----------
 quinn/src/lib.rs      |  3 +++
 2 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/quinn/src/endpoint.rs b/quinn/src/endpoint.rs
index 734410e80..316916cc1 100644
--- a/quinn/src/endpoint.rs
+++ b/quinn/src/endpoint.rs
@@ -29,7 +29,7 @@ use crate::{
     connection::Connecting,
     platform::{RecvMeta, UdpSocket, BATCH_SIZE},
     work_limiter::WorkLimiter,
-    ConnectionEvent, EndpointEvent, VarInt, IO_LOOP_BOUND, RECV_TIME_BOUND,
+    ConnectionEvent, EndpointEvent, VarInt, IO_LOOP_BOUND, RECV_TIME_BOUND, SEND_TIME_BOUND,
 };
 
 /// A QUIC endpoint.
@@ -284,6 +284,7 @@ where
     driver_lost: bool,
     recv_limiter: WorkLimiter,
     recv_buf: Box<[u8]>,
+    send_limiter: WorkLimiter,
     idle: Broadcast,
 }
 
@@ -355,35 +356,42 @@ where
     }
 
     fn drive_send(&mut self, cx: &mut Context) -> Result<bool, io::Error> {
-        let mut transmits = 0;
-        loop {
+        self.send_limiter.start_cycle();
+
+        let result = loop {
             while self.outgoing.len() < BATCH_SIZE {
                 match self.inner.poll_transmit() {
                     Some(x) => self.outgoing.push_back(x),
                     None => break,
                 }
             }
+
             if self.outgoing.is_empty() {
-                return Ok(false);
+                break Ok(false);
             }
+
+            if !self.send_limiter.allow_work() {
+                break Ok(true);
+            }
+
             match self.socket.poll_send(cx, self.outgoing.as_slices().0) {
                 Poll::Ready(Ok(n)) => {
                     self.outgoing.drain(..n);
                     // We count transmits instead of `poll_send` calls since the cost
                     // of a `sendmmsg` still linearily increases with number of packets.
-                    transmits += n;
-                    if transmits >= IO_LOOP_BOUND {
-                        return Ok(true);
-                    }
+                    self.send_limiter.record_work(n);
                 }
                 Poll::Pending => {
-                    return Ok(false);
+                    break Ok(false);
                 }
                 Poll::Ready(Err(e)) => {
-                    return Err(e);
+                    break Err(e);
                 }
             }
-        }
+        };
+
+        self.send_limiter.finish_cycle();
+        result
     }
 
     fn handle_events(&mut self, cx: &mut Context) -> bool {
@@ -537,6 +545,7 @@ where
             driver_lost: false,
             recv_buf: recv_buf.into(),
             recv_limiter: WorkLimiter::new(RECV_TIME_BOUND),
+            send_limiter: WorkLimiter::new(SEND_TIME_BOUND),
             idle: Broadcast::new(),
         })))
     }
diff --git a/quinn/src/lib.rs b/quinn/src/lib.rs
index e0b716438..e75d1725a 100644
--- a/quinn/src/lib.rs
+++ b/quinn/src/lib.rs
@@ -171,3 +171,6 @@ const IO_LOOP_BOUND: usize = 160;
 /// Going much lower does not yield any noticeable difference, since a single `recvmmsg`
 /// batch of size 32 was observed to take 30us on some systems.
 const RECV_TIME_BOUND: Duration = Duration::from_micros(50);
+
+/// The maximum amount of time that should be spent in `sendmsg()` calls per endpoint iteration
+const SEND_TIME_BOUND: Duration = Duration::from_micros(50);