use futures_core::ready; use pin_project::pin_project; use std::{ fmt, future::Future, pin::Pin, task::{Context, Poll}, }; use tower_service::Service; /// A `Future` consuming a `Service` and request, waiting until the `Service` /// is ready, and then calling `Service::call` with the request, and /// waiting for that `Future`. #[pin_project] #[derive(Debug)] pub struct Oneshot, Req> { #[pin] state: State, } #[pin_project(project = StateProj)] enum State, Req> { NotReady(S, Option), Called(#[pin] S::Future), Done, } impl fmt::Debug for State where S: Service + fmt::Debug, Req: fmt::Debug, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { State::NotReady(s, Some(req)) => f .debug_tuple("State::NotReady") .field(s) .field(req) .finish(), State::NotReady(_, None) => unreachable!(), State::Called(_) => f.debug_tuple("State::Called").field(&"S::Future").finish(), State::Done => f.debug_tuple("State::Done").finish(), } } } impl Oneshot where S: Service, { #[allow(missing_docs)] pub fn new(svc: S, req: Req) -> Self { Oneshot { state: State::NotReady(svc, Some(req)), } } } impl Future for Oneshot where S: Service, { type Output = Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let mut this = self.project(); loop { match this.state.as_mut().project() { StateProj::NotReady(svc, req) => { let _ = ready!(svc.poll_ready(cx))?; let f = svc.call(req.take().expect("already called")); this.state.set(State::Called(f)); } StateProj::Called(fut) => { let res = ready!(fut.poll(cx))?; this.state.set(State::Done); return Poll::Ready(Ok(res)); } StateProj::Done => panic!("polled after complete"), } } } }