从零写一个跨三个时刻的中间件

tower-service = "0.3"
tower = { version = "=0.5.3", features = ["util"] }   # 只用来在 main 里 service_fn、oneshot
tokio = { version = "=1.54.1", features = ["macros", "rt", "time"] }
http = "1"
http-body = "1"
http-body-util = "0.1"
bytes = "1"
futures = "0.3"
pin-project-lite = "0.2"

enum Classified {
  Ready(Verdict),
  RequiresEos,
}

// 5xx 当场判失败；带 `x-status: trailer` 的响应要等 trailers
fn classify_head(res: &http::response::Parts) -> Classified {
  if res.status.is_server_error() {
    Classified::Ready(Verdict::Failed(format!("status {}", res.status.as_u16())))
  } else if res.headers.get("x-status").is_some_and(|v| v == "trailer") {
    Classified::RequiresEos
  } else {
    Classified::Ready(Verdict::Ok)
  }
}

fn classify_eos(trailers: Option<&HeaderMap>) -> Verdict {
  match trailers.and_then(|t| t.get("x-result")) {
    Some(v) if v == "0" => Verdict::Ok,
    Some(v) => Verdict::Failed(format!("x-result {}", v.to_str().unwrap_or("?"))),
    None => Verdict::Failed("no trailers".into()),
  }
}

struct Guard {
  uri: String,
  start: Instant,
  stage: &'static str,
  done: bool,
}

impl Guard {
  fn finish(&mut self, what: String) {
    self.done = true;
    println!("  [{}] {what} ({} ms)", self.uri, self.start.elapsed().as_millis());
  }
}

impl Drop for Guard {
  fn drop(&mut self) {
    if !self.done {
      // drop 里不要 panic：可能正在 unwind，再 panic 会 abort
      println!("  [{}] abandoned during {}", self.uri, self.stage);
    }
  }
}

impl<S, B, ResBody> Service<Request<B>> for Observe<S>
where
  S: Service<Request<B>, Response = Response<ResBody>>,
  S::Error: std::fmt::Display,
  ResBody: Body,
  ResBody::Error: std::fmt::Display,
{
  type Response = Response<ObserveBody<ResBody>>;
  type Error = S::Error;
  type Future = ObserveFuture<S::Future>;

  fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), S::Error>> {
    self.inner.poll_ready(cx)
  }

  fn call(&mut self, req: Request<B>) -> Self::Future {
    let uri = req.uri().path().to_owned();
    println!("  [{uri}] call {}", req.method());
    let guard = Guard { uri, start: Instant::now(), stage: "response future", done: false };
    ObserveFuture { inner: self.inner.call(req), guard: Some(guard) }
  }
}

pin_project! {
  struct ObserveFuture<F> {
    #[pin]
    inner: F,
    guard: Option<Guard>,
  }
}

  fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
    let this = self.project();
    let result = ready!(this.inner.poll(cx));
    let Some(mut guard) = this.guard.take() else {
      unreachable!("ObserveFuture polled after it returned Ready");
    };
    let res = match result {
      Ok(res) => res,
      Err(e) => {
        guard.finish(format!("failed: service error {e}"));
        return Poll::Ready(Err(e));
      }
    };
    let (parts, body) = res.into_parts();
    let eos = match classify_head(&parts) {
      Classified::Ready(v) => { /* 打印结论 */ false }
      Classified::RequiresEos => { /* 打印「等 trailers」 */ true }
    };
    guard.stage = "body";
    let mut body = ObserveBody { inner: body, bytes: 0, eos, guard };
    if body.inner.is_end_stream() {
      body.guard.finish(end_msg(0, eos, None));
    }
    Poll::Ready(Ok(Response::from_parts(parts, body)))
  }

  fn poll_frame(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Result<Frame<B::Data>, B::Error>>> {
    let mut this = self.project();
    if this.guard.done {
      return this.inner.poll_frame(cx);
    }
    let frame = ready!(this.inner.as_mut().poll_frame(cx));
    match &frame {
      None => this.guard.finish(end_msg(*this.bytes, *this.eos, None)),
      Some(Err(e)) => this.guard.finish(format!("failed: body error {e}")),
      Some(Ok(f)) => {
        if let Some(t) = f.trailers_ref() {
          this.guard.finish(end_msg(*this.bytes, *this.eos, Some(t)));
        } else if let Some(d) = f.data_ref() {
          *this.bytes += d.remaining();
          if this.inner.is_end_stream() {
            this.guard.finish(end_msg(*this.bytes, *this.eos, None));
          }
        }
      }
    }
    Poll::Ready(frame)
  }

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let svc = Observe { inner: tower::service_fn(handler) };
  for path in ["/ok", "/boom", "/grpc"] {
    println!("GET {path}");
    let res = svc.clone().oneshot(Request::get(path).body(()).unwrap()).await.unwrap();
    let _ = res.into_body().collect().await;
  }
  println!("GET /stream, drop after one frame");
  let res = svc.clone().oneshot(Request::get("/stream").body(()).unwrap()).await.unwrap();
  let mut body = res.into_body();
  let _ = body.frame().await;
  drop(body);
  println!("GET /slow, give up after 50 ms");
  let _ = tokio::time::timeout(Duration::from_millis(50), svc.clone().oneshot(Request::get("/slow").body(()).unwrap())).await;
}

GET /ok
  [/ok] call GET
  [/ok] head 200 -> Ok
  [/ok] end of stream, 5 bytes (0 ms)
GET /boom
  [/boom] call GET
  [/boom] head 503 -> Failed("status 503")
  [/boom] end of stream, 4 bytes (0 ms)
GET /grpc
  [/grpc] call GET
  [/grpc] head 200 -> wait for trailers
  [/grpc] end of stream, 3 bytes, Failed("x-result 13") (0 ms)
GET /stream, drop after one frame
  [/stream] call GET
  [/stream] head 200 -> Ok
  [/stream] abandoned during body
GET /slow, give up after 50 ms
  [/slow] call GET
  [/slow] abandoned during response future