backport supress stateless packets to 0.9.x (#1601)

quinn-proto/src/endpoint.rs

@@ -61,8 +61,20 @@ pub struct Endpoint {
     local_cid_generator: Box<dyn ConnectionIdGenerator>,
     config: Arc<EndpointConfig>,
     server_config: Option<Arc<ServerConfig>>,
+    /// The contents length for packets in the transmits queue
+    transmit_queue_contents_len: usize,
+    /// The socket buffer aggregated contents length
+    /// `transmit_queue_contents_len` + `socket_buffer_fill` represents the total contents length
+    /// of outstanding outgoing packets.
+    socket_buffer_fill: usize,
 }
 
+/// The maximum size of content length of packets in the outgoing transmit queue. Transmit packets
+/// generated from the endpoint (retry, initial close, stateless reset and version negotiation)
+/// can be dropped when this limit is being execeeded.
+/// Chose to represent 100 MB of data.
+const MAX_TRANSMIT_QUEUE_CONTENTS_LEN: usize = 100_000_000;
+
 impl Endpoint {
     /// Create a new endpoint
     ///
@@ -79,13 +91,17 @@ impl Endpoint {
             local_cid_generator: (config.connection_id_generator_factory.as_ref())(),
             config,
             server_config,
+            transmit_queue_contents_len: 0,
+            socket_buffer_fill: 0,
         }
     }
 
     /// Get the next packet to transmit
     #[must_use]
     pub fn poll_transmit(&mut self) -> Option<Transmit> {
-        self.transmits.pop_front()
+        let t = self.transmits.pop_front();
+        self.decrement_transmit_queue_contents_len(t.as_ref().map_or(0, |t| t.contents.len()));
+        t
     }
 
     /// Replace the server configuration, affecting new incoming connections only
@@ -166,6 +182,9 @@ impl Endpoint {
                     debug!("dropping packet with unsupported version");
                     return None;
                 }
+                if self.stateless_packets_supressed() {
+                    return None;
+                }
                 trace!("sending version negotiation");
                 // Negotiate versions
                 let mut buf = Vec::<u8>::new();
@@ -184,6 +203,7 @@ impl Endpoint {
                 for &version in &self.config.supported_versions {
                     buf.write(version);
                 }
+                self.increment_transmit_queue_contents_len(buf.len());
                 self.transmits.push_back(Transmit {
                     destination: remote,
                     ecn: None,
@@ -314,6 +334,9 @@ impl Endpoint {
         addresses: FourTuple,
         dst_cid: &ConnectionId,
     ) {
+        if self.stateless_packets_supressed() {
+            return;
+        }
         /// Minimum amount of padding for the stateless reset to look like a short-header packet
         const MIN_PADDING_LEN: usize = 5;
 
@@ -346,7 +369,7 @@ impl Endpoint {
         buf.extend_from_slice(&ResetToken::new(&*self.config.reset_key, dst_cid));
 
         debug_assert!(buf.len() < inciting_dgram_len);
-
+        self.increment_transmit_queue_contents_len(buf.len());
         self.transmits.push_back(Transmit {
             destination: addresses.remote,
             ecn: None,
@@ -435,6 +458,35 @@ impl Endpoint {
         }
     }
 
+    /// Limiting the memory usage for items queued in the outgoing queue from endpoint
+    /// generated packets. Otherwise, we may see a build-up of the queue under test with
+    /// flood of initial packets against the endpoint. The sender with the sender-limiter
+    /// may not keep up the pace of these packets queued into the queue.
+    fn stateless_packets_supressed(&self) -> bool {
+        self.transmit_queue_contents_len
+            .saturating_add(self.socket_buffer_fill)
+            >= MAX_TRANSMIT_QUEUE_CONTENTS_LEN
+    }
+
+    /// Increment the contents length in the transmit queue.
+    fn increment_transmit_queue_contents_len(&mut self, contents_len: usize) {
+        self.transmit_queue_contents_len = self
+            .transmit_queue_contents_len
+            .saturating_add(contents_len);
+    }
+
+    /// Decrement the contents length in the transmit queue.
+    fn decrement_transmit_queue_contents_len(&mut self, contents_len: usize) {
+        self.transmit_queue_contents_len = self
+            .transmit_queue_contents_len
+            .saturating_sub(contents_len);
+    }
+
+    /// Set the `socket_buffer_fill` to the input `len`
+    pub fn set_socket_buffer_fill(&mut self, len: usize) {
+        self.socket_buffer_fill = len;
+    }
+
     fn handle_first_packet(
         &mut self,
         now: Instant,
@@ -509,6 +561,9 @@ impl Endpoint {
 
         let (retry_src_cid, orig_dst_cid) = if server_config.use_retry {
             if token.is_empty() {
+                if self.stateless_packets_supressed() {
+                    return None;
+                }
                 // First Initial
                 let mut random_bytes = vec![0u8; RetryToken::RANDOM_BYTES_LEN];
                 self.rng.fill_bytes(&mut random_bytes);
@@ -532,6 +587,7 @@ impl Endpoint {
                 buf.extend_from_slice(&server_config.crypto.retry_tag(version, &dst_cid, &buf));
                 encode.finish(&mut buf, &*crypto.header.local, None);
 
+                self.increment_transmit_queue_contents_len(buf.len());
                 self.transmits.push_back(Transmit {
                     destination: addresses.remote,
                     ecn: None,
@@ -674,6 +730,9 @@ impl Endpoint {
         local_id: &ConnectionId,
         reason: TransportError,
     ) {
+        if self.stateless_packets_supressed() {
+            return;
+        }
         let number = PacketNumber::U8(0);
         let header = Header::Initial {
             dst_cid: *remote_id,
@@ -695,6 +754,7 @@ impl Endpoint {
             &*crypto.header.local,
             Some((0, &*crypto.packet.local)),
         );
+        self.increment_transmit_queue_contents_len(buf.len());
         self.transmits.push_back(Transmit {
             destination: addresses.remote,
             ecn: None,

quinn/src/endpoint.rs

@@ -361,6 +361,8 @@ pub(crate) struct State {
     recv_buf: Box<[u8]>,
     send_limiter: WorkLimiter,
     runtime: Arc<dyn Runtime>,
+    /// The packet contents length in the outgoing queue.
+    outgoing_queue_contents_len: usize,
 }
 
 #[derive(Debug)]
@@ -447,7 +449,7 @@ impl State {
         let result = loop {
             while self.outgoing.len() < BATCH_SIZE {
                 match self.inner.poll_transmit() {
-                    Some(x) => self.outgoing.push_back(x),
+                    Some(t) => self.queue_transmit(t),
                     None => break,
                 }
             }
@@ -465,7 +467,9 @@ impl State {
                 .poll_send(&self.udp_state, cx, self.outgoing.as_slices().0)
             {
                 Poll::Ready(Ok(n)) => {
-                    self.outgoing.drain(..n);
+                    let contents_len: usize =
+                        self.outgoing.drain(..n).map(|t| t.contents.len()).sum();
+                    self.decrement_outgoing_contents_len(contents_len);
                     // We count transmits instead of `poll_send` calls since the cost
                     // of a `sendmmsg` still linearily increases with number of packets.
                     self.send_limiter.record_work(n);
@@ -506,7 +510,7 @@ impl State {
                                 .send(ConnectionEvent::Proto(event));
                         }
                     }
-                    Transmit(t) => self.outgoing.push_back(t),
+                    Transmit(t) => self.queue_transmit(t),
                 },
                 Poll::Ready(None) => unreachable!("EndpointInner owns one sender"),
                 Poll::Pending => {
@@ -517,6 +521,28 @@ impl State {
 
         true
     }
+
+    fn queue_transmit(&mut self, t: proto::Transmit) {
+        let contents_len = t.contents.len();
+        self.increment_outgoing_queue_contents_len(contents_len);
+        self.outgoing.push_back(t);
+    }
+
+    fn increment_outgoing_queue_contents_len(&mut self, contents_len: usize) {
+        self.outgoing_queue_contents_len = self
+            .outgoing_queue_contents_len
+            .saturating_add(contents_len);
+        self.inner
+            .set_socket_buffer_fill(self.outgoing_queue_contents_len);
+    }
+
+    fn decrement_outgoing_contents_len(&mut self, contents_len: usize) {
+        self.outgoing_queue_contents_len = self
+            .outgoing_queue_contents_len
+            .saturating_sub(contents_len);
+        self.inner
+            .set_socket_buffer_fill(self.outgoing_queue_contents_len);
+    }
 }
 
 #[derive(Debug)]
@@ -640,6 +666,7 @@ impl EndpointRef {
                 recv_limiter: WorkLimiter::new(RECV_TIME_BOUND),
                 send_limiter: WorkLimiter::new(SEND_TIME_BOUND),
                 runtime,
+                outgoing_queue_contents_len: 0,
             }),
         }))
     }