h2/codec/
framed_write.rs

1use crate::codec::UserError;
2use crate::codec::UserError::*;
3use crate::frame::{self, Frame, FrameSize};
4use crate::hpack;
5
6use bytes::{Buf, BufMut, BytesMut};
7use std::pin::Pin;
8use std::task::{Context, Poll};
9use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
10use tokio_util::io::poll_write_buf;
11
12use std::io::{self, Cursor};
13
14// A macro to get around a method needing to borrow &mut self
15macro_rules! limited_write_buf {
16    ($self:expr) => {{
17        let limit = $self.max_frame_size() + frame::HEADER_LEN;
18        $self.buf.get_mut().limit(limit)
19    }};
20}
21
22#[derive(Debug)]
23pub struct FramedWrite<T, B> {
24    /// Upstream `AsyncWrite`
25    inner: T,
26    final_flush_done: bool,
27
28    encoder: Encoder<B>,
29}
30
31#[derive(Debug)]
32struct Encoder<B> {
33    /// HPACK encoder
34    hpack: hpack::Encoder,
35
36    /// Write buffer
37    ///
38    /// TODO: Should this be a ring buffer?
39    buf: Cursor<BytesMut>,
40
41    /// Next frame to encode
42    next: Option<Next<B>>,
43
44    /// Last data frame
45    last_data_frame: Option<frame::Data<B>>,
46
47    /// Max frame size, this is specified by the peer
48    max_frame_size: FrameSize,
49
50    /// Chain payloads bigger than this.
51    chain_threshold: usize,
52
53    /// Min buffer required to attempt to write a frame
54    min_buffer_capacity: usize,
55}
56
57#[derive(Debug)]
58enum Next<B> {
59    Data(frame::Data<B>),
60    Continuation(frame::Continuation),
61}
62
63/// Initialize the connection with this amount of write buffer.
64///
65/// The minimum MAX_FRAME_SIZE is 16kb, so always be able to send a HEADERS
66/// frame that big.
67const DEFAULT_BUFFER_CAPACITY: usize = 16 * 1_024;
68
69/// Chain payloads bigger than this when vectored I/O is enabled. The remote
70/// will never advertise a max frame size less than this (well, the spec says
71/// the max frame size can't be less than 16kb, so not even close).
72const CHAIN_THRESHOLD: usize = 256;
73
74/// Chain payloads bigger than this when vectored I/O is **not** enabled.
75/// A larger value in this scenario will reduce the number of small and
76/// fragmented data being sent, and hereby improve the throughput.
77const CHAIN_THRESHOLD_WITHOUT_VECTORED_IO: usize = 1024;
78
79// TODO: Make generic
80impl<T, B> FramedWrite<T, B>
81where
82    T: AsyncWrite + Unpin,
83    B: Buf,
84{
85    pub fn new(inner: T) -> FramedWrite<T, B> {
86        let chain_threshold = if inner.is_write_vectored() {
87            CHAIN_THRESHOLD
88        } else {
89            CHAIN_THRESHOLD_WITHOUT_VECTORED_IO
90        };
91        FramedWrite {
92            inner,
93            final_flush_done: false,
94            encoder: Encoder {
95                hpack: hpack::Encoder::default(),
96                buf: Cursor::new(BytesMut::with_capacity(DEFAULT_BUFFER_CAPACITY)),
97                next: None,
98                last_data_frame: None,
99                max_frame_size: frame::DEFAULT_MAX_FRAME_SIZE,
100                chain_threshold,
101                min_buffer_capacity: chain_threshold + frame::HEADER_LEN,
102            },
103        }
104    }
105
106    /// Returns `Ready` when `send` is able to accept a frame
107    ///
108    /// Calling this function may result in the current contents of the buffer
109    /// to be flushed to `T`.
110    pub fn poll_ready(&mut self, cx: &mut Context) -> Poll<io::Result<()>> {
111        if !self.encoder.has_capacity() {
112            // Try flushing
113            ready!(self.flush(cx))?;
114
115            if !self.encoder.has_capacity() {
116                return Poll::Pending;
117            }
118        }
119
120        Poll::Ready(Ok(()))
121    }
122
123    /// Returns whether a frame can be buffered without first flushing the
124    /// underlying I/O object.
125    pub(crate) fn has_capacity(&self) -> bool {
126        self.encoder.has_capacity()
127    }
128
129    /// Buffer a frame.
130    ///
131    /// `poll_ready` must be called first to ensure that a frame may be
132    /// accepted.
133    pub fn buffer(&mut self, item: Frame<B>) -> Result<(), UserError> {
134        self.encoder.buffer(item)
135    }
136
137    /// Flush buffered data to the wire
138    pub fn flush(&mut self, cx: &mut Context) -> Poll<io::Result<()>> {
139        let span = tracing::trace_span!("FramedWrite::flush");
140        let _e = span.enter();
141
142        loop {
143            while !self.encoder.is_empty() {
144                let n = match self.encoder.next {
145                    Some(Next::Data(ref mut frame)) => {
146                        tracing::trace!(queued_data_frame = true);
147                        let mut buf = (&mut self.encoder.buf).chain(frame.payload_mut());
148                        ready!(poll_write_buf(Pin::new(&mut self.inner), cx, &mut buf))?
149                    }
150                    _ => {
151                        tracing::trace!(queued_data_frame = false);
152                        ready!(poll_write_buf(
153                            Pin::new(&mut self.inner),
154                            cx,
155                            &mut self.encoder.buf
156                        ))?
157                    }
158                };
159                if n == 0 {
160                    // No progress is possible; retrying would busy-loop.
161                    tracing::trace!("write returned zero, but non-zero bytes remaining");
162                    return Poll::Ready(Err(io::ErrorKind::WriteZero.into()));
163                }
164            }
165
166            match self.encoder.unset_frame() {
167                ControlFlow::Continue => (),
168                ControlFlow::Break => break,
169            }
170        }
171
172        tracing::trace!("flushing buffer");
173        // Flush the upstream
174        ready!(Pin::new(&mut self.inner).poll_flush(cx))?;
175
176        Poll::Ready(Ok(()))
177    }
178
179    /// Close the codec
180    pub fn shutdown(&mut self, cx: &mut Context) -> Poll<io::Result<()>> {
181        if !self.final_flush_done {
182            ready!(self.flush(cx))?;
183            self.final_flush_done = true;
184        }
185        Pin::new(&mut self.inner).poll_shutdown(cx)
186    }
187}
188
189#[must_use]
190enum ControlFlow {
191    Continue,
192    Break,
193}
194
195impl<B> Encoder<B>
196where
197    B: Buf,
198{
199    fn unset_frame(&mut self) -> ControlFlow {
200        // Clear internal buffer
201        self.buf.set_position(0);
202        self.buf.get_mut().clear();
203
204        // The data frame has been written, so unset it
205        match self.next.take() {
206            Some(Next::Data(frame)) => {
207                self.last_data_frame = Some(frame);
208                debug_assert!(self.is_empty());
209                ControlFlow::Break
210            }
211            Some(Next::Continuation(frame)) => {
212                // Buffer the continuation frame, then try to write again
213                let mut buf = limited_write_buf!(self);
214                if let Some(continuation) = frame.encode(&mut buf) {
215                    self.next = Some(Next::Continuation(continuation));
216                }
217                ControlFlow::Continue
218            }
219            None => ControlFlow::Break,
220        }
221    }
222
223    fn buffer(&mut self, item: Frame<B>) -> Result<(), UserError> {
224        // Ensure that we have enough capacity to accept the write.
225        assert!(self.has_capacity());
226        let span = tracing::trace_span!("FramedWrite::buffer", frame = ?item);
227        let _e = span.enter();
228
229        tracing::debug!(frame = ?item, "send");
230
231        match item {
232            Frame::Data(mut v) => {
233                // Ensure that the payload is not greater than the max frame.
234                let len = v.payload().remaining();
235
236                if len > self.max_frame_size() {
237                    return Err(PayloadTooBig);
238                }
239
240                if len >= self.chain_threshold {
241                    let head = v.head();
242
243                    // Encode the frame head to the buffer
244                    head.encode(len, self.buf.get_mut());
245
246                    if self.buf.get_ref().remaining() < self.chain_threshold {
247                        let extra_bytes = self.chain_threshold - self.buf.remaining();
248                        self.buf.get_mut().put(v.payload_mut().take(extra_bytes));
249                    }
250
251                    // Save the data frame
252                    self.next = Some(Next::Data(v));
253                } else {
254                    v.encode_chunk(self.buf.get_mut());
255
256                    // The chunk has been fully encoded, so there is no need to
257                    // keep it around
258                    assert_eq!(v.payload().remaining(), 0, "chunk not fully encoded");
259
260                    // Save off the last frame...
261                    self.last_data_frame = Some(v);
262                }
263            }
264            Frame::Headers(v) => {
265                let mut buf = limited_write_buf!(self);
266                if let Some(continuation) = v.encode(&mut self.hpack, &mut buf) {
267                    self.next = Some(Next::Continuation(continuation));
268                }
269            }
270            Frame::PushPromise(v) => {
271                let mut buf = limited_write_buf!(self);
272                if let Some(continuation) = v.encode(&mut self.hpack, &mut buf) {
273                    self.next = Some(Next::Continuation(continuation));
274                }
275            }
276            Frame::Settings(v) => {
277                v.encode(self.buf.get_mut());
278                tracing::trace!(rem = self.buf.remaining(), "encoded settings");
279            }
280            Frame::GoAway(v) => {
281                v.encode(self.buf.get_mut());
282                tracing::trace!(rem = self.buf.remaining(), "encoded go_away");
283            }
284            Frame::Ping(v) => {
285                v.encode(self.buf.get_mut());
286                tracing::trace!(rem = self.buf.remaining(), "encoded ping");
287            }
288            Frame::WindowUpdate(v) => {
289                v.encode(self.buf.get_mut());
290                tracing::trace!(rem = self.buf.remaining(), "encoded window_update");
291            }
292
293            Frame::Priority(_) => {
294                /*
295                v.encode(self.buf.get_mut());
296                tracing::trace!("encoded priority; rem={:?}", self.buf.remaining());
297                */
298                unimplemented!();
299            }
300            Frame::Reset(v) => {
301                v.encode(self.buf.get_mut());
302                tracing::trace!(rem = self.buf.remaining(), "encoded reset");
303            }
304        }
305
306        Ok(())
307    }
308
309    fn has_capacity(&self) -> bool {
310        self.next.is_none()
311            && (self.buf.get_ref().capacity() - self.buf.get_ref().len()
312                >= self.min_buffer_capacity)
313    }
314
315    fn is_empty(&self) -> bool {
316        match self.next {
317            Some(Next::Data(ref frame)) => !frame.payload().has_remaining(),
318            _ => !self.buf.has_remaining(),
319        }
320    }
321}
322
323impl<B> Encoder<B> {
324    fn max_frame_size(&self) -> usize {
325        self.max_frame_size as usize
326    }
327}
328
329impl<T, B> FramedWrite<T, B> {
330    /// Returns the max frame size that can be sent
331    pub fn max_frame_size(&self) -> usize {
332        self.encoder.max_frame_size()
333    }
334
335    /// Set the peer's max frame size.
336    pub fn set_max_frame_size(&mut self, val: usize) {
337        assert!(val <= frame::MAX_MAX_FRAME_SIZE as usize);
338        self.encoder.max_frame_size = val as FrameSize;
339    }
340
341    /// Set the peer's header table size.
342    pub fn set_header_table_size(&mut self, val: usize) {
343        self.encoder.hpack.update_max_size(val);
344    }
345
346    /// Retrieve the last data frame that has been sent
347    pub fn take_last_data_frame(&mut self) -> Option<frame::Data<B>> {
348        self.encoder.last_data_frame.take()
349    }
350
351    pub fn get_mut(&mut self) -> &mut T {
352        &mut self.inner
353    }
354}
355
356impl<T: AsyncRead + Unpin, B> AsyncRead for FramedWrite<T, B> {
357    fn poll_read(
358        mut self: Pin<&mut Self>,
359        cx: &mut Context<'_>,
360        buf: &mut ReadBuf,
361    ) -> Poll<io::Result<()>> {
362        Pin::new(&mut self.inner).poll_read(cx, buf)
363    }
364}
365
366// We never project the Pin to `B`.
367impl<T: Unpin, B> Unpin for FramedWrite<T, B> {}
368
369#[cfg(feature = "unstable")]
370mod unstable {
371    use super::*;
372
373    impl<T, B> FramedWrite<T, B> {
374        pub fn get_ref(&self) -> &T {
375            &self.inner
376        }
377    }
378}
h2/codec/framed_write.rs

h2/codec/
framed_write.rs
