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