一个请求穿过 axum

axum = { version = "=0.8.9", default-features = false, features = ["matched-path"] }
tower = { version = "=0.5.3", features = ["util"] }
http-body-util = "=0.1.5"
tokio = { version = "=1.54.1", features = ["macros", "rt"] }

use std::any::type_name_of_val;
use std::sync::atomic::{AtomicUsize, Ordering};

use axum::body::Body;
use axum::extract::{FromRequestParts, MatchedPath, Path, Request, State};
use axum::http::{request::Parts, StatusCode};
use axum::middleware::{from_fn, Next};
use axum::response::Response;
use axum::routing::get;
use axum::Router;
use http_body_util::BodyExt;
use tower::{Layer, ServiceExt};

// 数 AppState 被 clone 了几次
static CLONES: AtomicUsize = AtomicUsize::new(0);
// 数 Layer::layer 被调用了几次
static LAYERS: AtomicUsize = AtomicUsize::new(0);

fn clones() -> usize {
  CLONES.load(Ordering::SeqCst)
}

struct AppState {
  name: &'static str,
}

impl Clone for AppState {
  fn clone(&self) -> Self {
    CLONES.fetch_add(1, Ordering::SeqCst);
    Self { name: self.name }
  }
}

// 自定义提取器：从 x-caller 头里取调用者
struct Caller(String);

impl<S: Send + Sync> FromRequestParts<S> for Caller {
  type Rejection = (StatusCode, &'static str);

  async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
    println!("      Caller::from_request_parts");
    let value = parts.headers.get("x-caller").and_then(|v| v.to_str().ok());
    match value {
      Some(v) => Ok(Caller(v.to_owned())),
      None => Err((StatusCode::BAD_REQUEST, "missing x-caller")),
    }
  }
}

// 包住 from_fn 的 Layer，只为了打印 layer() 什么时候被调用
#[derive(Clone)]
struct LogLayer<L>(L);

impl<L: Layer<I>, I> Layer<I> for LogLayer<L> {
  type Service = L::Service;

  fn layer(&self, inner: I) -> Self::Service {
    let n = LAYERS.fetch_add(1, Ordering::SeqCst) + 1;
    println!("      Layer::layer #{n}");
    self.0.layer(inner)
  }
}

async fn log_mw(req: Request, next: Next) -> Response {
  let matched = req.extensions().get::<MatchedPath>().map(|m| m.as_str().to_owned());
  println!("    middleware before: {} {} matched={:?}", req.method(), req.uri(), matched);
  let res = next.run(req).await;
  println!("    middleware after:  {}", res.status());
  res
}

async fn get_user(State(state): State<AppState>, caller: Caller, Path(id): Path<u32>) -> String {
  println!("      handler get_user(id={id})");
  format!("user {id} on {}, asked by {}", state.name, caller.0)
}

async fn health() -> &'static str {
  println!("      handler health");
  "ok"
}

async fn send(app: &Router, method: &str, uri: &str) {
  let before = clones();
  println!("\n[request] {method} {uri}");
  let req = Request::builder()
    .method(method)
    .uri(uri)
    .header("x-caller", "alice")
    .body(Body::empty())
    .unwrap();
  let res = app.clone().oneshot(req).await.unwrap();
  let status = res.status();
  let allow = res.headers().get("allow").map(|v| v.to_str().unwrap().to_owned());
  let body = res.into_body().collect().await.unwrap().to_bytes();
  println!("  -> {status} allow={allow:?} body={body:?}");
  println!("  AppState clones during this request: {}", clones() - before);
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  println!("[build] start");
  let state = AppState { name: "demo" };

  let users = get(get_user);
  println!("[build] get(get_user)      : {}", type_name_of_val(&users));

  let app = Router::new()
    .route("/users/{id}", users)
    .route("/health", get(health));
  println!("[build] .route x2          : {}", type_name_of_val(&app));

  println!("[build] .layer(...)");
  let app = app.layer(LogLayer(from_fn(log_mw)));

  println!("[build] .with_state(...)");
  let app: Router = app.with_state(state);
  println!("[build] after with_state   : {}", type_name_of_val(&app));
  println!(
    "[build] done: Layer::layer calls = {}, AppState clones = {}",
    LAYERS.load(Ordering::SeqCst),
    clones()
  );

  let layers_before = LAYERS.load(Ordering::SeqCst);
  send(&app, "GET", "/users/7").await;
  send(&app, "GET", "/users/abc").await;
  send(&app, "GET", "/health").await;
  send(&app, "POST", "/health").await;
  send(&app, "GET", "/nope").await;
  println!(
    "\n[after 5 requests] Layer::layer calls = {}",
    LAYERS.load(Ordering::SeqCst) - layers_before
  );
}

[build] start
[build] get(get_user)      : axum::routing::method_routing::MethodRouter<walk::AppState>
[build] .route x2          : axum::routing::Router<walk::AppState>
[build] .layer(...)
      Layer::layer #1
      Layer::layer #2
      Layer::layer #3
      Layer::layer #4
      Layer::layer #5
[build] .with_state(...)
      Layer::layer #6
      Layer::layer #7
[build] after with_state   : axum::routing::Router
[build] done: Layer::layer calls = 7, AppState clones = 6

[request] GET /users/7
    middleware before: GET /users/7 matched=Some("/users/{id}")
      Caller::from_request_parts
      handler get_user(id=7)
    middleware after:  200 OK
  -> 200 OK allow=None body=b"user 7 on demo, asked by alice"
  AppState clones during this request: 5

[request] GET /users/abc
    middleware before: GET /users/abc matched=Some("/users/{id}")
      Caller::from_request_parts
    middleware after:  400 Bad Request
  -> 400 Bad Request allow=None body=b"Invalid URL: Cannot parse `abc` to a `u32`"
  AppState clones during this request: 5

[request] GET /health
    middleware before: GET /health matched=Some("/health")
      handler health
    middleware after:  200 OK
  -> 200 OK allow=None body=b"ok"
  AppState clones during this request: 4

[request] POST /health
    middleware before: POST /health matched=Some("/health")
    middleware after:  405 Method Not Allowed
  -> 405 Method Not Allowed allow=Some("GET,HEAD") body=b""
  AppState clones during this request: 0

[request] GET /nope
    middleware before: GET /nope matched=None
    middleware after:  404 Not Found
  -> 404 Not Found allow=None body=b""
  AppState clones during this request: 0

[after 5 requests] Layer::layer calls = 0

        self.on_endpoint(
            filter,
            MethodEndpoint::BoxedHandler(BoxedIntoRoute::from_handler(handler)),
        )

        map_inner!(self, this => RouterInner {
            path_router: this.path_router.layer(layer.clone()),
            fallback_router: this.fallback_router.layer(layer.clone()),
            default_fallback: this.default_fallback,
            catch_all_fallback: this.catch_all_fallback.map(|route| route.layer(layer)),
        })

    fn into_route(self: Box<Self>, state: S) -> Route<E2> {
        (self.layer)(self.inner.into_route(state))
    }

            fn call(self, req: Request, state: S) -> Self::Future {
                let (mut parts, body) = req.into_parts();
                Box::pin(async move {
                    $(
                        let $ty = match $ty::from_request_parts(&mut parts, &state).await {
                            Ok(value) => value,
                            Err(rejection) => return rejection.into_response(),
                        };
                    )*

                    let req = Request::from_parts(parts, body);

                    let $last = match $last::from_request(req, &state).await {
                        Ok(value) => value,
                        Err(rejection) => return rejection.into_response(),
                    };

                    self($($ty,)* $last,).await.into_response()
                })
            }

        fn call(&mut self, _req: serve::IncomingStream<'_, L>) -> Self::Future {
            // call `Router::with_state` such that everything is turned into `Route` eagerly
            // rather than doing that per request
            std::future::ready(Ok(self.clone().with_state(())))
        }

axum = { version = "=0.8.9", default-features = false, features = ["tokio", "http1"] }
tower = { version = "=0.5.3", features = ["util"] }
hyper = { version = "=1.12.0", features = ["client", "http1"] }
hyper-util = { version = "=0.1.21", features = ["tokio"] }
tokio = { version = "=1.54.1", features = ["macros", "rt", "net"] }

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let app: Router = Router::new()
    .route("/", get(|| async { "hi" }))
    .layer(LogLayer(from_fn(mw)));
  println!("after build: {}", layers());

  let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
  let addr = listener.local_addr().unwrap();
  tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });

  for conn_no in 1..=2 {
    let before = layers();
    let stream = tokio::net::TcpStream::connect(addr).await.unwrap();
    let (mut sender, conn) = hyper::client::conn::http1::handshake(TokioIo::new(stream)).await.unwrap();
    tokio::spawn(conn);
    for _ in 0..3 {
      let req = Request::get("/").header("host", "x").body(Body::empty()).unwrap();
      let res = sender.send_request(req).await.unwrap();
      assert_eq!(res.status(), 200);
    }
    println!("connection {conn_no}, 3 requests: Layer::layer calls = {}", layers() - before);
  }
}

after build: 4
connection 1, 3 requests: Layer::layer calls = 1
connection 2, 3 requests: Layer::layer calls = 1

after into_make_service: 5
connection 1, 3 requests: Layer::layer calls = 0
connection 2, 3 requests: Layer::layer calls = 0

after build: Layer::layer calls = 4
without with_state: Layer::layer calls during 2 requests = 2
after with_state(()): Layer::layer calls during 2 requests = 0