1 | // TODO: Eventually to be replaced with tower_util::Oneshot. |
2 | |
3 | use pin_project_lite::pin_project; |
4 | use tower_service::Service; |
5 | |
6 | use crate::common::{task, Future, Pin, Poll}; |
7 | |
8 | pub(crate) fn oneshot<S, Req>(svc: S, req: Req) -> Oneshot<S, Req> |
9 | where |
10 | S: Service<Req>, |
11 | { |
12 | Oneshot { |
13 | state: State::NotReady { svc, req }, |
14 | } |
15 | } |
16 | |
17 | pin_project! { |
18 | // A `Future` consuming a `Service` and request, waiting until the `Service` |
19 | // is ready, and then calling `Service::call` with the request, and |
20 | // waiting for that `Future`. |
21 | #[allow(missing_debug_implementations)] |
22 | pub struct Oneshot<S: Service<Req>, Req> { |
23 | #[pin] |
24 | state: State<S, Req>, |
25 | } |
26 | } |
27 | |
28 | pin_project! { |
29 | #[project = StateProj] |
30 | #[project_replace = StateProjOwn] |
31 | enum State<S: Service<Req>, Req> { |
32 | NotReady { |
33 | svc: S, |
34 | req: Req, |
35 | }, |
36 | Called { |
37 | #[pin] |
38 | fut: S::Future, |
39 | }, |
40 | Tmp, |
41 | } |
42 | } |
43 | |
44 | impl<S, Req> Future for Oneshot<S, Req> |
45 | where |
46 | S: Service<Req>, |
47 | { |
48 | type Output = Result<S::Response, S::Error>; |
49 | |
50 | fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> { |
51 | let mut me = self.project(); |
52 | |
53 | loop { |
54 | match me.state.as_mut().project() { |
55 | StateProj::NotReady { ref mut svc, .. } => { |
56 | ready!(svc.poll_ready(cx))?; |
57 | // fallthrough out of the match's borrow |
58 | } |
59 | StateProj::Called { fut } => { |
60 | return fut.poll(cx); |
61 | } |
62 | StateProj::Tmp => unreachable!(), |
63 | } |
64 | |
65 | match me.state.as_mut().project_replace(State::Tmp) { |
66 | StateProjOwn::NotReady { mut svc, req } => { |
67 | me.state.set(State::Called { fut: svc.call(req) }); |
68 | } |
69 | _ => unreachable!(), |
70 | } |
71 | } |
72 | } |
73 | } |
74 | |