InFlightRequests：在途计数

pub struct InFlightRequestsCounter {
    count: Arc<AtomicUsize>,
}

impl InFlightRequestsCounter {
    pub fn get(&self) -> usize {
        self.count.load(Ordering::Relaxed)
    }

    fn increment(&self) -> IncrementGuard {
        self.count.fetch_add(1, Ordering::Relaxed);
        IncrementGuard {
            count: self.count.clone(),
        }
    }
    // ...
}

struct IncrementGuard {
    count: Arc<AtomicUsize>,
}

impl Drop for IncrementGuard {
    fn drop(&mut self) {
        self.count.fetch_sub(1, Ordering::Relaxed);
    }
}

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        let response = ready!(this.inner.poll(cx))?;
        let guard = this.guard.take().unwrap();
        let response = response.map(move |body| ResponseBody { inner: body, guard });

        Poll::Ready(Ok(response))
    }

tower-http = { version = "=0.7.1", features = ["metrics"] }
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"

use bytes::Bytes;
use http::{Request, Response};
use http_body::Frame;
use http_body_util::{combinators::UnsyncBoxBody, BodyExt, Full, StreamBody};
use std::{convert::Infallible, time::Duration};
use tokio::time::{sleep, Instant};
use tower::{BoxError, ServiceBuilder, ServiceExt};
use tower_http::metrics::InFlightRequestsLayer;

type Body = UnsyncBoxBody<Bytes, BoxError>;

async fn handler(req: Request<()>) -> Result<Response<Body>, Infallible> {
  sleep(Duration::from_millis(100)).await;
  let body = if req.uri().path() == "/sse" {
    // 像 SSE 一样：发一帧，然后一直不结束
    let first = futures::stream::iter([Ok::<_, BoxError>(Frame::data(Bytes::from("data: hi\n\n")))]);
    StreamBody::new(futures::StreamExt::chain(first, futures::stream::pending())).boxed_unsync()
  } else {
    Full::from("hello").map_err(|e| match e {}).boxed_unsync()
  };
  Ok(Response::new(body))
}

#[tokio::main(flavor = "current_thread", start_paused = true)]
async fn main() {
  let start = Instant::now();
  let (layer, counter) = InFlightRequestsLayer::pair();
  // 每秒上报一次；所有别的 counter 副本都没了，它自己退出
  let emitter = tokio::spawn(counter.clone().run_emitter(Duration::from_secs(1), move |n| async move {
    println!("  [emit {:>4}ms] in flight = {n}", start.elapsed().as_millis());
  }));
  let svc = ServiceBuilder::new().layer(layer).service_fn(handler);
  let t = || start.elapsed().as_millis();

  let fut = svc.clone().oneshot(Request::get("/").body(()).unwrap());
  let mut fut = Box::pin(fut);
  let _ = tokio::time::timeout(Duration::from_millis(10), &mut fut).await;
  println!("{:>4}ms future pending:         {}", t(), counter.get());
  let res = fut.await.unwrap();
  println!("{:>4}ms response head:          {}", t(), counter.get());
  let mut body = res.into_body();
  while body.frame().await.is_some() {}
  println!("{:>4}ms body read to the end:   {} (is_end_stream: {})", t(), counter.get(), http_body::Body::is_end_stream(&body));
  drop(body);
  println!("{:>4}ms body dropped:           {}", t(), counter.get());

  let res = svc.clone().oneshot(Request::get("/sse").body(()).unwrap()).await.unwrap();
  let mut body = res.into_body();
  let _ = body.frame().await;
  println!("{:>4}ms sse first event read:   {}", t(), counter.get());
  sleep(Duration::from_millis(2500)).await;
  drop(body);
  println!("{:>4}ms sse body dropped:       {}", t(), counter.get());
  sleep(Duration::from_millis(1000)).await;

  drop(svc);
  drop(counter);
  emitter.await.unwrap();
  println!("{:>4}ms emitter returned", t());
}

  [emit    0ms] in flight = 1
  10ms future pending:         1
 100ms response head:          1
 100ms body read to the end:   1 (is_end_stream: true)
 100ms body dropped:           0
 200ms sse first event read:   1
  [emit 1000ms] in flight = 1
  [emit 2000ms] in flight = 1
2700ms sse body dropped:       0
  [emit 3000ms] in flight = 0
  [emit 4000ms] in flight = 0
4000ms emitter returned

        let mut interval = tokio::time::interval(interval);

        loop {
            // if all producers have gone away we don't need to emit anymore
            match Arc::try_unwrap(self.count) {
                Ok(_) => return,
                Err(shared_count) => {
                    self = Self {
                        count: shared_count,
                    }
                }
            }

            interval.tick().await;
            emit(self.get()).await;
        }