noq_proto/connection/
datagrams.rs

1use std::collections::VecDeque;
2
3use bytes::Bytes;
4use thiserror::Error;
5use tracing::{debug, trace};
6
7use super::Connection;
8use crate::{
9    FrameStats, TransportError,
10    connection::PacketBuilder,
11    frame::{Datagram, FrameStruct},
12};
13
14/// API to control datagram traffic
15pub struct Datagrams<'a> {
16    pub(super) conn: &'a mut Connection,
17}
18
19impl Datagrams<'_> {
20    /// Queue an unreliable, unordered datagram for immediate transmission
21    ///
22    /// If `drop` is true, previously queued datagrams which are still unsent may be discarded to
23    /// make space for this datagram, in order of oldest to newest. If `drop` is false, and there
24    /// isn't enough space due to previously queued datagrams, this function will return
25    /// `SendDatagramError::Blocked`. `Event::DatagramsUnblocked` will be emitted once datagrams
26    /// have been sent.
27    ///
28    /// Returns `Err` iff a `len`-byte datagram cannot currently be sent.
29    pub fn send(&mut self, data: Bytes, drop: bool) -> Result<(), SendDatagramError> {
30        if self.conn.config.datagram_receive_buffer_size.is_none() {
31            return Err(SendDatagramError::Disabled);
32        }
33        let max = self
34            .max_size()
35            .ok_or(SendDatagramError::UnsupportedByPeer)?;
36        let send_buffer_size = self.conn.config.datagram_send_buffer_size;
37        if data.len() > Ord::min(max, send_buffer_size) {
38            return Err(SendDatagramError::TooLarge);
39        }
40        if drop {
41            self.conn
42                .datagrams
43                .make_space_for(data.len(), send_buffer_size);
44        } else if !self
45            .conn
46            .datagrams
47            .has_send_buffer_space(data.len(), send_buffer_size)
48        {
49            self.conn.datagrams.send_blocked = true;
50            return Err(SendDatagramError::Blocked(data));
51        }
52        self.conn.datagrams.outgoing_total += data.len();
53        self.conn.datagrams.outgoing.push_back(Datagram { data });
54        Ok(())
55    }
56
57    /// Queue many unreliable, unordered datagrams for transmission in a single call.
58    ///
59    /// This is the batch analogue of [`Self::send`], avoiding repeated connection checks.
60    ///
61    /// The batch is rejected atomically with the [`TooLarge`] error if any datagram
62    /// in the batch is too large.
63    ///
64    /// `drop` selects the backpressure behaviour, matching [`Self::send`]:
65    ///
66    /// - `drop = true` drops the oldest queued datagrams to make room, so every element is queued
67    ///   and `Ok(datagrams.len())` is returned.
68    /// - `drop = false` queues elements until the send buffer is full, then stops and returns
69    ///   `Ok(n)` for the `n` elements queued. The remaining elements are the caller's to retry once
70    ///   space frees up.
71    ///
72    /// Returns `Err` if datagrams are unsupported by the peer or disabled locally.
73    ///
74    /// [`TooLarge`]: SendDatagramError::TooLarge
75    pub fn send_many(
76        &mut self,
77        datagrams: &[Bytes],
78        drop: bool,
79    ) -> Result<usize, SendDatagramError> {
80        if self.conn.config.datagram_receive_buffer_size.is_none() {
81            return Err(SendDatagramError::Disabled);
82        }
83        let max = self
84            .max_size()
85            .ok_or(SendDatagramError::UnsupportedByPeer)?;
86        let send_buffer_size = self.conn.config.datagram_send_buffer_size;
87        if datagrams
88            .iter()
89            .any(|data| data.len() > Ord::min(max, send_buffer_size))
90        {
91            return Err(SendDatagramError::TooLarge);
92        }
93
94        let mut queued = 0usize;
95        for data in datagrams {
96            if drop {
97                self.conn
98                    .datagrams
99                    .make_space_for(data.len(), send_buffer_size);
100            } else if !self
101                .conn
102                .datagrams
103                .has_send_buffer_space(data.len(), send_buffer_size)
104            {
105                self.conn.datagrams.send_blocked = true;
106                break;
107            }
108            self.conn.datagrams.outgoing_total += data.len();
109            self.conn
110                .datagrams
111                .outgoing
112                .push_back(Datagram { data: data.clone() });
113            queued += 1;
114        }
115
116        Ok(queued)
117    }
118
119    /// Compute the maximum size of datagrams that may be passed to `send_datagram`
120    ///
121    /// Returns `None` if datagrams are unsupported by the peer or disabled locally.
122    ///
123    /// This may change over the lifetime of a connection according to variation in the path MTU
124    /// estimate. The peer can also enforce an arbitrarily small fixed limit, but if the peer's
125    /// limit is large this is guaranteed to be a little over a kilobyte at minimum.
126    ///
127    /// Not necessarily the maximum size of received datagrams.
128    ///
129    /// When multipath is enabled, this is calculated using the smallest MTU across all
130    /// available paths.
131    pub fn max_size(&self) -> Option<usize> {
132        // We use the conservative overhead bound for any packet number, reducing the budget by at
133        // most 3 bytes, so that PN size fluctuations don't cause users sending maximum-size
134        // datagrams to suffer avoidable packet loss.
135        let max_size = self.conn.current_mtu() as usize
136            - self.conn.predict_1rtt_overhead_no_pn()
137            - Datagram::SIZE_BOUND;
138        let limit = self
139            .conn
140            .peer_params
141            .max_datagram_frame_size?
142            .into_inner()
143            .saturating_sub(Datagram::SIZE_BOUND as u64);
144        Some(limit.min(max_size as u64) as usize)
145    }
146
147    /// Receive an unreliable, unordered datagram
148    pub fn recv(&mut self) -> Option<Bytes> {
149        self.conn.datagrams.recv()
150    }
151
152    /// Drain up to `out.len()` buffered datagrams into `out`, in arrival order.
153    ///
154    /// This is the batch analogue of [`Self::recv`]: a single call takes many
155    /// datagrams at once. `out` is filled from the front and overwritten in place;
156    /// pass a slice of empty `Bytes` sized to the batch you want. Returns the number
157    /// of datagrams written, which may be less than `out.len()` if fewer are buffered
158    /// (0 if none). Any remaining datagrams stay queued for the next call.
159    pub fn recv_many(&mut self, out: &mut [Bytes]) -> usize {
160        self.conn.datagrams.recv_many(out)
161    }
162
163    /// Bytes available in the outgoing datagram buffer
164    ///
165    /// When greater than zero, [`send`](Self::send)ing a datagram of at most this size is
166    /// guaranteed not to cause older datagrams to be dropped.
167    pub fn send_buffer_space(&self) -> usize {
168        self.conn
169            .config
170            .datagram_send_buffer_size
171            .saturating_sub(self.conn.datagrams.outgoing_total)
172    }
173}
174
175#[derive(Default)]
176pub(super) struct DatagramState {
177    /// Number of bytes of datagrams that have been received by the local transport but not
178    /// delivered to the application
179    pub(super) recv_buffered: usize,
180    pub(super) incoming: VecDeque<Datagram>,
181    pub(super) outgoing: VecDeque<Datagram>,
182    pub(super) outgoing_total: usize,
183    pub(super) send_blocked: bool,
184}
185
186impl DatagramState {
187    pub(super) fn received(
188        &mut self,
189        datagram: Datagram,
190        window: &Option<usize>,
191    ) -> Result<bool, TransportError> {
192        let window = match window {
193            None => {
194                return Err(TransportError::PROTOCOL_VIOLATION(
195                    "unexpected DATAGRAM frame",
196                ));
197            }
198            Some(x) => *x,
199        };
200
201        if datagram.data.len() > window {
202            return Err(TransportError::PROTOCOL_VIOLATION("oversized datagram"));
203        }
204
205        let was_empty = self.recv_buffered == 0;
206        while datagram.data.len() + self.recv_buffered > window {
207            debug!("dropping stale datagram");
208            self.recv();
209        }
210
211        self.recv_buffered += datagram.data.len();
212        self.incoming.push_back(datagram);
213        Ok(was_empty)
214    }
215
216    fn make_space_for(&mut self, datagram_len: usize, send_buffer_size: usize) {
217        while !self.has_send_buffer_space(datagram_len, send_buffer_size) {
218            let Some(prev) = self.outgoing.pop_front() else {
219                break;
220            };
221            trace!(len = prev.data.len(), "dropping outgoing datagram");
222            self.outgoing_total -= prev.data.len();
223        }
224    }
225
226    fn has_send_buffer_space(&self, datagram_len: usize, send_buffer_size: usize) -> bool {
227        let Some(total) = self.outgoing_total.checked_add(datagram_len) else {
228            return false;
229        };
230
231        total <= send_buffer_size
232    }
233
234    /// Discard outgoing datagrams with a payload larger than `max_payload` bytes
235    ///
236    /// Returns whether any datagrams were dropped.
237    ///
238    /// Used to ensure that reductions in MTU don't get us stuck in a state where we have a datagram
239    /// queued but can't send it.
240    pub(super) fn drop_oversized(&mut self, max_payload: usize) -> bool {
241        let mut dropped_any = false;
242        self.outgoing.retain(|datagram| {
243            let result = datagram.data.len() < max_payload;
244            if !result {
245                trace!(
246                    "dropping {} byte datagram violating {} byte limit",
247                    datagram.data.len(),
248                    max_payload
249                );
250                self.outgoing_total -= datagram.data.len();
251                dropped_any = true;
252            }
253            result
254        });
255        dropped_any
256    }
257
258    /// Attempt to write a datagram frame into `buf`, consuming it from `self.outgoing`
259    ///
260    /// Returns whether a frame was written. At most `max_size` bytes will be written, including
261    /// framing.
262    pub(super) fn write<'a, 'b>(
263        &mut self,
264        buf: &mut PacketBuilder<'a, 'b>,
265        stat: &mut FrameStats,
266    ) -> bool {
267        let Some(datagram) = self.outgoing.pop_front() else {
268            return false;
269        };
270
271        if buf.frame_space_remaining() < datagram.size(true) {
272            // Future work: we could be more clever about cramming small datagrams into
273            // mostly-full packets when a larger one is queued first
274            self.outgoing.push_front(datagram);
275            return false;
276        }
277
278        self.outgoing_total -= datagram.data.len();
279        buf.write_frame(datagram, stat);
280        true
281    }
282
283    pub(super) fn recv(&mut self) -> Option<Bytes> {
284        let x = self.incoming.pop_front()?.data;
285        self.recv_buffered -= x.len();
286        Some(x)
287    }
288
289    /// Drain up to `out.len()` buffered datagrams into `out`, in arrival order.
290    ///
291    /// Returns the number of datagrams written into `out` (which may be less than
292    /// `out.len()` if fewer are buffered). Remaining datagrams stay queued.
293    pub(super) fn recv_many(&mut self, out: &mut [Bytes]) -> usize {
294        let n = out.len().min(self.incoming.len());
295        let mut received_bytes = 0;
296        for (i, d) in self.incoming.drain(..n).enumerate() {
297            received_bytes += d.data.len();
298            out[i] = d.data;
299        }
300        self.recv_buffered -= received_bytes;
301        n
302    }
303}
304
305#[cfg(test)]
306mod tests {
307    use super::*;
308
309    #[test]
310    fn make_space_for_accounts_for_new_datagram() {
311        let mut state = DatagramState::default();
312        state.outgoing.push_back(Datagram {
313            data: Bytes::from_static(&[0; 7]),
314        });
315        state.outgoing.push_back(Datagram {
316            data: Bytes::from_static(&[0; 2]),
317        });
318        state.outgoing_total = 9;
319
320        state.make_space_for(4, 10);
321
322        assert_eq!(state.outgoing.len(), 1);
323        assert_eq!(state.outgoing[0].data.len(), 2);
324        assert_eq!(state.outgoing_total, 2);
325    }
326
327    #[test]
328    fn make_space_for_handles_overflowing_capacity_check() {
329        let mut state = DatagramState::default();
330        state.outgoing.push_back(Datagram {
331            data: Bytes::from_static(&[0]),
332        });
333        state.outgoing_total = usize::MAX - 1;
334
335        state.make_space_for(2, usize::MAX);
336
337        assert!(state.outgoing.is_empty());
338        assert_eq!(state.outgoing_total, usize::MAX - 2);
339    }
340}
341
342/// Errors that can arise when sending a datagram
343#[derive(Debug, Error, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
344pub enum SendDatagramError {
345    /// The peer does not support receiving datagram frames
346    #[error("datagrams not supported by peer")]
347    UnsupportedByPeer,
348    /// Datagram support is disabled locally
349    #[error("datagram support disabled")]
350    Disabled,
351    /// The datagram is larger than the connection can currently accommodate
352    ///
353    /// Indicates that the path MTU minus overhead or the limit advertised by the peer has been
354    /// exceeded.
355    #[error("datagram too large")]
356    TooLarge,
357    /// Send would block
358    #[error("datagram send blocked")]
359    Blocked(Bytes),
360}
noq_proto/connection/datagrams.rs

noq_proto/connection/
datagrams.rs
