From 4edeb5ce47280d4daf1352faefc832420ecf409b Mon Sep 17 00:00:00 2001
From: fakeshadow <24548779@qq.com>
Date: Thu, 14 Jan 2021 09:43:44 +0800
Subject: [PATCH] optimize ErrorHandler middleware (#1902)

---
 src/middleware/err_handlers.rs | 77 ++++++++++++++++++++++++----------
 1 file changed, 55 insertions(+), 22 deletions(-)

diff --git a/src/middleware/err_handlers.rs b/src/middleware/err_handlers.rs
index dfd9a7dc5..44962aa98 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<B> = dyn Fn(ServiceResponse<B>) -> Result<ErrorHandlerResponse
 ///     ));
 /// ```
 pub struct ErrorHandlers<B> {
-    handlers: Rc<AHashMap<StatusCode, Box<ErrorHandler<B>>>>,
+    handlers: Handlers<B>,
 }
 
+type Handlers<B> = Rc<AHashMap<StatusCode, Box<ErrorHandler<B>>>>;
+
 impl<B> Default for ErrorHandlers<B> {
     fn default() -> Self {
         ErrorHandlers {
@@ -82,7 +89,7 @@ impl<B> ErrorHandlers<B> {
 
 impl<S, B> Transform<S, ServiceRequest> for ErrorHandlers<B>
 where
-    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
+    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
     S::Future: 'static,
     B: 'static,
 {
@@ -90,20 +97,18 @@ where
     type Error = Error;
     type Transform = ErrorHandlersMiddleware<S, B>;
     type InitError = ();
-    type Future = Ready<Result<Self::Transform, Self::InitError>>;
+    type Future = LocalBoxFuture<'static, Result<Self::Transform, Self::InitError>>;
 
     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<S, B> {
     service: S,
-    handlers: Rc<AHashMap<StatusCode, Box<ErrorHandler<B>>>>,
+    handlers: Handlers<B>,
 }
 
 impl<S, B> Service<ServiceRequest> for ErrorHandlersMiddleware<S, B>
@@ -114,35 +119,63 @@ where
 {
     type Response = ServiceResponse<B>;
     type Error = Error;
-    type Future = LocalBoxFuture<'static, Result<Self::Response, Self::Error>>;
+    type Future = ErrorHandlersFuture<S::Future, B>;
 
     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<Fut, B>
+where
+    Fut: Future,
+{
+    ServiceFuture {
+        #[pin]
+        fut: Fut,
+        handlers: Handlers<B>,
+    },
+    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<Fut, B> Future for ErrorHandlersFuture<Fut, B>
+where
+    Fut: Future<Output = Result<ServiceResponse<B>, Error>>,
+{
+    type Output = Fut::Output;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        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};