1 use futures_core::ready;
2 use pin_project_lite::pin_project;
3 use std::{
4     fmt,
5     future::Future,
6     pin::Pin,
7     task::{Context, Poll},
8 };
9 use tower_service::Service;
10 
11 pin_project! {
12     /// A [`Future`] consuming a [`Service`] and request, waiting until the [`Service`]
13     /// is ready, and then calling [`Service::call`] with the request, and
14     /// waiting for that [`Future`].
15     #[derive(Debug)]
16     pub struct Oneshot<S: Service<Req>, Req> {
17         #[pin]
18         state: State<S, Req>,
19     }
20 }
21 
22 pin_project! {
23     #[project = StateProj]
24     enum State<S: Service<Req>, Req> {
25         NotReady {
26             svc: S,
27             req: Option<Req>,
28         },
29         Called {
30             #[pin]
31             fut: S::Future,
32         },
33         Done,
34     }
35 }
36 
37 impl<S: Service<Req>, Req> State<S, Req> {
not_ready(svc: S, req: Option<Req>) -> Self38     fn not_ready(svc: S, req: Option<Req>) -> Self {
39         Self::NotReady { svc, req }
40     }
41 
called(fut: S::Future) -> Self42     fn called(fut: S::Future) -> Self {
43         Self::Called { fut }
44     }
45 }
46 
47 impl<S, Req> fmt::Debug for State<S, Req>
48 where
49     S: Service<Req> + fmt::Debug,
50     Req: fmt::Debug,
51 {
fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result52     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
53         match self {
54             State::NotReady {
55                 svc,
56                 req: Some(req),
57             } => f
58                 .debug_tuple("State::NotReady")
59                 .field(svc)
60                 .field(req)
61                 .finish(),
62             State::NotReady { req: None, .. } => unreachable!(),
63             State::Called { .. } => f.debug_tuple("State::Called").field(&"S::Future").finish(),
64             State::Done => f.debug_tuple("State::Done").finish(),
65         }
66     }
67 }
68 
69 impl<S, Req> Oneshot<S, Req>
70 where
71     S: Service<Req>,
72 {
73     #[allow(missing_docs)]
new(svc: S, req: Req) -> Self74     pub fn new(svc: S, req: Req) -> Self {
75         Oneshot {
76             state: State::not_ready(svc, Some(req)),
77         }
78     }
79 }
80 
81 impl<S, Req> Future for Oneshot<S, Req>
82 where
83     S: Service<Req>,
84 {
85     type Output = Result<S::Response, S::Error>;
86 
poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>87     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
88         let mut this = self.project();
89         loop {
90             match this.state.as_mut().project() {
91                 StateProj::NotReady { svc, req } => {
92                     let _ = ready!(svc.poll_ready(cx))?;
93                     let f = svc.call(req.take().expect("already called"));
94                     this.state.set(State::called(f));
95                 }
96                 StateProj::Called { fut } => {
97                     let res = ready!(fut.poll(cx))?;
98                     this.state.set(State::Done);
99                     return Poll::Ready(Ok(res));
100                 }
101                 StateProj::Done => panic!("polled after complete"),
102             }
103         }
104     }
105 }
106