1 use futures_util::{future, FutureExt};
2 use std::{
3     fmt,
4     future::Future,
5     task::{Context, Poll},
6 };
7 use tower_layer::Layer;
8 use tower_service::Service;
9 
10 /// [`Service`] returned by the [`then`] combinator.
11 ///
12 /// [`then`]: crate::util::ServiceExt::then
13 #[derive(Clone)]
14 pub struct Then<S, F> {
15     inner: S,
16     f: F,
17 }
18 
19 impl<S, F> fmt::Debug for Then<S, F>
20 where
21     S: fmt::Debug,
22 {
fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result23     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
24         f.debug_struct("Then")
25             .field("inner", &self.inner)
26             .field("f", &format_args!("{}", std::any::type_name::<F>()))
27             .finish()
28     }
29 }
30 
31 /// A [`Layer`] that produces a [`Then`] service.
32 ///
33 /// [`Layer`]: tower_layer::Layer
34 #[derive(Debug, Clone)]
35 pub struct ThenLayer<F> {
36     f: F,
37 }
38 
39 impl<S, F> Then<S, F> {
40     /// Creates a new `Then` service.
new(inner: S, f: F) -> Self41     pub fn new(inner: S, f: F) -> Self {
42         Then { f, inner }
43     }
44 
45     /// Returns a new [`Layer`] that produces [`Then`] services.
46     ///
47     /// This is a convenience function that simply calls [`ThenLayer::new`].
48     ///
49     /// [`Layer`]: tower_layer::Layer
layer(f: F) -> ThenLayer<F>50     pub fn layer(f: F) -> ThenLayer<F> {
51         ThenLayer { f }
52     }
53 }
54 
55 opaque_future! {
56     /// Response future from [`Then`] services.
57     ///
58     /// [`Then`]: crate::util::Then
59     pub type ThenFuture<F1, F2, N> = future::Then<F1, F2, N>;
60 }
61 
62 impl<S, F, Request, Response, Error, Fut> Service<Request> for Then<S, F>
63 where
64     S: Service<Request>,
65     S::Error: Into<Error>,
66     F: FnOnce(Result<S::Response, S::Error>) -> Fut + Clone,
67     Fut: Future<Output = Result<Response, Error>>,
68 {
69     type Response = Response;
70     type Error = Error;
71     type Future = ThenFuture<S::Future, Fut, F>;
72 
73     #[inline]
poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>>74     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
75         self.inner.poll_ready(cx).map_err(Into::into)
76     }
77 
78     #[inline]
call(&mut self, request: Request) -> Self::Future79     fn call(&mut self, request: Request) -> Self::Future {
80         ThenFuture::new(self.inner.call(request).then(self.f.clone()))
81     }
82 }
83 
84 impl<F> ThenLayer<F> {
85     /// Creates a new [`ThenLayer`] layer.
new(f: F) -> Self86     pub fn new(f: F) -> Self {
87         ThenLayer { f }
88     }
89 }
90 
91 impl<S, F> Layer<S> for ThenLayer<F>
92 where
93     F: Clone,
94 {
95     type Service = Then<S, F>;
96 
layer(&self, inner: S) -> Self::Service97     fn layer(&self, inner: S) -> Self::Service {
98         Then {
99             f: self.f.clone(),
100             inner,
101         }
102     }
103 }
104