actix_web_lab/
util.rs

1//! Utilities for working with Actix Web types.
2
3// stuff in here comes in and out of usage
4#![allow(dead_code)]
5
6use std::{
7    convert::Infallible,
8    io,
9    pin::Pin,
10    task::{Context, Poll, ready},
11};
12
13use actix_http::{BoxedPayloadStream, error::PayloadError};
14use actix_web::{dev, web::BufMut};
15use futures_core::Stream;
16use futures_util::StreamExt as _;
17use local_channel::mpsc;
18
19/// Returns an effectively cloned payload that supports streaming efficiently.
20///
21/// The cloned payload:
22/// - yields identical chunks;
23/// - does not poll ahead of the original;
24/// - does not poll significantly slower than the original;
25/// - receives an error signal if the original errors, but details are opaque to the copy.
26///
27/// If the payload is forked in one of the extractors used in a handler, then the original _must_ be
28/// read in another extractor or else the request will hang.
29pub fn fork_request_payload(orig_payload: &mut dev::Payload) -> dev::Payload {
30    const TARGET: &str = concat!(module_path!(), "::fork_request_payload");
31
32    let payload = orig_payload.take();
33
34    let (tx, rx) = mpsc::channel();
35
36    let proxy_stream: BoxedPayloadStream = Box::pin(payload.inspect(move |res| {
37        match res {
38            Ok(chunk) => {
39                tracing::trace!(target: TARGET, "yielding {} byte chunk", chunk.len());
40                tx.send(Ok(chunk.clone())).unwrap();
41            }
42
43            Err(err) => tx
44                .send(Err(PayloadError::Io(io::Error::other(format!(
45                    "error from original stream: {err}"
46                )))))
47                .unwrap(),
48        }
49    }));
50
51    tracing::trace!(target: TARGET, "creating proxy payload");
52    *orig_payload = dev::Payload::from(proxy_stream);
53
54    dev::Payload::Stream {
55        payload: Box::pin(rx),
56    }
57}
58
59/// An `io::Write`r that only requires mutable reference and assumes that there is space available
60/// in the buffer for every write operation or that it can be extended implicitly (like
61/// `bytes::BytesMut`, for example).
62///
63/// This is slightly faster (~10%) than `bytes::buf::Writer` in such cases because it does not
64/// perform a remaining length check before writing.
65pub(crate) struct MutWriter<'a, B>(pub(crate) &'a mut B);
66
67impl<B> MutWriter<'_, B> {
68    pub fn get_ref(&self) -> &B {
69        self.0
70    }
71}
72
73impl<B: BufMut> io::Write for MutWriter<'_, B> {
74    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
75        self.0.put_slice(buf);
76        Ok(buf.len())
77    }
78
79    fn flush(&mut self) -> io::Result<()> {
80        Ok(())
81    }
82}
83
84pin_project_lite::pin_project! {
85    /// Converts stream with item `T` into `Result<T, Infallible>`.
86    pub struct InfallibleStream<S> {
87        #[pin]
88        stream: S,
89    }
90}
91
92impl<S> InfallibleStream<S> {
93    /// Constructs new `InfallibleStream` stream.
94    pub fn new(stream: S) -> Self {
95        Self { stream }
96    }
97}
98
99impl<S: Stream> Stream for InfallibleStream<S> {
100    type Item = Result<S::Item, Infallible>;
101
102    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
103        Poll::Ready(ready!(self.project().stream.poll_next(cx)).map(Ok))
104    }
105
106    fn size_hint(&self) -> (usize, Option<usize>) {
107        self.stream.size_hint()
108    }
109}
110
111#[cfg(test)]
112#[derive(Debug, Clone, Default)]
113pub(crate) struct PollSeq<T> {
114    seq: std::collections::VecDeque<T>,
115}
116
117#[cfg(test)]
118mod poll_seq_impls {
119    use std::collections::VecDeque;
120
121    use futures_util::stream;
122
123    use super::*;
124
125    impl<T> PollSeq<T> {
126        pub fn new(seq: VecDeque<T>) -> Self {
127            Self { seq }
128        }
129    }
130
131    impl<T> PollSeq<Poll<Option<T>>> {
132        pub fn into_stream(mut self) -> impl Stream<Item = T> {
133            stream::poll_fn(move |_cx| match self.seq.pop_front() {
134                Some(item) => item,
135                None => Poll::Ready(None),
136            })
137        }
138    }
139
140    impl<T, const N: usize> From<[T; N]> for PollSeq<T> {
141        fn from(seq: [T; N]) -> Self {
142            Self::new(VecDeque::from(seq))
143        }
144    }
145}
actix_web_lab/util.rs

actix_web_lab/
util.rs
