Either 与 FutureService

#[derive(Clone, Copy, Debug)]
pub enum Either<A, B> {
    #[allow(missing_docs)]
    Left(A),
    #[allow(missing_docs)]
    Right(B),
}

impl<A, B, Request> Service<Request> for Either<A, B>
where
    A: Service<Request>,
    B: Service<Request, Response = A::Response, Error = A::Error>,
{
    type Response = A::Response;
    type Error = A::Error;
    type Future = EitherResponseFuture<A::Future, B::Future>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        match self {
            Either::Left(service) => service.poll_ready(cx),
            Either::Right(service) => service.poll_ready(cx),
        }
    }

    fn call(&mut self, request: Request) -> Self::Future {
        match self {
            Either::Left(service) => EitherResponseFuture {
                kind: Kind::Left {
                    inner: service.call(request),
                },
            },
            Either::Right(service) => EitherResponseFuture {
                kind: Kind::Right {
                    inner: service.call(request),
                },
            },
        }
    }
}

tower = { version = "=0.5.3", features = ["limit", "timeout", "util"] }
tokio = { version = "=1.54.1", features = ["macros", "rt-multi-thread"] }

use std::time::Duration;
use tower::util::Either;
use tower::{service_fn, BoxError, Service, ServiceBuilder, ServiceExt};

async fn handle(req: String) -> Result<String, BoxError> {
  Ok(req.to_uppercase())
}

fn build(strict: bool) -> impl Service<String, Response = String, Error = BoxError> {
  if strict {
    Either::Left(
      ServiceBuilder::new()
        .timeout(Duration::from_secs(1))
        .concurrency_limit(10)
        .service_fn(handle),
    )
  } else {
    Either::Right(service_fn(handle))
  }
}

#[tokio::main]
async fn main() {
  for strict in [true, false] {
    let r = build(strict).oneshot("hi".to_string()).await.unwrap();
    println!("strict={strict}: {r}");
  }
}

strict=true: HI
strict=false: HI

#[derive(Clone)]
pub struct FutureService<F, S> {
    state: State<F, S>,
}

#[derive(Clone)]
enum State<F, S> {
    Future(F),
    Service(S),
}

impl<F, S, R, E> Service<R> for FutureService<F, S>
where
    F: Future<Output = Result<S, E>> + Unpin,
    S: Service<R, Error = E>,
{
    type Response = S::Response;
    type Error = E;
    type Future = S::Future;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        loop {
            self.state = match &mut self.state {
                State::Future(fut) => {
                    let fut = Pin::new(fut);
                    let svc = std::task::ready!(fut.poll(cx)?);
                    State::Service(svc)
                }
                State::Service(svc) => return svc.poll_ready(cx),
            };
        }
    }

    fn call(&mut self, req: R) -> Self::Future {
        if let State::Service(svc) = &mut self.state {
            svc.call(req)
        } else {
            panic!("FutureService::call was called before FutureService::poll_ready")
        }
    }
}

tower = { version = "=0.5.3", features = ["util"] }
tower-test = "0.4"
tokio = { version = "=1.54.1", features = ["sync"] }
tokio-test = "0.4"

use std::panic::{catch_unwind, AssertUnwindSafe};
use tokio::sync::oneshot;
use tokio_test::task;
use tower::util::FutureService;
use tower::Service;
use tower_test::mock;

fn main() {
  std::panic::set_hook(Box::new(|_| {}));
  let (svc, mut handle) = mock::pair::<u32, u32>();
  handle.allow(0); // inner service: not ready yet

  let (tx, rx) = oneshot::channel();
  let mut fs = FutureService::new(Box::pin(async move {
    let svc = rx.await.map_err(|_| "maker dropped")?;
    Ok::<_, tower::BoxError>(svc)
  }));

  let mut t = task::spawn(());
  let p = catch_unwind(AssertUnwindSafe(|| fs.call(1))).err().unwrap();
  println!("call before ready: panic {:?}", p.downcast_ref::<&str>().unwrap());
  println!("1 poll_ready: {:?}", t.enter(|cx, _| fs.poll_ready(cx)).map(|r| r.is_ok()));
  tx.send(svc).ok().unwrap();
  println!("  woken after send: {}", t.is_woken());
  println!("2 poll_ready: {:?}", t.enter(|cx, _| fs.poll_ready(cx)).map(|r| r.is_ok()));
  handle.allow(1);
  println!("  woken after allow: {}", t.is_woken());
  println!("3 poll_ready: {:?}", t.enter(|cx, _| fs.poll_ready(cx)).map(|r| r.is_ok()));
}

call before ready: panic "FutureService::call was called before FutureService::poll_ready"
1 poll_ready: Pending
  woken after send: true
2 poll_ready: Pending
  woken after allow: true
3 poll_ready: Ready(true)

poll_ready #1: Ready(Err("connect failed"))
poll_ready #2: panic "`async fn` resumed after completion"