From ab4e889f9603525731ecd388acaa24efe1f94e04 Mon Sep 17 00:00:00 2001 From: Nikolay Kim Date: Sun, 29 Apr 2018 20:50:38 -0700 Subject: [PATCH] add middleware finished handler for route middleware --- src/resource.rs | 3 ++ src/route.rs | 101 +++++++++++++++++++++++++++++++++++++++++------- src/scope.rs | 2 +- 3 files changed, 90 insertions(+), 16 deletions(-) diff --git a/src/resource.rs b/src/resource.rs index a6e6b731e..e4dfbb2df 100644 --- a/src/resource.rs +++ b/src/resource.rs @@ -187,6 +187,9 @@ impl ResourceHandler { /// /// This is similar to `App's` middlewares, but /// middlewares get invoked on resource level. + /// + /// *Note* `Middleware::finish()` fires right after response get + /// prepared. It does not wait until body get sent to peer. pub fn middleware>(&mut self, mw: M) { Rc::get_mut(&mut self.middlewares) .unwrap() diff --git a/src/route.rs b/src/route.rs index 27f0133b0..6c4ba4197 100644 --- a/src/route.rs +++ b/src/route.rs @@ -10,8 +10,8 @@ use handler::{AsyncHandler, FromRequest, Handler, Reply, ReplyItem, Responder, use http::StatusCode; use httprequest::HttpRequest; use httpresponse::HttpResponse; -use middleware::{Middleware, Response as MiddlewareResponse, - Started as MiddlewareStarted}; +use middleware::{Finished as MiddlewareFinished, Middleware, + Response as MiddlewareResponse, Started as MiddlewareStarted}; use pred::Predicate; use with::{ExtractorConfig, With, With2, With3}; @@ -274,7 +274,8 @@ enum ComposeState { Starting(StartMiddlewares), Handler(WaitingResponse), RunMiddlewares(RunMiddlewares), - Response(Response), + Finishing(FinishingMiddlewares), + Completed(Response), } impl ComposeState { @@ -283,7 +284,8 @@ impl ComposeState { ComposeState::Starting(ref mut state) => state.poll(info), ComposeState::Handler(ref mut state) => state.poll(info), ComposeState::RunMiddlewares(ref mut state) => state.poll(info), - ComposeState::Response(_) => None, + ComposeState::Finishing(ref mut state) => state.poll(info), + ComposeState::Completed(_) => None, } } } @@ -310,7 +312,7 @@ impl Future for Compose { fn poll(&mut self) -> Poll { loop { - if let ComposeState::Response(ref mut resp) = self.state { + if let ComposeState::Completed(ref mut resp) = self.state { let resp = resp.resp.take().unwrap(); return Ok(Async::Ready(resp)); } @@ -357,9 +359,9 @@ impl StartMiddlewares { } info.count += 1; } - Err(err) => return Response::init(err.into()), + Err(err) => return FinishingMiddlewares::init(info, err.into()), }, - Err(err) => return Response::init(err.into()), + Err(err) => return FinishingMiddlewares::init(info, err.into()), } } } @@ -389,12 +391,17 @@ impl StartMiddlewares { self.fut = Some(fut); continue 'outer; } - Err(err) => return Some(Response::init(err.into())), + Err(err) => { + return Some(FinishingMiddlewares::init( + info, + err.into(), + )) + } } } } } - Err(err) => return Some(Response::init(err.into())), + Err(err) => return Some(FinishingMiddlewares::init(info, err.into())), } } } @@ -443,12 +450,12 @@ impl RunMiddlewares { resp = match info.mws[curr].response(&mut info.req, resp) { Err(err) => { info.count = curr + 1; - return Response::init(err.into()); + return FinishingMiddlewares::init(info, err.into()); } Ok(MiddlewareResponse::Done(r)) => { curr += 1; if curr == len { - return Response::init(r); + return FinishingMiddlewares::init(info, r); } else { r } @@ -475,15 +482,17 @@ impl RunMiddlewares { self.curr += 1; resp } - Err(err) => return Some(Response::init(err.into())), + Err(err) => return Some(FinishingMiddlewares::init(info, err.into())), }; loop { if self.curr == len { - return Some(Response::init(resp)); + return Some(FinishingMiddlewares::init(info, resp)); } else { match info.mws[self.curr].response(&mut info.req, resp) { - Err(err) => return Some(Response::init(err.into())), + Err(err) => { + return Some(FinishingMiddlewares::init(info, err.into())) + } Ok(MiddlewareResponse::Done(r)) => { self.curr += 1; resp = r @@ -499,6 +508,68 @@ impl RunMiddlewares { } } +/// Middlewares start executor +struct FinishingMiddlewares { + resp: Option, + fut: Option>>, + _s: PhantomData, +} + +impl FinishingMiddlewares { + fn init(info: &mut ComposeInfo, resp: HttpResponse) -> ComposeState { + if info.count == 0 { + Response::init(resp) + } else { + let mut state = FinishingMiddlewares { + resp: Some(resp), + fut: None, + _s: PhantomData, + }; + if let Some(st) = state.poll(info) { + st + } else { + ComposeState::Finishing(state) + } + } + } + + fn poll(&mut self, info: &mut ComposeInfo) -> Option> { + loop { + // poll latest fut + let not_ready = if let Some(ref mut fut) = self.fut { + match fut.poll() { + Ok(Async::NotReady) => true, + Ok(Async::Ready(())) => false, + Err(err) => { + error!("Middleware finish error: {}", err); + false + } + } + } else { + false + }; + if not_ready { + return None; + } + self.fut = None; + info.count -= 1; + + match info.mws[info.count as usize] + .finish(&mut info.req, self.resp.as_ref().unwrap()) + { + MiddlewareFinished::Done => { + if info.count == 0 { + return Some(Response::init(self.resp.take().unwrap())); + } + } + MiddlewareFinished::Future(fut) => { + self.fut = Some(fut); + } + } + } + } +} + struct Response { resp: Option, _s: PhantomData, @@ -506,7 +577,7 @@ struct Response { impl Response { fn init(resp: HttpResponse) -> ComposeState { - ComposeState::Response(Response { + ComposeState::Completed(Response { resp: Some(resp), _s: PhantomData, }) diff --git a/src/scope.rs b/src/scope.rs index b671eaaca..29d51518a 100644 --- a/src/scope.rs +++ b/src/scope.rs @@ -169,7 +169,7 @@ impl Scope { /// This is similar to `App's` middlewares, but /// middlewares get invoked on scope level. /// - /// *Note* `Middleware::finish()` is fired right after response get + /// *Note* `Middleware::finish()` fires right after response get /// prepared. It does not wait until body get sent to peer. pub fn middleware>(mut self, mw: M) -> Scope { Rc::get_mut(&mut self.middlewares)