tokio_util/codec/
framed_impl.rs

1use crate::codec::decoder::Decoder;
2use crate::codec::encoder::Encoder;
3
4use futures_core::Stream;
5use tokio::io::{AsyncRead, AsyncWrite};
6
7use bytes::BytesMut;
8use futures_sink::Sink;
9use pin_project_lite::pin_project;
10use std::borrow::{Borrow, BorrowMut};
11use std::io;
12use std::pin::Pin;
13use std::task::{ready, Context, Poll};
14
15pin_project! {
16    #[derive(Debug)]
17    pub(crate) struct FramedImpl<T, U, State> {
18        #[pin]
19        pub(crate) inner: T,
20        pub(crate) state: State,
21        pub(crate) codec: U,
22    }
23}
24
25const INITIAL_CAPACITY: usize = 8 * 1024;
26
27#[derive(Debug)]
28pub(crate) struct ReadFrame {
29    pub(crate) eof: bool,
30    pub(crate) is_readable: bool,
31    pub(crate) buffer: BytesMut,
32    pub(crate) has_errored: bool,
33}
34
35pub(crate) struct WriteFrame {
36    pub(crate) buffer: BytesMut,
37    pub(crate) backpressure_boundary: usize,
38}
39
40#[derive(Default)]
41pub(crate) struct RWFrames {
42    pub(crate) read: ReadFrame,
43    pub(crate) write: WriteFrame,
44}
45
46impl Default for ReadFrame {
47    fn default() -> Self {
48        Self {
49            eof: false,
50            is_readable: false,
51            buffer: BytesMut::with_capacity(INITIAL_CAPACITY),
52            has_errored: false,
53        }
54    }
55}
56
57impl Default for WriteFrame {
58    fn default() -> Self {
59        Self {
60            buffer: BytesMut::with_capacity(INITIAL_CAPACITY),
61            backpressure_boundary: INITIAL_CAPACITY,
62        }
63    }
64}
65
66impl From<BytesMut> for ReadFrame {
67    fn from(mut buffer: BytesMut) -> Self {
68        let size = buffer.capacity();
69        if size < INITIAL_CAPACITY {
70            buffer.reserve(INITIAL_CAPACITY - size);
71        }
72
73        Self {
74            buffer,
75            is_readable: size > 0,
76            eof: false,
77            has_errored: false,
78        }
79    }
80}
81
82impl From<BytesMut> for WriteFrame {
83    fn from(mut buffer: BytesMut) -> Self {
84        let size = buffer.capacity();
85        if size < INITIAL_CAPACITY {
86            buffer.reserve(INITIAL_CAPACITY - size);
87        }
88
89        Self {
90            buffer,
91            backpressure_boundary: INITIAL_CAPACITY,
92        }
93    }
94}
95
96impl Borrow<ReadFrame> for RWFrames {
97    fn borrow(&self) -> &ReadFrame {
98        &self.read
99    }
100}
101impl BorrowMut<ReadFrame> for RWFrames {
102    fn borrow_mut(&mut self) -> &mut ReadFrame {
103        &mut self.read
104    }
105}
106impl Borrow<WriteFrame> for RWFrames {
107    fn borrow(&self) -> &WriteFrame {
108        &self.write
109    }
110}
111impl BorrowMut<WriteFrame> for RWFrames {
112    fn borrow_mut(&mut self) -> &mut WriteFrame {
113        &mut self.write
114    }
115}
116impl<T, U, R> Stream for FramedImpl<T, U, R>
117where
118    T: AsyncRead,
119    U: Decoder,
120    R: BorrowMut<ReadFrame>,
121{
122    type Item = Result<U::Item, U::Error>;
123
124    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
125        use crate::util::poll_read_buf;
126
127        let mut pinned = self.project();
128        let state: &mut ReadFrame = pinned.state.borrow_mut();
129        // The following loops implements a state machine with each state corresponding
130        // to a combination of the `is_readable` and `eof` flags. States persist across
131        // loop entries and most state transitions occur with a return.
132        //
133        // The initial state is `reading`.
134        //
135        // | state   | eof   | is_readable | has_errored |
136        // |---------|-------|-------------|-------------|
137        // | reading | false | false       | false       |
138        // | framing | false | true        | false       |
139        // | pausing | true  | true        | false       |
140        // | paused  | true  | false       | false       |
141        // | errored | <any> | <any>       | true        |
142        //                                                       `decode_eof` returns Err
143        //                                          ┌────────────────────────────────────────────────────────┐
144        //                   `decode_eof` returns   │                                                        │
145        //                             `Ok(Some)`   │                                                        │
146        //                                 ┌─────┐  │     `decode_eof` returns               After returning │
147        //                Read 0 bytes     ├─────▼──┴┐    `Ok(None)`          ┌────────┐ ◄───┐ `None`    ┌───▼─────┐
148        //               ┌────────────────►│ Pausing ├───────────────────────►│ Paused ├─┐   └───────────┤ Errored │
149        //               │                 └─────────┘                        └─┬──▲───┘ │               └───▲───▲─┘
150        // Pending read  │                                                      │  │     │                   │   │
151        //     ┌──────┐  │            `decode` returns `Some`                   │  └─────┘                   │   │
152        //     │      │  │                   ┌──────┐                           │  Pending                   │   │
153        //     │ ┌────▼──┴─┐ Read n>0 bytes ┌┴──────▼─┐     read n>0 bytes      │  read                      │   │
154        //     └─┤ Reading ├───────────────►│ Framing │◄────────────────────────┘                            │   │
155        //       └──┬─▲────┘                └─────┬──┬┘                                                      │   │
156        //          │ │                           │  │                 `decode` returns Err                  │   │
157        //          │ └───decode` returns `None`──┘  └───────────────────────────────────────────────────────┘   │
158        //          │                             read returns Err                                               │
159        //          └────────────────────────────────────────────────────────────────────────────────────────────┘
160        loop {
161            // Return `None` if we have encountered an error from the underlying decoder
162            // See: https://github.com/tokio-rs/tokio/issues/3976
163            if state.has_errored {
164                // preparing has_errored -> paused
165                trace!("Returning None and setting paused");
166                state.is_readable = false;
167                state.has_errored = false;
168                return Poll::Ready(None);
169            }
170
171            // Repeatedly call `decode` or `decode_eof` while the buffer is "readable",
172            // i.e. it _might_ contain data consumable as a frame or closing frame.
173            // Both signal that there is no such data by returning `None`.
174            //
175            // If `decode` couldn't read a frame and the upstream source has returned eof,
176            // `decode_eof` will attempt to decode the remaining bytes as closing frames.
177            //
178            // If the underlying AsyncRead is resumable, we may continue after an EOF,
179            // but must finish emitting all of it's associated `decode_eof` frames.
180            // Furthermore, we don't want to emit any `decode_eof` frames on retried
181            // reads after an EOF unless we've actually read more data.
182            if state.is_readable {
183                // pausing or framing
184                if state.eof {
185                    // pausing
186                    let frame = pinned.codec.decode_eof(&mut state.buffer).map_err(|err| {
187                        trace!("Got an error, going to errored state");
188                        state.has_errored = true;
189                        err
190                    })?;
191                    if frame.is_none() {
192                        state.is_readable = false; // prepare pausing -> paused
193                    }
194                    // implicit pausing -> pausing or pausing -> paused
195                    return Poll::Ready(frame.map(Ok));
196                }
197
198                // framing
199                trace!("attempting to decode a frame");
200
201                if let Some(frame) = pinned.codec.decode(&mut state.buffer).map_err(|op| {
202                    trace!("Got an error, going to errored state");
203                    state.has_errored = true;
204                    op
205                })? {
206                    trace!("frame decoded from buffer");
207                    // implicit framing -> framing
208                    return Poll::Ready(Some(Ok(frame)));
209                }
210
211                // framing -> reading
212                state.is_readable = false;
213            }
214            // reading or paused
215            // If we can't build a frame yet, try to read more data and try again.
216            // Make sure we've got room for at least one byte to read to ensure
217            // that we don't get a spurious 0 that looks like EOF.
218            state.buffer.reserve(1);
219            #[allow(clippy::blocks_in_conditions)]
220            let bytect = match poll_read_buf(pinned.inner.as_mut(), cx, &mut state.buffer).map_err(
221                |err| {
222                    trace!("Got an error, going to errored state");
223                    state.has_errored = true;
224                    err
225                },
226            )? {
227                Poll::Ready(ct) => ct,
228                // implicit reading -> reading or implicit paused -> paused
229                Poll::Pending => return Poll::Pending,
230            };
231            if bytect == 0 {
232                if state.eof {
233                    // We're already at an EOF, and since we've reached this path
234                    // we're also not readable. This implies that we've already finished
235                    // our `decode_eof` handling, so we can simply return `None`.
236                    // implicit paused -> paused
237                    return Poll::Ready(None);
238                }
239                // prepare reading -> paused
240                state.eof = true;
241            } else {
242                // prepare paused -> framing or noop reading -> framing
243                state.eof = false;
244            }
245
246            // paused -> framing or reading -> framing or reading -> pausing
247            state.is_readable = true;
248        }
249    }
250}
251
252impl<T, I, U, W> Sink<I> for FramedImpl<T, U, W>
253where
254    T: AsyncWrite,
255    U: Encoder<I>,
256    U::Error: From<io::Error>,
257    W: BorrowMut<WriteFrame>,
258{
259    type Error = U::Error;
260
261    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
262        if self.state.borrow().buffer.len() >= self.state.borrow().backpressure_boundary {
263            self.as_mut().poll_flush(cx)
264        } else {
265            Poll::Ready(Ok(()))
266        }
267    }
268
269    fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> {
270        let pinned = self.project();
271        pinned
272            .codec
273            .encode(item, &mut pinned.state.borrow_mut().buffer)?;
274        Ok(())
275    }
276
277    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
278        use crate::util::poll_write_buf;
279        trace!("flushing framed transport");
280        let mut pinned = self.project();
281
282        while !pinned.state.borrow_mut().buffer.is_empty() {
283            let WriteFrame { buffer, .. } = pinned.state.borrow_mut();
284            trace!(remaining = buffer.len(), "writing;");
285
286            let n = ready!(poll_write_buf(pinned.inner.as_mut(), cx, buffer))?;
287
288            if n == 0 {
289                return Poll::Ready(Err(io::Error::new(
290                    io::ErrorKind::WriteZero,
291                    "failed to \
292                     write frame to transport",
293                )
294                .into()));
295            }
296        }
297
298        // Try flushing the underlying IO
299        ready!(pinned.inner.poll_flush(cx))?;
300
301        trace!("framed transport flushed");
302        Poll::Ready(Ok(()))
303    }
304
305    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
306        ready!(self.as_mut().poll_flush(cx))?;
307        ready!(self.project().inner.poll_shutdown(cx))?;
308
309        Poll::Ready(Ok(()))
310    }
311}
tokio_util/codec/framed_impl.rs

tokio_util/codec/
framed_impl.rs
