从零写一个 tower

tokio = { version = "=1.54.1", features = ["rt", "macros", "sync", "time"] }
pin-project-lite = "0.2"

type BoxError = Box<dyn Error + Send + Sync>;

trait Service<Req> {
  type Response;
  type Error;
  type Future: Future<Output = Result<Self::Response, Self::Error>>;
  fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>>;
  fn call(&mut self, req: Req) -> Self::Future;
}

trait Layer<S> {
  type Service;
  fn layer(&self, inner: S) -> Self::Service;
}

struct Identity;

impl<S> Layer<S> for Identity {
  type Service = S;
  fn layer(&self, inner: S) -> S {
    inner
  }
}

struct Stack<Inner, Outer> {
  inner: Inner,
  outer: Outer,
}

impl<S, Inner: Layer<S>, Outer: Layer<Inner::Service>> Layer<S> for Stack<Inner, Outer> {
  type Service = Outer::Service;
  fn layer(&self, service: S) -> Self::Service {
    self.outer.layer(self.inner.layer(service))
  }
}

struct Builder<L>(L);

impl Builder<Identity> {
  fn new() -> Self {
    Builder(Identity)
  }
}

impl<L> Builder<L> {
  // 新加的 layer 放在 Stack 的 inner 一侧：先加的在外层
  fn layer<T>(self, layer: T) -> Builder<Stack<T, L>> {
    Builder(Stack { inner: layer, outer: self.0 })
  }
  fn service<S>(&self, service: S) -> L::Service
  where
    L: Layer<S>,
  {
    self.0.layer(service)
  }
}

impl<S, Req> Service<Req> for Timeout<S>
where
  S: Service<Req>,
  S::Error: Into<BoxError>,
{
  type Response = S::Response;
  type Error = BoxError;
  type Future = TimeoutFuture<S::Future>;

  fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), BoxError>> {
    // 就绪主线：原样转发，等待就绪的时间不计入超时
    self.inner.poll_ready(cx).map_err(Into::into)
  }

  fn call(&mut self, req: Req) -> Self::Future {
    // 请求主线：从 call 开始计时
    TimeoutFuture { inner: self.inner.call(req), sleep: tokio::time::sleep(self.timeout) }
  }
}

pin_project! {
  struct TimeoutFuture<F> {
    #[pin]
    inner: F,
    #[pin]
    sleep: tokio::time::Sleep,
  }
}

impl<F, T, E> Future for TimeoutFuture<F>
where
  F: Future<Output = Result<T, E>>,
  E: Into<BoxError>,
{
  type Output = Result<T, BoxError>;

  fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
    let this = self.project();
    // 先 poll 内层：同一轮里两边都好了，算成功
    if let Poll::Ready(res) = this.inner.poll(cx) {
      return Poll::Ready(res.map_err(Into::into));
    }
    ready!(this.sleep.poll(cx));
    Poll::Ready(Err(Elapsed.into()))
  }
}

type Acquire = Pin<Box<dyn Future<Output = Result<OwnedSemaphorePermit, AcquireError>> + Send>>;

struct ConcurrencyLimit<S> {
  inner: S,
  semaphore: Arc<Semaphore>,
  acquiring: Option<Acquire>,
  permit: Option<OwnedSemaphorePermit>,
}

impl<S, Req> Service<Req> for ConcurrencyLimit<S>
where
  S: Service<Req>,
{
  type Response = S::Response;
  type Error = S::Error;
  type Future = LimitFuture<S::Future>;

  fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), S::Error>> {
    if self.permit.is_none() {
      let sem = &self.semaphore;
      let acquiring = self.acquiring.get_or_insert_with(|| Box::pin(sem.clone().acquire_owned()));
      // 取不到许可就 Pending；acquire_owned 的 future 记下了 waker
      let permit = ready!(acquiring.as_mut().poll(cx)).expect("semaphore is never closed");
      self.acquiring = None;
      self.permit = Some(permit);
    }
    self.inner.poll_ready(cx)
  }

  fn call(&mut self, req: Req) -> Self::Future {
    let permit = self.permit.take().expect("poll_ready must be called first");
    LimitFuture { inner: self.inner.call(req), _permit: permit }
  }
}

pin_project! {
  struct LimitFuture<F> {
    #[pin]
    inner: F,
    _permit: OwnedSemaphorePermit, // future 被 drop 时归还
  }
}

type Responder = oneshot::Sender<Result<String, BoxError>>;

#[derive(Clone)]
struct Mock {
  tx: mpsc::UnboundedSender<(u32, Responder)>,
}

struct Handle {
  rx: mpsc::UnboundedReceiver<(u32, Responder)>,
}

fn mock() -> (Mock, Handle) {
  let (tx, rx) = mpsc::unbounded_channel();
  (Mock { tx }, Handle { rx })
}

impl Service<u32> for Mock {
  type Response = String;
  type Error = BoxError;
  type Future = Pin<Box<dyn Future<Output = Result<String, BoxError>> + Send>>;

  fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), BoxError>> {
    Poll::Ready(Ok(()))
  }

  fn call(&mut self, req: u32) -> Self::Future {
    // 请求交给测试代码，响应由测试代码决定什么时候给
    let (resp_tx, resp_rx) = oneshot::channel();
    let _ = self.tx.send((req, resp_tx));
    Box::pin(async move { resp_rx.await.map_err(|_| BoxError::from("mock dropped"))? })
  }
}

impl Handle {
  fn try_next(&mut self) -> Option<(u32, Responder)> {
    self.rx.try_recv().ok()
  }
}

async fn run<S>(mut svc: S, mut handle: Handle)
where
  S: Service<u32, Response = String, Error = BoxError>,
{
  let flag = Arc::new(Flag(AtomicBool::new(false)));
  println!("1 poll_ready: {}", poll_once(&mut svc, &flag));
  let first = svc.call(1);
  println!("2 poll_ready: {}", poll_once(&mut svc, &flag));

  let (req, responder) = handle.try_next().expect("first request reached the mock");
  println!("3 mock got {req}, more queued: {}", handle.try_next().is_some());

  responder.send(Ok(format!("resp {req}"))).unwrap();
  let first = first.await.map_err(|e| e.to_string());
  println!("4 first: {first:?}, waker woken: {}", flag.0.load(Ordering::SeqCst));
  println!("5 poll_ready: {}", poll_once(&mut svc, &flag));

  let second = svc.call(2);
  let (_req, _responder) = handle.try_next().expect("second request reached the mock");
  println!("6 second: {:?}", second.await.map_err(|e| e.to_string()));
  ready(&mut svc).await.unwrap();
  println!("7 ready again after timeout");
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  println!("Timeout outside ConcurrencyLimit:");
  let (inner, handle) = mock();
  let svc = Builder::new()
    .layer(TimeoutLayer(Duration::from_millis(50)))
    .layer(ConcurrencyLimitLayer(1))
    .service(inner);
  run(svc, handle).await;

  println!("ConcurrencyLimit outside Timeout:");
  let (inner, handle) = mock();
  let svc = Builder::new()
    .layer(ConcurrencyLimitLayer(1))
    .layer(TimeoutLayer(Duration::from_millis(50)))
    .service(inner);
  run(svc, handle).await;
}

Timeout outside ConcurrencyLimit:
1 poll_ready: Ready(Ok)
2 poll_ready: Pending
3 mock got 1, more queued: false
4 first: Ok("resp 1"), waker woken: true
5 poll_ready: Ready(Ok)
6 second: Err("request timed out")
7 ready again after timeout

Timeout outside ConcurrencyLimit:
1 poll_ready: Ready(Ok)
2 poll_ready: Pending
3 mock got 1, more queued: false
4 first: Ok("resp 1"), waker woken: true
5 poll_ready: Ready(Ok)
6 second: Err("request timed out")
7 ready again after timeout
ConcurrencyLimit outside Timeout:
1 poll_ready: Ready(Ok)
2 poll_ready: Pending
3 mock got 1, more queued: false
4 first: Ok("resp 1"), waker woken: true
5 poll_ready: Ready(Ok)
6 second: Err("request timed out")
7 ready again after timeout