hyper_util/service/
oneshot.rs

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