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