1 use futures_core::TryFuture;
2 use futures_util::{future, TryFutureExt};
3 use std::fmt;
4 use std::future::Future;
5 use std::pin::Pin;
6 use std::task::{Context, Poll};
7 use tower_layer::Layer;
8 use tower_service::Service;
9 
10 /// Service returned by the [`and_then`] combinator.
11 ///
12 /// [`and_then`]: crate::util::ServiceExt::and_then
13 #[derive(Clone)]
14 pub struct AndThen<S, F> {
15     inner: S,
16     f: F,
17 }
18 
19 impl<S, F> fmt::Debug for AndThen<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("AndThen")
25             .field("inner", &self.inner)
26             .field("f", &format_args!("{}", std::any::type_name::<F>()))
27             .finish()
28     }
29 }
30 
31 pin_project_lite::pin_project! {
32     /// Response future from [`AndThen`] services.
33     ///
34     /// [`AndThen`]: crate::util::AndThen
35     pub struct AndThenFuture<F1, F2: TryFuture, N> {
36         #[pin]
37         inner: future::AndThen<future::ErrInto<F1, F2::Error>, F2, N>,
38     }
39 }
40 
41 impl<F1, F2: TryFuture, N> AndThenFuture<F1, F2, N> {
new(inner: future::AndThen<future::ErrInto<F1, F2::Error>, F2, N>) -> Self42     pub(crate) fn new(inner: future::AndThen<future::ErrInto<F1, F2::Error>, F2, N>) -> Self {
43         Self { inner }
44     }
45 }
46 
47 impl<F1, F2: TryFuture, N> std::fmt::Debug for AndThenFuture<F1, F2, N> {
fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result48     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
49         f.debug_tuple("AndThenFuture")
50             .field(&format_args!("..."))
51             .finish()
52     }
53 }
54 
55 impl<F1, F2: TryFuture, N> Future for AndThenFuture<F1, F2, N>
56 where
57     future::AndThen<future::ErrInto<F1, F2::Error>, F2, N>: Future,
58 {
59     type Output = <future::AndThen<future::ErrInto<F1, F2::Error>, F2, N> as Future>::Output;
60 
61     #[inline]
poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>62     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
63         self.project().inner.poll(cx)
64     }
65 }
66 
67 /// A [`Layer`] that produces a [`AndThen`] service.
68 ///
69 /// [`Layer`]: tower_layer::Layer
70 #[derive(Clone, Debug)]
71 pub struct AndThenLayer<F> {
72     f: F,
73 }
74 
75 impl<S, F> AndThen<S, F> {
76     /// Creates a new `AndThen` service.
new(inner: S, f: F) -> Self77     pub fn new(inner: S, f: F) -> Self {
78         AndThen { f, inner }
79     }
80 
81     /// Returns a new [`Layer`] that produces [`AndThen`] services.
82     ///
83     /// This is a convenience function that simply calls [`AndThenLayer::new`].
84     ///
85     /// [`Layer`]: tower_layer::Layer
layer(f: F) -> AndThenLayer<F>86     pub fn layer(f: F) -> AndThenLayer<F> {
87         AndThenLayer { f }
88     }
89 }
90 
91 impl<S, F, Request, Fut> Service<Request> for AndThen<S, F>
92 where
93     S: Service<Request>,
94     S::Error: Into<Fut::Error>,
95     F: FnOnce(S::Response) -> Fut + Clone,
96     Fut: TryFuture,
97 {
98     type Response = Fut::Ok;
99     type Error = Fut::Error;
100     type Future = AndThenFuture<S::Future, Fut, F>;
101 
poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>>102     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
103         self.inner.poll_ready(cx).map_err(Into::into)
104     }
105 
call(&mut self, request: Request) -> Self::Future106     fn call(&mut self, request: Request) -> Self::Future {
107         AndThenFuture::new(self.inner.call(request).err_into().and_then(self.f.clone()))
108     }
109 }
110 
111 impl<F> AndThenLayer<F> {
112     /// Creates a new [`AndThenLayer`] layer.
new(f: F) -> Self113     pub fn new(f: F) -> Self {
114         AndThenLayer { f }
115     }
116 }
117 
118 impl<S, F> Layer<S> for AndThenLayer<F>
119 where
120     F: Clone,
121 {
122     type Service = AndThen<S, F>;
123 
layer(&self, inner: S) -> Self::Service124     fn layer(&self, inner: S) -> Self::Service {
125         AndThen {
126             f: self.f.clone(),
127             inner,
128         }
129     }
130 }
131