Trace 的内部

pub struct Trace<
    S,
    M,
    MakeSpan = DefaultMakeSpan,
    OnRequest = DefaultOnRequest,
    OnResponse = DefaultOnResponse,
    OnBodyChunk = DefaultOnBodyChunk,
    OnEos = DefaultOnEos,
    OnFailure = DefaultOnFailure,
> {
    pub(crate) inner: S,
    pub(crate) make_classifier: M,
    pub(crate) make_span: MakeSpan,
    pub(crate) on_request: OnRequest,
    // ...
}

    pub fn on_request<NewOnRequest>(
        self,
        new_on_request: NewOnRequest,
    ) -> Trace<S, M, MakeSpan, NewOnRequest, OnResponse, OnBodyChunk, OnEos, OnFailure> {

        event_dynamic_lvl!(parent: span, self.level, "started processing request");

        let result = ready!(this.inner.as_mut().poll_frame(cx));

        let latency = this.start.elapsed();
        *this.start = Instant::now();

        match result {
            Some(Ok(frame)) => {
                // 数据帧：on_body_chunk
                // trailers 帧：classify_eos(Some(&trailers))，失败就 on_failure；然后 on_eos
                // ...
                if this.inner.is_end_stream() {
                    // classify_eos(None)，失败就 on_failure；然后 on_eos
                }
                Poll::Ready(Some(Ok(frame)))
            }
            Some(Err(err)) => {
                if let Some((classify_eos, mut on_failure)) =
                    this.classify_eos.take().zip(this.on_failure.take())
                {
                    let failure_class = classify_eos.classify_error(&err);
                    on_failure.on_failure(failure_class, latency, this.span);
                }

                Poll::Ready(Some(Err(err)))
            }
            None => {
                // classify_eos(None)，失败就 on_failure；然后 on_eos
            }
        }

tower-http = { version = "=0.7.1", features = ["trace"] }
tower = { version = "=0.5.3", features = ["util"] }
tokio = { version = "=1.54.1", features = ["macros", "rt", "time", "test-util"] }
http = "1"
http-body = "1"
http-body-util = "0.1"
bytes = "1"
futures = "0.3"
tracing = "0.1"

use bytes::Bytes;
use http::{HeaderMap, HeaderValue, Request, Response};
use http_body::Frame;
use http_body_util::{combinators::UnsyncBoxBody, BodyExt, Full, StreamBody};
use std::{convert::Infallible, time::Duration};
use tower::{service_fn, BoxError, Layer, ServiceExt};
use tower_http::{
  classify::{GrpcFailureClass, ServerErrorsFailureClass},
  trace::TraceLayer,
};
use tracing::Span;

type Body = UnsyncBoxBody<Bytes, BoxError>;

// 两种 body：一个以 grpc-status: 13 的 trailers 结束，一个读到一半出错
async fn handler(req: Request<Full<Bytes>>) -> Result<Response<Body>, Infallible> {
  tokio::time::sleep(Duration::from_secs(1)).await;
  let frames: Vec<Result<Frame<Bytes>, BoxError>> = if req.uri().path() == "/trailers" {
    let mut t = HeaderMap::new();
    t.insert("grpc-status", HeaderValue::from_static("13"));
    vec![Ok(Frame::data(Bytes::from("msg"))), Ok(Frame::trailers(t))]
  } else {
    vec![Ok(Frame::data(Bytes::from("part"))), Err("disk gone".into())]
  };
  let mut res = Response::new(StreamBody::new(futures::stream::iter(frames)).boxed_unsync());
  res.headers_mut().insert("content-type", HeaderValue::from_static("application/grpc"));
  Ok(res)
}

#[tokio::main(flavor = "current_thread", start_paused = true)]
async fn main() {
  macro_rules! callbacks {
    ($layer:expr, $class:ty) => {
      $layer
        .on_response(|r: &Response<Body>, l: Duration, _: &Span| println!("    on_response {} (latency < 10ms: {})", r.status(), l < Duration::from_millis(10)))
        .on_body_chunk(|c: &Bytes, _: Duration, _: &Span| println!("    on_body_chunk {:?}", c))
        .on_eos(|t: Option<&HeaderMap>, _: Duration, _: &Span| println!("    on_eos trailers={:?}", t.map(|t| t.get("grpc-status"))))
        .on_failure(|f: $class, _: Duration, _: &Span| println!("    on_failure {f}"))
    };
  }
  for path in ["/trailers", "/error"] {
    println!("new_for_http {path}");
    let svc = callbacks!(TraceLayer::new_for_http(), ServerErrorsFailureClass).layer(service_fn(handler));
    let res = svc.oneshot(Request::get(path).body(Full::default()).unwrap()).await.unwrap();
    let _ = res.into_body().collect().await;
    println!("new_for_grpc {path}");
    let svc = callbacks!(TraceLayer::new_for_grpc(), GrpcFailureClass).layer(service_fn(handler));
    let res = svc.oneshot(Request::get(path).body(Full::default()).unwrap()).await.unwrap();
    let _ = res.into_body().collect().await;
  }
}

new_for_http /trailers
    on_response 200 OK (latency < 10ms: true)
    on_body_chunk b"msg"
    on_eos trailers=Some(Some("13"))
new_for_grpc /trailers
    on_response 200 OK (latency < 10ms: true)
    on_body_chunk b"msg"
    on_failure Status: Internal
    on_eos trailers=Some(Some("13"))
new_for_http /error
    on_response 200 OK (latency < 10ms: true)
    on_body_chunk b"part"
new_for_grpc /error
    on_response 200 OK (latency < 10ms: true)
    on_body_chunk b"part"
    on_failure Error: disk gone