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 | { |
23 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
24 | f&mut DebugStruct<'_, '_>.debug_struct("Then" ) |
25 | .field("inner" , &self.inner) |
26 | .field(name:"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. |
41 | pub const 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 |
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 ] |
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 ] |
79 | fn call(&mut self, request: Request) -> Self::Future { |
80 | ThenFuture::new(self.inner.call(req:request).then(self.f.clone())) |
81 | } |
82 | } |
83 | |
84 | impl<F> ThenLayer<F> { |
85 | /// Creates a new [`ThenLayer`] layer. |
86 | pub const 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 | |
97 | fn layer(&self, inner: S) -> Self::Service { |
98 | Then { |
99 | f: self.f.clone(), |
100 | inner, |
101 | } |
102 | } |
103 | } |
104 | |