Route 与类型擦除

pub struct Route<E = Infallible>(BoxCloneSyncService<Request, Response, E>);

impl<E> Route<E> {
    pub(crate) fn new<T>(svc: T) -> Self
    where
        T: Service<Request, Error = E> + Clone + Send + Sync + 'static,
        T::Response: IntoResponse + 'static,
        T::Future: Send + 'static,
    {
        Self(BoxCloneSyncService::new(MapIntoResponse::new(svc)))
    }

        let layer = (MapErrLayer::new(Into::into), layer);

        Route::new(layer.layer(self))

    /// Variant of [`Route::call`] that takes ownership of the route to avoid cloning.
    pub(crate) fn call_owned(self, req: Request<Body>) -> RouteFuture<E> {
        let req = req.map(Body::new);
        self.oneshot_inner_owned(req).not_top_level()
    }

    pub(crate) fn oneshot_inner(&mut self, req: Request) -> RouteFuture<E> {
        let method = req.method().clone();
        RouteFuture::new(method, self.0.clone().oneshot(req))
    }

    /// Variant of [`Route::oneshot_inner`] that takes ownership of the route to avoid cloning.
    pub(crate) fn oneshot_inner_owned(self, req: Request) -> RouteFuture<E> {
        let method = req.method().clone();
        RouteFuture::new(method, self.0.oneshot(req))
    }

axum = { version = "=0.8.9", default-features = false }
tower = "=0.5.3"
tokio = { version = "=1.54.1", features = ["macros", "rt"] }

use std::alloc::{GlobalAlloc, Layout, System};
use std::convert::Infallible;
use std::future::{ready, Ready};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::task::{Context, Poll};

use axum::body::Body;
use axum::extract::Request;
use axum::routing::get_service;
use axum::Router;
use tower::layer::util::Identity;
use tower::Service;

// 数堆分配次数
struct Counting;
static ALLOCS: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for Counting {
  unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
    ALLOCS.fetch_add(1, Ordering::SeqCst);
    unsafe { System.alloc(layout) }
  }
  unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
    unsafe { System.dealloc(ptr, layout) }
  }
}

#[global_allocator]
static GLOBAL: Counting = Counting;

// 数 Svc 被 clone、被 call 了几次
static CLONES: AtomicUsize = AtomicUsize::new(0);
static CALLS: AtomicUsize = AtomicUsize::new(0);

// 带一个字段：零大小类型装箱不分配，量不出来
struct Svc(u64);

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

impl Service<Request> for Svc {
  type Response = &'static str;
  type Error = Infallible;
  type Future = Ready<Result<&'static str, Infallible>>;

  fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Infallible>> {
    Poll::Ready(Ok(()))
  }

  fn call(&mut self, _req: Request) -> Self::Future {
    CALLS.fetch_add(1, Ordering::SeqCst);
    ready(Ok("svc"))
  }
}

fn load(c: &AtomicUsize) -> usize {
  c.load(Ordering::SeqCst)
}

async fn measure(label: &str, app: Router) {
  // 先跑一次热身，再量第二次
  for round in 0..2 {
    let req = Request::get("/").body(Body::empty()).unwrap();
    let mut svc = app.clone();
    let (c0, k0, a0) = (load(&CLONES), load(&CALLS), load(&ALLOCS));
    let res = svc.call(req).await.unwrap();
    let (c1, k1, a1) = (load(&CLONES), load(&CALLS), load(&ALLOCS));
    assert_eq!(res.status(), 200);
    if round == 1 {
      println!("{label:<34} clones={} calls={} allocs={}", c1 - c0, k1 - k0, a1 - a0);
    }
  }
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  measure("route_service(Svc)", Router::new().route_service("/", Svc(1))).await;
  measure("get_service(Svc)", Router::new().route("/", get_service(Svc(1)))).await;
  measure(
    "get_service(Svc) + 1 layer",
    Router::new().route("/", get_service(Svc(1))).layer(Identity::new()),
  )
  .await;
}

route_service(Svc)                 clones=1 calls=1 allocs=8
get_service(Svc)                   clones=1 calls=1 allocs=9
get_service(Svc) + 1 layer         clones=2 calls=1 allocs=12

pub(crate) struct BoxedIntoRoute<S, E>(Box<dyn ErasedIntoRoute<S, E>>);

// ...

pub(crate) trait ErasedIntoRoute<S, E>: Send + Sync {
    fn clone_box(&self) -> Box<dyn ErasedIntoRoute<S, E>>;

    fn into_route(self: Box<Self>, state: S) -> Route<E>;

    #[allow(dead_code)]
    fn call_with_state(self: Box<Self>, request: Request, state: S) -> RouteFuture<E>;
}

pub(crate) struct MakeErasedHandler<H, S> {
    pub(crate) handler: H,
    pub(crate) into_route: fn(H, S) -> Route,
}

    pub struct RouteFuture<E> {
        #[pin]
        inner: Oneshot<BoxCloneSyncService<Request, Response, E>, Request>,
        method: Method,
        allow_header: Option<Bytes>,
        top_level: bool,
    }

        let mut res = ready!(this.inner.poll(cx))?;

        if *this.method == Method::CONNECT && res.status().is_success() {
            // From https://httpwg.org/specs/rfc9110.html#CONNECT:
            // > A server MUST NOT send any Transfer-Encoding or
            // > Content-Length header fields in a 2xx (Successful)
            // > response to CONNECT.
            if res.headers().contains_key(&CONTENT_LENGTH)
                || res.headers().contains_key(&header::TRANSFER_ENCODING)
                || res.size_hint().lower() != 0
            {
                error!("response to CONNECT with nonempty body");
                res = res.map(|_| Body::empty());
            }
        } else if *this.top_level {
            set_allow_header(res.headers_mut(), this.allow_header);

            // make sure to set content-length before removing the body
            set_content_length(res.size_hint(), res.headers_mut());

            if *this.method == Method::HEAD {
                *res.body_mut() = Body::empty();
            }
        }

axum = { version = "=0.8.9", default-features = false, features = ["tokio", "http1"] }
tower = { version = "=0.5.3", features = ["util"] }
http-body-util = "=0.1.5"
futures = "=0.3.34"
tokio = { version = "=1.54.1", features = ["macros", "rt", "net", "io-util"] }

use std::convert::Infallible;

use axum::body::Body;
use axum::extract::Request;
use axum::middleware::map_response;
use axum::response::Response;
use axum::routing::get;
use axum::Router;
use http_body_util::BodyExt;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tower::{service_fn, ServiceExt};

fn app() -> Router {
  let svc = service_fn(|_req: Request| async { Ok::<_, Infallible>("hello world") });
  let chunks = || async {
    let s = futures::stream::iter([Ok::<_, Infallible>("hello "), Ok("world")]);
    Body::from_stream(s)
  };
  Router::new()
    .route("/get", get(|| async { "hello world" }))
    .route_service("/svc", svc)
    .route("/stream", get(chunks))
    .route("/mw", get(|| async { "hi" }).layer(map_response(longer)))
}

// 中间件在 handler 之后把 body 换成更长的
async fn longer(mut res: Response) -> Response {
  *res.body_mut() = Body::from("hello, middleware");
  res
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let paths = ["/get", "/svc", "/stream", "/mw"];

  println!("== oneshot (no server)");
  for path in paths {
    let req = Request::head(path).body(Body::empty()).unwrap();
    let res = app().oneshot(req).await.unwrap();
    let cl = res.headers().get("content-length").cloned();
    let body = res.into_body().collect().await.unwrap().to_bytes();
    println!("HEAD {path:<8} content-length={cl:?} body={body:?}");
  }

  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() });

  println!("== on the wire");
  for path in paths {
    let mut stream = tokio::net::TcpStream::connect(addr).await.unwrap();
    let head = format!("HEAD {path} HTTP/1.1\r\nHost: x\r\nConnection: close\r\n\r\n");
    stream.write_all(head.as_bytes()).await.unwrap();
    let mut raw = Vec::new();
    stream.read_to_end(&mut raw).await.unwrap();
    let text = String::from_utf8_lossy(&raw);
    // date 每次不同，不打印
    let lines: Vec<&str> = text.split("\r\n").filter(|l| !l.starts_with("date:")).collect();
    println!("HEAD {path}: {} bytes read\n  {lines:?}", raw.len());
  }
}

== oneshot (no server)
HEAD /get     content-length=Some("11") body=b""
HEAD /svc     content-length=None body=b"hello world"
HEAD /stream  content-length=None body=b""
HEAD /mw      content-length=Some("17") body=b""
== on the wire
HEAD /get: 136 bytes read
  ["HTTP/1.1 200 OK", "content-type: text/plain; charset=utf-8", "content-length: 11", "connection: close", "", ""]
HEAD /svc: 136 bytes read
  ["HTTP/1.1 200 OK", "content-type: text/plain; charset=utf-8", "connection: close", "content-length: 11", "", ""]
HEAD /stream: 75 bytes read
  ["HTTP/1.1 200 OK", "connection: close", "", ""]
HEAD /mw: 136 bytes read
  ["HTTP/1.1 200 OK", "content-type: text/plain; charset=utf-8", "content-length: 17", "connection: close", "", ""]

get_service(Svc) + 2 layers        clones=3 calls=1 allocs=16