OnEarlyDrop：请求被放弃

pub struct OnEarlyDropGuard<Callback: OnDropCallback> {
    callback: Option<Callback>,
}

impl<Callback: OnDropCallback> OnEarlyDropGuard<Callback> {
    pub fn completed(&mut self) {
        self.callback = None;
    }
}

impl<Callback: OnDropCallback> Drop for OnEarlyDropGuard<Callback> {
    fn drop(&mut self) {
        // Panicking in Drop aborts the process if we are already unwinding,
        // so avoid `expect` here.
        if let Some(callback) = self.callback.take() {
            callback.on_drop();
        }
    }
}

        // Inner resolved: suppress the future-drop guard for both Ok and Err.
        if let Some(guard) = this.future_guard.as_mut() {
            guard.completed();
        }
        *this.future_guard = None;

    pub(crate) fn new(body: B, callback: Callback) -> Self
    where
        B: Body,
    {
        let mut guard = OnEarlyDropGuard::new(callback);
        if body.is_end_stream() {
            guard.completed();
        }
        Self { inner: body, guard }
    }
    // poll_frame 里：
        if matches!(result, None | Some(Err(_))) {
            this.guard.completed();
        }
        if matches!(result, Some(Ok(_))) && this.inner.is_end_stream() {
            this.guard.completed();
        }

    fn make(&mut self, _request: &Request<ReqB>) -> Self::Callback {
        FutureDropFailureCallback {
            start: Instant::now(),
            on_failure: self.on_failure.clone(),
            span: Span::current(),
        }
    }

tower-http = { version = "=0.7.1", features = ["on-early-drop", "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"
tracing-subscriber = "0.3"

use bytes::Bytes;
use http::{Request, Response, StatusCode};
use http_body::Frame;
use http_body_util::{combinators::UnsyncBoxBody, BodyExt, Empty, Full, StreamBody};
use std::{convert::Infallible, time::Duration};
use tower::{BoxError, ServiceBuilder, ServiceExt};
use tower_http::{
  on_early_drop::{EarlyDropsAsFailures, OnBodyDropFn, OnEarlyDropLayer},
  trace::{DefaultMakeSpan, DefaultOnFailure, TraceLayer},
};

type Body = UnsyncBoxBody<Bytes, BoxError>;

async fn handler(req: Request<Full<Bytes>>) -> Result<Response<Body>, Infallible> {
  let res = match req.uri().path() {
    "/slow" => {
      tokio::time::sleep(Duration::from_secs(5)).await;
      Response::new(Empty::new().map_err(|e| match e {}).boxed_unsync())
    }
    "/stream" => {
      let frames = [Ok::<_, BoxError>(Frame::data(Bytes::from("one"))), Ok(Frame::data(Bytes::from("two")))];
      Response::new(StreamBody::new(futures::stream::iter(frames)).boxed_unsync())
    }
    "/no-content" => {
      let mut r = Response::new(Empty::new().map_err(|e| match e {}).boxed_unsync());
      *r.status_mut() = StatusCode::NO_CONTENT;
      r
    }
    _ => Response::new(Full::from("hello").map_err(|e| match e {}).boxed_unsync()),
  };
  Ok(res)
}

#[tokio::main(flavor = "current_thread", start_paused = true)]
async fn main() {
  tracing_subscriber::fmt().without_time().with_target(false).with_ansi(false).init();

  // 1. 直接装闭包：future 一路是「请求时造回调」，body 一路多一层「响应头出来时」
  let svc = ServiceBuilder::new()
    .layer(
      OnEarlyDropLayer::builder()
        .on_future_drop(|req: &Request<Full<Bytes>>| {
          let uri = req.uri().clone();
          move || println!("    future dropped: {uri}")
        })
        .on_body_drop(OnBodyDropFn::new(|req: &Request<Full<Bytes>>| {
          let uri = req.uri().clone();
          move |parts: &http::response::Parts| {
            let status = parts.status;
            move || println!("    body dropped: {uri} ({status})")
          }
        })),
    )
    .service_fn(handler);

  println!("/slow, caller gives up after 10ms");
  let r = tokio::time::timeout(Duration::from_millis(10), svc.clone().oneshot(Request::get("/slow").body(Full::default()).unwrap())).await;
  println!("  timed out: {}", r.is_err());

  println!("/stream, read one frame then drop");
  let mut body = svc.clone().oneshot(Request::get("/stream").body(Full::default()).unwrap()).await.unwrap().into_body();
  let _ = body.frame().await;
  drop(body);

  println!("/stream, read to the end");
  let body = svc.clone().oneshot(Request::get("/stream").body(Full::default()).unwrap()).await.unwrap().into_body();
  let _ = body.collect().await;

  println!("/ (Full), read the only frame, never poll again");
  let mut body = svc.clone().oneshot(Request::get("/").body(Full::default()).unwrap()).await.unwrap().into_body();
  let _ = body.frame().await;
  drop(body);

  println!("/no-content, dropped unread");
  let body = svc.clone().oneshot(Request::get("/no-content").body(Full::default()).unwrap()).await.unwrap().into_body();
  drop(body);

  // 2. 桥到 trace::OnFailure：放在 TraceLayer 里面和外面各一次
  let early = OnEarlyDropLayer::new(EarlyDropsAsFailures::new(DefaultOnFailure::default()));
  // 默认 span 是 DEBUG 级，fmt 默认只开到 INFO，span 会被关掉；这里把 span 提到 INFO
  let trace = TraceLayer::new_for_http().make_span_with(DefaultMakeSpan::new().level(tracing::Level::INFO));
  println!("EarlyDropsAsFailures inside TraceLayer");
  let inside = ServiceBuilder::new().layer(trace.clone()).layer(early.clone()).service_fn(handler);
  let _ = tokio::time::timeout(Duration::from_millis(10), inside.oneshot(Request::get("/slow").body(Full::default()).unwrap())).await;
  println!("EarlyDropsAsFailures outside TraceLayer");
  let outside = ServiceBuilder::new().layer(early).layer(trace).service_fn(handler);
  let _ = tokio::time::timeout(Duration::from_millis(10), outside.oneshot(Request::get("/slow").body(Full::default()).unwrap())).await;
}

/slow, caller gives up after 10ms
    future dropped: /slow
  timed out: true
/stream, read one frame then drop
    body dropped: /stream (200 OK)
/stream, read to the end
/ (Full), read the only frame, never poll again
/no-content, dropped unread
EarlyDropsAsFailures inside TraceLayer
ERROR request{method=GET uri=/slow version=HTTP/1.1}: response failed classification=response future dropped before completion latency=0 ms
EarlyDropsAsFailures outside TraceLayer
ERROR response failed classification=response future dropped before completion latency=0 ms