Skip to main content

tower_fallback/
future.rs

1//! Future types for the `Fallback` middleware.
2
3use std::{
4    fmt::Debug,
5    future::Future,
6    pin::Pin,
7    task::{Context, Poll},
8};
9
10use futures_core::ready;
11use pin_project::pin_project;
12use tower::Service;
13
14use crate::{BoxedError, FallbackPolicy, OnError};
15
16/// Future that completes either with the first service's successful response, or
17/// with the second service's response.
18#[pin_project]
19pub struct ResponseFuture<S1, S2, Request, F = OnError>
20where
21    S1: Service<Request>,
22    S2: Service<Request, Response = <S1 as Service<Request>>::Response>,
23    F: FallbackPolicy<<S1 as Service<Request>>::Response>,
24    S1::Error: Into<BoxedError>,
25    S2::Error: Into<BoxedError>,
26{
27    #[pin]
28    state: ResponseState<S1, S2, Request>,
29    policy: F,
30}
31
32#[pin_project(project_replace = __ResponseStateProjectionOwned, project = ResponseStateProj)]
33enum ResponseState<S1, S2, Request>
34where
35    S1: Service<Request>,
36    S2: Service<Request>,
37    S2::Error: Into<BoxedError>,
38{
39    PollResponse1 {
40        #[pin]
41        fut: S1::Future,
42        req: Request,
43        svc2: S2,
44    },
45    PollReady2 {
46        req: Request,
47        svc2: S2,
48    },
49    PollResponse2 {
50        #[pin]
51        fut: S2::Future,
52    },
53    // Placeholder value to swap into the pin projection of the enum so we can take ownership of the fields.
54    Tmp,
55}
56
57impl<S1, S2, Request, F> ResponseFuture<S1, S2, Request, F>
58where
59    S1: Service<Request>,
60    S2: Service<Request, Response = <S1 as Service<Request>>::Response>,
61    F: FallbackPolicy<<S1 as Service<Request>>::Response>,
62    S1::Error: Into<BoxedError>,
63    S2::Error: Into<BoxedError>,
64{
65    pub(crate) fn new(fut: S1::Future, req: Request, svc2: S2, policy: F) -> Self {
66        ResponseFuture {
67            state: ResponseState::PollResponse1 { fut, req, svc2 },
68            policy,
69        }
70    }
71}
72
73impl<S1, S2, Request, F> Future for ResponseFuture<S1, S2, Request, F>
74where
75    S1: Service<Request>,
76    S2: Service<Request, Response = <S1 as Service<Request>>::Response>,
77    F: FallbackPolicy<<S1 as Service<Request>>::Response>,
78    S1::Error: Into<BoxedError>,
79    S2::Error: Into<BoxedError>,
80{
81    type Output = Result<<S1 as Service<Request>>::Response, BoxedError>;
82
83    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
84        let mut this = self.project();
85        // CORRECTNESS
86        //
87        // The current task must be scheduled for wakeup every time we return
88        // `Poll::Pending`.
89        //
90        // This loop ensures that the task is scheduled as required, because it
91        // only returns Pending when a future or service returns Pending.
92        loop {
93            match this.state.as_mut().project() {
94                ResponseStateProj::PollResponse1 { fut, .. } => {
95                    let result = ready!(fut.poll(cx)).map_err(Into::into);
96
97                    if this.policy.should_fallback(&result) {
98                        tracing::debug!("fallback policy selected svc2");
99                        if let __ResponseStateProjectionOwned::PollResponse1 { req, svc2, .. } =
100                            this.state.as_mut().project_replace(ResponseState::Tmp)
101                        {
102                            this.state.set(ResponseState::PollReady2 { req, svc2 });
103                        } else {
104                            unreachable!();
105                        }
106                    } else {
107                        return Poll::Ready(result);
108                    }
109                }
110                ResponseStateProj::PollReady2 { svc2, .. } => match ready!(svc2.poll_ready(cx)) {
111                    Err(e) => return Poll::Ready(Err(e.into())),
112                    Ok(()) => {
113                        if let __ResponseStateProjectionOwned::PollReady2 { mut svc2, req } =
114                            this.state.as_mut().project_replace(ResponseState::Tmp)
115                        {
116                            this.state.set(ResponseState::PollResponse2 {
117                                fut: svc2.call(req),
118                            });
119                        } else {
120                            unreachable!();
121                        }
122                    }
123                },
124                ResponseStateProj::PollResponse2 { fut } => {
125                    return fut.poll(cx).map_err(Into::into)
126                }
127                ResponseStateProj::Tmp => unreachable!(),
128            }
129        }
130    }
131}
132
133impl<S1, S2, Request, F> Debug for ResponseFuture<S1, S2, Request, F>
134where
135    S1: Service<Request>,
136    S2: Service<Request, Response = <S1 as Service<Request>>::Response>,
137    F: FallbackPolicy<<S1 as Service<Request>>::Response>,
138    S1::Error: Into<BoxedError>,
139    Request: Debug,
140    S1::Future: Debug,
141    S2: Debug,
142    S2::Future: Debug,
143    S2::Error: Into<BoxedError>,
144{
145    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
146        f.debug_struct("ResponseFuture")
147            .field("state", &self.state)
148            .finish()
149    }
150}
151
152impl<S1, S2, Request> Debug for ResponseState<S1, S2, Request>
153where
154    S1: Service<Request>,
155    S2: Service<Request, Response = <S1 as Service<Request>>::Response>,
156    Request: Debug,
157    S1::Future: Debug,
158    S2: Debug,
159    S2::Future: Debug,
160    S2::Error: Into<BoxedError>,
161{
162    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
163        match self {
164            ResponseState::PollResponse1 { fut, req, svc2 } => f
165                .debug_struct("ResponseState::PollResponse1")
166                .field("fut", fut)
167                .field("req", req)
168                .field("svc2", svc2)
169                .finish(),
170            ResponseState::PollReady2 { req, svc2 } => f
171                .debug_struct("ResponseState::PollReady2")
172                .field("req", req)
173                .field("svc2", svc2)
174                .finish(),
175            ResponseState::PollResponse2 { fut } => f
176                .debug_struct("ResponseState::PollResponse2")
177                .field("fut", fut)
178                .finish(),
179            ResponseState::Tmp => unreachable!(),
180        }
181    }
182}