diff --git a/src/middleware/err_handlers.rs b/src/middleware/err_handlers.rs index dfd9a7dc..44962aa9 100644 --- a/src/middleware/err_handlers.rs +++ b/src/middleware/err_handlers.rs @@ -1,10 +1,15 @@ //! For middleware documentation, see [`ErrorHandlers`]. -use std::rc::Rc; +use std::{ + future::Future, + pin::Pin, + rc::Rc, + task::{Context, Poll}, +}; use actix_service::{Service, Transform}; use ahash::AHashMap; -use futures_util::future::{ok, FutureExt, LocalBoxFuture, Ready}; +use futures_core::{future::LocalBoxFuture, ready}; use crate::{ dev::{ServiceRequest, ServiceResponse}, @@ -51,9 +56,11 @@ type ErrorHandler = dyn Fn(ServiceResponse) -> Result { - handlers: Rc>>>, + handlers: Handlers, } +type Handlers = Rc>>>; + impl Default for ErrorHandlers { fn default() -> Self { ErrorHandlers { @@ -82,7 +89,7 @@ impl ErrorHandlers { impl Transform for ErrorHandlers where - S: Service, Error = Error>, + S: Service, Error = Error> + 'static, S::Future: 'static, B: 'static, { @@ -90,20 +97,18 @@ where type Error = Error; type Transform = ErrorHandlersMiddleware; type InitError = (); - type Future = Ready>; + type Future = LocalBoxFuture<'static, Result>; fn new_transform(&self, service: S) -> Self::Future { - ok(ErrorHandlersMiddleware { - service, - handlers: self.handlers.clone(), - }) + let handlers = self.handlers.clone(); + Box::pin(async move { Ok(ErrorHandlersMiddleware { service, handlers }) }) } } #[doc(hidden)] pub struct ErrorHandlersMiddleware { service: S, - handlers: Rc>>>, + handlers: Handlers, } impl Service for ErrorHandlersMiddleware @@ -114,35 +119,63 @@ where { type Response = ServiceResponse; type Error = Error; - type Future = LocalBoxFuture<'static, Result>; + type Future = ErrorHandlersFuture; actix_service::forward_ready!(service); fn call(&mut self, req: ServiceRequest) -> Self::Future { let handlers = self.handlers.clone(); let fut = self.service.call(req); + ErrorHandlersFuture::ServiceFuture { fut, handlers } + } +} - async move { - let res = fut.await?; +#[pin_project::pin_project(project = ErrorHandlersProj)] +pub enum ErrorHandlersFuture +where + Fut: Future, +{ + ServiceFuture { + #[pin] + fut: Fut, + handlers: Handlers, + }, + HandlerFuture { + fut: LocalBoxFuture<'static, Fut::Output>, + }, +} - if let Some(handler) = handlers.get(&res.status()) { - match handler(res) { - Ok(ErrorHandlerResponse::Response(res)) => Ok(res), - Ok(ErrorHandlerResponse::Future(fut)) => fut.await, - Err(e) => Err(e), +impl Future for ErrorHandlersFuture +where + Fut: Future, Error>>, +{ + type Output = Fut::Output; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + match self.as_mut().project() { + ErrorHandlersProj::ServiceFuture { fut, handlers } => { + let res = ready!(fut.poll(cx))?; + match handlers.get(&res.status()) { + Some(handler) => match handler(res)? { + ErrorHandlerResponse::Response(res) => Poll::Ready(Ok(res)), + ErrorHandlerResponse::Future(fut) => { + self.as_mut() + .set(ErrorHandlersFuture::HandlerFuture { fut }); + self.poll(cx) + } + }, + None => Poll::Ready(Ok(res)), } - } else { - Ok(res) } + ErrorHandlersProj::HandlerFuture { fut } => fut.as_mut().poll(cx), } - .boxed_local() } } #[cfg(test)] mod tests { use actix_service::IntoService; - use futures_util::future::ok; + use futures_util::future::{ok, FutureExt}; use super::*; use crate::http::{header::CONTENT_TYPE, HeaderValue, StatusCode};