From d546b905701ea83ac44d69eede5d6ccb7fb5e51d Mon Sep 17 00:00:00 2001 From: Gus Caplan Date: Sat, 13 Apr 2024 16:19:06 +1000 Subject: [PATCH] add drop param to datagram send --- quinn-proto/src/connection/datagrams.rs | 41 +++++++++++++++++-------- quinn-proto/src/connection/mod.rs | 8 +++++ quinn-proto/src/tests/mod.rs | 24 ++++++++++----- quinn/src/connection.rs | 4 ++- 4 files changed, 57 insertions(+), 20 deletions(-) diff --git a/quinn-proto/src/connection/datagrams.rs b/quinn-proto/src/connection/datagrams.rs index abadbe916..5d722cdef 100644 --- a/quinn-proto/src/connection/datagrams.rs +++ b/quinn-proto/src/connection/datagrams.rs @@ -19,27 +19,40 @@ pub struct Datagrams<'a> { impl<'a> Datagrams<'a> { /// Queue an unreliable, unordered datagram for immediate transmission /// - /// Returns `Err` iff a `len`-byte datagram cannot currently be sent - pub fn send(&mut self, data: Bytes) -> Result<(), SendDatagramError> { + /// If `drop` is true, previously queued datagrams which are still unsent may be discarded to + /// make space for this datagram, in order of oldest to newest. If `drop` is false, and there + /// isn't enough space due to previously queued datagrams, this function will return + /// `SendDatagramError::Blocked`. `Event::DatagramsUnblocked` will be emitted once datagrams + /// have been sent. + /// + /// Returns `Err` iff a `len`-byte datagram cannot currently be sent. + pub fn send(&mut self, data: Bytes, drop: bool) -> Result<(), SendDatagramError> { if self.conn.config.datagram_receive_buffer_size.is_none() { return Err(SendDatagramError::Disabled); } let max = self .max_size() .ok_or(SendDatagramError::UnsupportedByPeer)?; - while self.conn.datagrams.outgoing_total > self.conn.config.datagram_send_buffer_size { - let prev = self - .conn - .datagrams - .outgoing - .pop_front() - .expect("datagrams.outgoing_total desynchronized"); - trace!(len = prev.data.len(), "dropping outgoing datagram"); - self.conn.datagrams.outgoing_total -= prev.data.len(); - } if data.len() > max { return Err(SendDatagramError::TooLarge); } + if drop { + while self.conn.datagrams.outgoing_total > self.conn.config.datagram_send_buffer_size { + let prev = self + .conn + .datagrams + .outgoing + .pop_front() + .expect("datagrams.outgoing_total desynchronized"); + trace!(len = prev.data.len(), "dropping outgoing datagram"); + self.conn.datagrams.outgoing_total -= prev.data.len(); + } + } else if self.conn.datagrams.outgoing_total + data.len() + > self.conn.config.datagram_send_buffer_size + { + self.conn.datagrams.send_blocked = true; + return Err(SendDatagramError::Blocked(data)); + } self.conn.datagrams.outgoing_total += data.len(); self.conn.datagrams.outgoing.push_back(Datagram { data }); Ok(()) @@ -95,6 +108,7 @@ pub(super) struct DatagramState { pub(super) incoming: VecDeque, pub(super) outgoing: VecDeque, pub(super) outgoing_total: usize, + pub(super) send_blocked: bool, } impl DatagramState { @@ -167,4 +181,7 @@ pub enum SendDatagramError { /// exceeded. #[error("datagram too large")] TooLarge, + /// Send would block + #[error("datagram send blocked")] + Blocked(Bytes), } diff --git a/quinn-proto/src/connection/mod.rs b/quinn-proto/src/connection/mod.rs index a2dee54ec..29b5f5c5d 100644 --- a/quinn-proto/src/connection/mod.rs +++ b/quinn-proto/src/connection/mod.rs @@ -3149,15 +3149,21 @@ impl Connection { } // DATAGRAM + let mut sent_datagrams = false; while buf.len() + Datagram::SIZE_BOUND < max_size && space_id == SpaceId::Data { match self.datagrams.write(buf, max_size) { true => { + sent_datagrams = true; sent.non_retransmits = true; self.stats.frame_tx.datagram += 1; } false => break, } } + if self.datagrams.send_blocked && sent_datagrams { + self.events.push_back(Event::DatagramsUnblocked); + self.datagrams.send_blocked = false; + } // STREAM if space_id == SpaceId::Data { @@ -3632,6 +3638,8 @@ pub enum Event { Stream(StreamEvent), /// One or more application datagrams have been received DatagramReceived, + /// One or more application datagrams have been sent after blocking + DatagramsUnblocked, } fn instant_saturating_sub(x: Instant, y: Instant) -> Duration { diff --git a/quinn-proto/src/tests/mod.rs b/quinn-proto/src/tests/mod.rs index 62b553936..3d97dbac5 100644 --- a/quinn-proto/src/tests/mod.rs +++ b/quinn-proto/src/tests/mod.rs @@ -1633,7 +1633,9 @@ fn datagram_send_recv() { assert_matches!(pair.client_datagrams(client_ch).max_size(), Some(x) if x > 0); const DATA: &[u8] = b"whee"; - pair.client_datagrams(client_ch).send(DATA.into()).unwrap(); + pair.client_datagrams(client_ch) + .send(DATA.into(), true) + .unwrap(); pair.drive(); assert_matches!( pair.server_conn_mut(server_ch).poll(), @@ -1665,9 +1667,15 @@ fn datagram_recv_buffer_overflow() { const DATA1: &[u8] = &[0xAB; (WINDOW / 3) + 1]; const DATA2: &[u8] = &[0xBC; (WINDOW / 3) + 1]; const DATA3: &[u8] = &[0xCD; (WINDOW / 3) + 1]; - pair.client_datagrams(client_ch).send(DATA1.into()).unwrap(); - pair.client_datagrams(client_ch).send(DATA2.into()).unwrap(); - pair.client_datagrams(client_ch).send(DATA3.into()).unwrap(); + pair.client_datagrams(client_ch) + .send(DATA1.into(), true) + .unwrap(); + pair.client_datagrams(client_ch) + .send(DATA2.into(), true) + .unwrap(); + pair.client_datagrams(client_ch) + .send(DATA3.into(), true) + .unwrap(); pair.drive(); assert_matches!( pair.server_conn_mut(server_ch).poll(), @@ -1677,7 +1685,9 @@ fn datagram_recv_buffer_overflow() { assert_eq!(pair.server_datagrams(server_ch).recv().unwrap(), DATA3); assert_matches!(pair.server_datagrams(server_ch).recv(), None); - pair.client_datagrams(client_ch).send(DATA1.into()).unwrap(); + pair.client_datagrams(client_ch) + .send(DATA1.into(), true) + .unwrap(); pair.drive(); assert_eq!(pair.server_datagrams(server_ch).recv().unwrap(), DATA1); assert_matches!(pair.server_datagrams(server_ch).recv(), None); @@ -1698,7 +1708,7 @@ fn datagram_unsupported() { assert_matches!(pair.server_conn_mut(server_ch).poll(), None); assert_matches!(pair.client_datagrams(client_ch).max_size(), None); - match pair.client_datagrams(client_ch).send(Bytes::new()) { + match pair.client_datagrams(client_ch).send(Bytes::new(), true) { Err(SendDatagramError::UnsupportedByPeer) => {} Err(e) => panic!("unexpected error: {e}"), Ok(_) => panic!("unexpected success"), @@ -2799,7 +2809,7 @@ fn pure_sender_voluntarily_acks() { for _ in 0..100 { const MSG: &[u8] = b"hello"; pair.client_datagrams(client_ch) - .send(Bytes::from_static(MSG)) + .send(Bytes::from_static(MSG), true) .unwrap(); pair.drive(); assert_eq!(pair.server_datagrams(server_ch).recv().unwrap(), MSG); diff --git a/quinn/src/connection.rs b/quinn/src/connection.rs index d3d620da8..76d8532b5 100644 --- a/quinn/src/connection.rs +++ b/quinn/src/connection.rs @@ -392,12 +392,13 @@ impl Connection { return Err(SendDatagramError::ConnectionLost(x.clone())); } use proto::SendDatagramError::*; - match conn.inner.datagrams().send(data) { + match conn.inner.datagrams().send(data, true) { Ok(()) => { conn.wake(); Ok(()) } Err(e) => Err(match e { + Blocked(..) => unreachable!(), UnsupportedByPeer => SendDatagramError::UnsupportedByPeer, Disabled => SendDatagramError::Disabled, TooLarge => SendDatagramError::TooLarge, @@ -971,6 +972,7 @@ impl State { DatagramReceived => { shared.datagrams.notify_waiters(); } + DatagramsUnblocked => {} Stream(StreamEvent::Readable { id }) => { if let Some(reader) = self.blocked_readers.remove(&id) { reader.wake();