quinn_proto/connection/streams/
send.rs

1use bytes::Bytes;
2use thiserror::Error;
3
4use crate::{VarInt, connection::send_buffer::SendBuffer, frame};
5
6#[derive(Debug)]
7pub(super) struct Send {
8    pub(super) max_data: u64,
9    pub(super) state: SendState,
10    pub(super) pending: SendBuffer,
11    pub(super) priority: i32,
12    /// Whether a frame containing a FIN bit must be transmitted, even if we don't have any new data
13    pub(super) fin_pending: bool,
14    /// Whether this stream is in the `connection_blocked` list of `Streams`
15    pub(super) connection_blocked: bool,
16    /// Value of `max_data` for which a `STREAM_DATA_BLOCKED` frame was most recently queued
17    ///
18    /// A new frame is only queued once the peer has raised the limit.
19    pub(super) data_blocked_limit: Option<u64>,
20    /// The reason the peer wants us to stop, if `STOP_SENDING` was received
21    pub(super) stop_reason: Option<VarInt>,
22}
23
24impl Send {
25    pub(super) fn new(max_data: VarInt) -> Box<Self> {
26        Box::new(Self {
27            max_data: max_data.into(),
28            state: SendState::Ready,
29            pending: SendBuffer::new(),
30            priority: 0,
31            fin_pending: false,
32            connection_blocked: false,
33            data_blocked_limit: None,
34            stop_reason: None,
35        })
36    }
37
38    /// Whether the stream has been reset
39    pub(super) fn is_reset(&self) -> bool {
40        matches!(self.state, SendState::ResetSent)
41    }
42
43    pub(super) fn finish(&mut self) -> Result<(), FinishError> {
44        if let Some(error_code) = self.stop_reason {
45            Err(FinishError::Stopped(error_code))
46        } else if self.state == SendState::Ready {
47            self.state = SendState::DataSent {
48                finish_acked: false,
49            };
50            self.fin_pending = true;
51            Ok(())
52        } else {
53            Err(FinishError::ClosedStream)
54        }
55    }
56
57    pub(super) fn write<S: BytesSource>(
58        &mut self,
59        source: &mut S,
60        limit: u64,
61    ) -> Result<Written, WriteError> {
62        if !self.is_writable() {
63            return Err(WriteError::ClosedStream);
64        }
65        if let Some(error_code) = self.stop_reason {
66            return Err(WriteError::Stopped(error_code));
67        }
68        let budget = self.max_data - self.pending.offset();
69        if budget == 0 {
70            return Err(WriteError::Blocked);
71        }
72        let mut limit = limit.min(budget) as usize;
73
74        let mut result = Written::default();
75        loop {
76            let (chunk, chunks_consumed) = source.pop_chunk(limit);
77            result.chunks += chunks_consumed;
78            result.bytes += chunk.len();
79
80            if chunk.is_empty() {
81                break;
82            }
83
84            limit -= chunk.len();
85            self.pending.write(chunk);
86        }
87
88        Ok(result)
89    }
90
91    /// Update stream state due to a reset sent by the local application
92    pub(super) fn reset(&mut self) {
93        use SendState::*;
94        if let DataSent { .. } | Ready = self.state {
95            self.state = ResetSent;
96            self.pending.discard();
97            self.fin_pending = false;
98        }
99    }
100
101    /// Handle STOP_SENDING
102    ///
103    /// Returns true if the stream was stopped due to this frame, and false
104    /// if it had been stopped before
105    pub(super) fn try_stop(&mut self, error_code: VarInt) -> bool {
106        if self.stop_reason.is_none() {
107            self.stop_reason = Some(error_code);
108            true
109        } else {
110            false
111        }
112    }
113
114    /// Returns whether the stream has been finished and all data has been acknowledged by the peer
115    pub(super) fn ack(&mut self, frame: frame::StreamMeta) -> bool {
116        self.pending.ack(frame.offsets);
117        match self.state {
118            SendState::DataSent {
119                ref mut finish_acked,
120            } => {
121                *finish_acked |= frame.fin;
122                *finish_acked && self.pending.is_fully_acked()
123            }
124            _ => false,
125        }
126    }
127
128    /// Handle increase to stream-level flow control limit
129    ///
130    /// Returns whether the stream was unblocked
131    pub(super) fn increase_max_data(&mut self, offset: u64) -> bool {
132        if offset <= self.max_data || self.state != SendState::Ready {
133            return false;
134        }
135        let was_blocked = self.pending.offset() == self.max_data;
136        self.max_data = offset;
137        was_blocked
138    }
139
140    pub(super) fn offset(&self) -> u64 {
141        self.pending.offset()
142    }
143
144    pub(super) fn is_pending(&self) -> bool {
145        self.pending.has_unsent_data() || self.fin_pending
146    }
147
148    pub(super) fn is_writable(&self) -> bool {
149        matches!(self.state, SendState::Ready)
150    }
151}
152
153/// A [`BytesSource`] implementation for `&'a mut [Bytes]`
154///
155/// The type allows to dequeue [`Bytes`] chunks from an array of chunks, up to
156/// a configured limit.
157pub(crate) struct BytesArray<'a> {
158    /// The wrapped slice of `Bytes`
159    chunks: &'a mut [Bytes],
160    /// The amount of chunks consumed from this source
161    consumed: usize,
162}
163
164impl<'a> BytesArray<'a> {
165    pub(crate) fn from_chunks(chunks: &'a mut [Bytes]) -> Self {
166        Self {
167            chunks,
168            consumed: 0,
169        }
170    }
171}
172
173impl BytesSource for BytesArray<'_> {
174    fn pop_chunk(&mut self, limit: usize) -> (Bytes, usize) {
175        // The loop exists to skip empty chunks while still marking them as
176        // consumed
177        let mut chunks_consumed = 0;
178
179        while self.consumed < self.chunks.len() {
180            let chunk = &mut self.chunks[self.consumed];
181
182            if chunk.len() <= limit {
183                let chunk = std::mem::take(chunk);
184                self.consumed += 1;
185                chunks_consumed += 1;
186                if chunk.is_empty() {
187                    continue;
188                }
189                return (chunk, chunks_consumed);
190            } else if limit > 0 {
191                let chunk = chunk.split_to(limit);
192                return (chunk, chunks_consumed);
193            } else {
194                break;
195            }
196        }
197
198        (Bytes::new(), chunks_consumed)
199    }
200}
201
202/// A [`BytesSource`] implementation for `&[u8]`
203///
204/// The type allows to dequeue a single [`Bytes`] chunk, which will be lazily
205/// created from a reference. This allows to defer the allocation until it is
206/// known how much data needs to be copied.
207pub(crate) struct ByteSlice<'a> {
208    /// The wrapped byte slice
209    data: &'a [u8],
210}
211
212impl<'a> ByteSlice<'a> {
213    pub(crate) fn from_slice(data: &'a [u8]) -> Self {
214        Self { data }
215    }
216}
217
218impl BytesSource for ByteSlice<'_> {
219    fn pop_chunk(&mut self, limit: usize) -> (Bytes, usize) {
220        let limit = limit.min(self.data.len());
221        if limit == 0 {
222            return (Bytes::new(), 0);
223        }
224
225        let chunk = Bytes::from(self.data[..limit].to_owned());
226        self.data = &self.data[chunk.len()..];
227
228        let chunks_consumed = usize::from(self.data.is_empty());
229        (chunk, chunks_consumed)
230    }
231}
232
233/// A source of one or more buffers which can be converted into `Bytes` buffers on demand
234///
235/// The purpose of this data type is to defer conversion as long as possible,
236/// so that no heap allocation is required in case no data is writable.
237pub(super) trait BytesSource {
238    /// Returns the next chunk from the source of owned chunks.
239    ///
240    /// This method will consume parts of the source.
241    /// Calling it will yield `Bytes` elements up to the configured `limit`.
242    ///
243    /// The method returns a tuple:
244    /// - The first item is the yielded `Bytes` element. The element will be
245    ///   empty if the limit is zero or no more data is available.
246    /// - The second item returns how many complete chunks inside the source had
247    ///   had been consumed. This can be less than 1, if a chunk inside the
248    ///   source had been truncated in order to adhere to the limit. It can also
249    ///   be more than 1, if zero-length chunks had been skipped.
250    fn pop_chunk(&mut self, limit: usize) -> (Bytes, usize);
251}
252
253/// Indicates how many bytes and chunks had been transferred in a write operation
254#[derive(Debug, Default, PartialEq, Eq, Clone, Copy)]
255pub struct Written {
256    /// The amount of bytes which had been written
257    pub bytes: usize,
258    /// The amount of full chunks which had been written
259    ///
260    /// If a chunk was only partially written, it will not be counted by this field.
261    pub chunks: usize,
262}
263
264/// Errors triggered while writing to a send stream
265#[derive(Debug, Error, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
266pub enum WriteError {
267    /// The peer is not able to accept additional data, or the connection is congested.
268    ///
269    /// If the peer issues additional flow control credit, a [`StreamEvent::Writable`] event will
270    /// be generated, indicating that retrying the write might succeed.
271    ///
272    /// [`StreamEvent::Writable`]: crate::StreamEvent::Writable
273    #[error("unable to accept further writes")]
274    Blocked,
275    /// The peer is no longer accepting data on this stream, and it has been implicitly reset. The
276    /// stream cannot be finished or further written to.
277    ///
278    /// Carries an application-defined error code.
279    ///
280    /// [`StreamEvent::Finished`]: crate::StreamEvent::Finished
281    #[error("stopped by peer: code {0}")]
282    Stopped(VarInt),
283    /// The stream has not been opened or has already been finished or reset
284    #[error("closed stream")]
285    ClosedStream,
286}
287
288#[derive(Debug, Copy, Clone, Eq, PartialEq)]
289pub(super) enum SendState {
290    /// Sending new data
291    Ready,
292    /// Stream was finished; now sending retransmits only
293    DataSent { finish_acked: bool },
294    /// Sent RESET
295    ResetSent,
296}
297
298/// Reasons why attempting to finish a stream might fail
299#[derive(Debug, Error, Clone, PartialEq, Eq)]
300pub enum FinishError {
301    /// The peer is no longer accepting data on this stream. No
302    /// [`StreamEvent::Finished`] event will be emitted for this stream.
303    ///
304    /// Carries an application-defined error code.
305    ///
306    /// [`StreamEvent::Finished`]: crate::StreamEvent::Finished
307    #[error("stopped by peer: code {0}")]
308    Stopped(VarInt),
309    /// The stream has not been opened or was already finished or reset
310    #[error("closed stream")]
311    ClosedStream,
312}
313
314#[cfg(test)]
315mod tests {
316    use super::*;
317
318    #[test]
319    fn bytes_array() {
320        let full = b"Hello World 123456789 ABCDEFGHJIJKLMNOPQRSTUVWXYZ".to_owned();
321        for limit in 0..full.len() {
322            let mut chunks = [
323                Bytes::from_static(b""),
324                Bytes::from_static(b"Hello "),
325                Bytes::from_static(b"Wo"),
326                Bytes::from_static(b""),
327                Bytes::from_static(b"r"),
328                Bytes::from_static(b"ld"),
329                Bytes::from_static(b""),
330                Bytes::from_static(b" 12345678"),
331                Bytes::from_static(b"9 ABCDE"),
332                Bytes::from_static(b"F"),
333                Bytes::from_static(b"GHJIJKLMNOPQRSTUVWXYZ"),
334            ];
335            let num_chunks = chunks.len();
336            let last_chunk_len = chunks[chunks.len() - 1].len();
337
338            let mut array = BytesArray::from_chunks(&mut chunks);
339
340            let mut buf = Vec::new();
341            let mut chunks_popped = 0;
342            let mut chunks_consumed = 0;
343            let mut remaining = limit;
344            loop {
345                let (chunk, consumed) = array.pop_chunk(remaining);
346                chunks_consumed += consumed;
347
348                if !chunk.is_empty() {
349                    buf.extend_from_slice(&chunk);
350                    remaining -= chunk.len();
351                    chunks_popped += 1;
352                } else {
353                    break;
354                }
355            }
356
357            assert_eq!(&buf[..], &full[..limit]);
358
359            if limit == full.len() {
360                // Full consumption of the last chunk
361                assert_eq!(chunks_consumed, num_chunks);
362                // Since there are empty chunks, we consume more than there are popped
363                assert_eq!(chunks_consumed, chunks_popped + 3);
364            } else if limit > full.len() - last_chunk_len {
365                // Partial consumption of the last chunk
366                assert_eq!(chunks_consumed, num_chunks - 1);
367                assert_eq!(chunks_consumed, chunks_popped + 2);
368            }
369        }
370    }
371
372    #[test]
373    fn byte_slice() {
374        let full = b"Hello World 123456789 ABCDEFGHJIJKLMNOPQRSTUVWXYZ".to_owned();
375        for limit in 0..full.len() {
376            let mut array = ByteSlice::from_slice(&full[..]);
377
378            let mut buf = Vec::new();
379            let mut chunks_popped = 0;
380            let mut chunks_consumed = 0;
381            let mut remaining = limit;
382            loop {
383                let (chunk, consumed) = array.pop_chunk(remaining);
384                chunks_consumed += consumed;
385
386                if !chunk.is_empty() {
387                    buf.extend_from_slice(&chunk);
388                    remaining -= chunk.len();
389                    chunks_popped += 1;
390                } else {
391                    break;
392                }
393            }
394
395            assert_eq!(&buf[..], &full[..limit]);
396            if limit != 0 {
397                assert_eq!(chunks_popped, 1);
398            } else {
399                assert_eq!(chunks_popped, 0);
400            }
401
402            if limit == full.len() {
403                assert_eq!(chunks_consumed, 1);
404            } else {
405                assert_eq!(chunks_consumed, 0);
406            }
407        }
408    }
409}
quinn_proto/connection/streams/send.rs

quinn_proto/connection/streams/
send.rs
