ServiceExt 与 Oneshot

impl<T: ?Sized, Request> ServiceExt<Request> for T where T: tower_service::Service<Request> {}

pub struct ReadyOneshot<T, Request> {
    inner: Option<T>,
    _p: PhantomData<fn() -> Request>,
}

// Safety: This is safe because `Services`'s are always `Unpin`.
impl<T, Request> Unpin for ReadyOneshot<T, Request> {}

impl<T, Request> Future for ReadyOneshot<T, Request>
where
    T: Service<Request>,
{
    type Output = Result<T, T::Error>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        ready!(self
            .inner
            .as_mut()
            .expect("poll after Poll::Ready")
            .poll_ready(cx))?;

        Poll::Ready(Ok(self.inner.take().expect("poll after Poll::Ready")))
    }
}

pub struct Ready<'a, T, Request>(ReadyOneshot<&'a mut T, Request>);

pin_project! {
    #[project = StateProj]
    enum State<S: Service<Req>, Req> {
        NotReady {
            svc: S,
            req: Option<Req>,
        },
        Called {
            #[pin]
            fut: S::Future,
        },
        Done,
    }
}

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut this = self.project();
        loop {
            match this.state.as_mut().project() {
                StateProj::NotReady { svc, req } => {
                    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"),
            }
        }
    }

tower = { version = "=0.5.3", features = ["util"] }
tokio = { version = "=1.54.1", features = ["full"] }

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use tower::{Service, ServiceExt};

struct Svc {
  polls: u32,
}

impl Service<u32> for Svc {
  type Response = u32;
  type Error = &'static str;
  type Future = Pin<Box<dyn Future<Output = Result<u32, &'static str>>>>;

  fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
    self.polls += 1;
    if self.polls < 3 {
      println!("poll_ready #{} -> Pending", self.polls);
      cx.waker().wake_by_ref();
      Poll::Pending
    } else {
      println!("poll_ready #{} -> Ready", self.polls);
      Poll::Ready(Ok(()))
    }
  }

  fn call(&mut self, req: u32) -> Self::Future {
    println!("call({req})");
    Box::pin(async move {
      tokio::task::yield_now().await;
      println!("response future done");
      Ok(req + 1)
    })
  }
}

impl Drop for Svc {
  fn drop(&mut self) {
    println!("Svc dropped");
  }
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let rsp = Svc { polls: 0 }.oneshot(41).await;
  println!("oneshot -> {rsp:?}");
}

poll_ready #1 -> Pending
poll_ready #2 -> Pending
poll_ready #3 -> Ready
call(41)
Svc dropped
response future done
oneshot -> Ok(42)

use std::future::{ready, Ready};
use std::task::{Context, Poll};
use tokio_test::task;
use tower::{Service, ServiceExt};

#[derive(Debug)]
struct Fails;

impl Service<u32> for Fails {
  type Response = u32;
  type Error = &'static str;
  type Future = Ready<Result<u32, &'static str>>;
  fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
    Poll::Ready(Ok(()))
  }
  fn call(&mut self, _req: u32) -> Self::Future {
    ready(Err("boom"))
  }
}

fn main() {
  let mut fut = task::spawn(Fails.oneshot(1));
  println!("first poll: {:?}", fut.poll());
  println!("{:?}", fut.enter(|_, f| format!("{:?}", f)));
  println!("second poll: {:?}", fut.poll());
}

first poll: Ready(Err("boom"))
"Oneshot { state: State::Called(\"S::Future\") }"

thread 'main' (16814438) panicked at ~/.rustup/toolchains/stable-aarch64-apple-darwin/lib/rustlib/src/rust/library/core/src/future/ready.rs:23:35:
`Ready` polled after completion
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace

        loop {
            // First, see if we have any responses to yield
            if let Poll::Ready(r) = this.queue.poll(cx) {
                if let Some(rsp) = r.transpose()? {
                    return Poll::Ready(Some(Ok(rsp)));
                }
            }

            // If there are no more requests coming, check if we're done
            if *this.eof {
                if this.queue.is_empty() {
                    return Poll::Ready(None);
                } else {
                    return Poll::Pending;
                }
            }

            // If not done, and we don't have a stored request, gather the next request from the
            // stream (if there is one), or return `Pending` if the stream is not ready.
            if this.curr_req.is_none() {
                *this.curr_req = match ready!(this.stream.as_mut().poll_next(cx)) {
                    Some(next_req) => Some(next_req),
                    None => {
                        // Mark that there will be no more requests.
                        *this.eof = true;
                        continue;
                    }
                };
            }

            // Then, see that the service is ready for another request
            let svc = this
                .service
                .as_mut()
                .expect("Using CallAll after extracting inner Service");

            if let Err(e) = ready!(svc.poll_ready(cx)) {
                // Set eof to prevent the service from being called again after a `poll_ready` error
                *this.eof = true;
                return Poll::Ready(Some(Err(e)));
            }

            // Unwrap: The check above always sets `this.curr_req` if none.
            this.queue.push(svc.call(this.curr_req.take().unwrap()));
        }

tower = { version = "=0.5.3", features = ["util"] }
tokio = { version = "=1.54.1", features = ["full", "test-util"] }
futures = "0.3"

use futures::stream::{self, StreamExt};
use std::convert::Infallible;
use std::time::Duration;
use tower::{service_fn, ServiceExt};

#[tokio::main(flavor = "current_thread", start_paused = true)]
async fn main() {
  // request n takes n * 10ms
  let svc = service_fn(|n: u64| async move {
    tokio::time::sleep(Duration::from_millis(n * 10)).await;
    Ok::<_, Infallible>(n)
  });
  let reqs = [3u64, 1, 2];

  let ordered: Vec<u64> = svc
    .call_all(stream::iter(reqs))
    .map(Result::unwrap)
    .collect()
    .await;
  println!("call_all:             {ordered:?}");

  let unordered: Vec<u64> = svc
    .call_all(stream::iter(reqs))
    .unordered()
    .map(Result::unwrap)
    .collect()
    .await;
  println!("call_all().unordered: {unordered:?}");
}

call_all:             [3, 1, 2]
call_all().unordered: [1, 2, 3]

Err("poll_ready failed")
Ok(1)
Err("request 2 failed")
poll_ready = 3, call = 2