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

actix_web_lab/
util.rs
