tower_util/
oneshot.rs

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

tower_util/
oneshot.rs
