1
0
mirror of https://github.com/fafhrd91/actix-web synced 2024-11-30 18:44:35 +01:00

Static either extract future (#2184)

This commit is contained in:
Ibraheem Ahmed 2021-04-30 22:02:56 -04:00 committed by GitHub
parent 1fcf92e11f
commit 3a0fb3f89e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -1,8 +1,14 @@
//! For either helper, see [`Either`]. //! For either helper, see [`Either`].
use std::{
future::Future,
mem,
pin::Pin,
task::{Context, Poll},
};
use bytes::Bytes; use bytes::Bytes;
use futures_core::future::LocalBoxFuture; use futures_core::ready;
use futures_util::{FutureExt as _, TryFutureExt as _};
use crate::{ use crate::{
dev, dev,
@ -180,41 +186,111 @@ where
R: FromRequest + 'static, R: FromRequest + 'static,
{ {
type Error = EitherExtractError<L::Error, R::Error>; type Error = EitherExtractError<L::Error, R::Error>;
type Future = LocalBoxFuture<'static, Result<Self, Self::Error>>; type Future = EitherExtractFut<L, R>;
type Config = (); type Config = ();
fn from_request(req: &HttpRequest, payload: &mut dev::Payload) -> Self::Future { fn from_request(req: &HttpRequest, payload: &mut dev::Payload) -> Self::Future {
let req2 = req.clone(); EitherExtractFut {
req: req.clone(),
Bytes::from_request(req, payload) state: EitherExtractState::Bytes {
.map_err(EitherExtractError::Bytes) bytes: Bytes::from_request(req, payload),
.and_then(|bytes| bytes_to_l_or_r(req2, bytes)) },
.boxed_local() }
} }
} }
async fn bytes_to_l_or_r<L, R>( #[pin_project::pin_project]
req: HttpRequest, pub struct EitherExtractFut<L, R>
bytes: Bytes,
) -> Result<Either<L, R>, EitherExtractError<L::Error, R::Error>>
where where
L: FromRequest + 'static, R: FromRequest,
R: FromRequest + 'static, L: FromRequest,
{ {
req: HttpRequest,
#[pin]
state: EitherExtractState<L, R>,
}
#[pin_project::pin_project(project = EitherExtractProj)]
pub enum EitherExtractState<L, R>
where
L: FromRequest,
R: FromRequest,
{
Bytes {
#[pin]
bytes: <Bytes as FromRequest>::Future,
},
Left {
#[pin]
left: L::Future,
fallback: Bytes,
},
Right {
#[pin]
right: R::Future,
left_err: Option<L::Error>,
},
}
impl<R, RF, RE, L, LF, LE> Future for EitherExtractFut<L, R>
where
L: FromRequest<Future = LF, Error = LE>,
R: FromRequest<Future = RF, Error = RE>,
LF: Future<Output = Result<L, LE>> + 'static,
RF: Future<Output = Result<R, RE>> + 'static,
LE: Into<Error>,
RE: Into<Error>,
{
type Output = Result<Either<L, R>, EitherExtractError<LE, RE>>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut this = self.project();
let ready = loop {
let next = match this.state.as_mut().project() {
EitherExtractProj::Bytes { bytes } => {
let res = ready!(bytes.poll(cx));
match res {
Ok(bytes) => {
let fallback = bytes.clone(); let fallback = bytes.clone();
let a_err; let left =
L::from_request(&this.req, &mut payload_from_bytes(bytes));
let mut pl = payload_from_bytes(bytes); EitherExtractState::Left { left, fallback }
match L::from_request(&req, &mut pl).await { }
Ok(a_data) => return Ok(Either::Left(a_data)), Err(err) => break Err(EitherExtractError::Bytes(err)),
// store A's error for returning if B also fails }
Err(err) => a_err = err, }
EitherExtractProj::Left { left, fallback } => {
let res = ready!(left.poll(cx));
match res {
Ok(extracted) => break Ok(Either::Left(extracted)),
Err(left_err) => {
let right = R::from_request(
&this.req,
&mut payload_from_bytes(mem::take(fallback)),
);
EitherExtractState::Right {
left_err: Some(left_err),
right,
}
}
}
}
EitherExtractProj::Right { right, left_err } => {
let res = ready!(right.poll(cx));
match res {
Ok(data) => break Ok(Either::Right(data)),
Err(err) => {
break Err(EitherExtractError::Extract(
left_err.take().unwrap(),
err,
));
}
}
}
}; };
this.state.set(next);
let mut pl = payload_from_bytes(fallback); };
match R::from_request(&req, &mut pl).await { Poll::Ready(ready)
Ok(b_data) => return Ok(Either::Right(b_data)),
Err(b_err) => Err(EitherExtractError::Extract(a_err, b_err)),
} }
} }