hyper_util/service/
oneshot.rs

1use futures_util::ready;
2use pin_project_lite::pin_project;
3use std::future::Future;
4use std::pin::Pin;
5use std::task::{Context, Poll};
6use tower_service::Service;
7
8// Vendored from tower::util to reduce dependencies, the code is small enough.
9
10// Not really pub, but used in a trait for bounds
11pin_project! {
12    #[project = OneshotProj]
13    #[derive(Debug)]
14    pub enum Oneshot<S: Service<Req>, Req> {
15        NotReady {
16            svc: S,
17            req: Option<Req>,
18        },
19        Called {
20            #[pin]
21            fut: S::Future,
22        },
23        Done,
24    }
25}
26
27impl<S, Req> Oneshot<S, Req>
28where
29    S: Service<Req>,
30{
31    pub(crate) const fn new(svc: S, req: Req) -> Self {
32        Oneshot::NotReady {
33            svc,
34            req: Some(req),
35        }
36    }
37}
38
39impl<S, Req> Future for Oneshot<S, Req>
40where
41    S: Service<Req>,
42{
43    type Output = Result<S::Response, S::Error>;
44
45    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
46        loop {
47            let this = self.as_mut().project();
48            match this {
49                OneshotProj::NotReady { svc, req } => {
50                    let _ = ready!(svc.poll_ready(cx))?;
51                    let fut = svc.call(req.take().expect("already called"));
52                    self.set(Oneshot::Called { fut });
53                }
54                OneshotProj::Called { fut } => {
55                    let res = ready!(fut.poll(cx))?;
56                    self.set(Oneshot::Done);
57                    return Poll::Ready(Ok(res));
58                }
59                OneshotProj::Done => panic!("polled after complete"),
60            }
61        }
62    }
63}