chore(transport): Replace axum with matchit

hyperium · Oct 30, 2023 · ea7f8ac · ea7f8ac
1 parent b3fca19
commit ea7f8ac
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 67 deletions.
diff --git a/examples/src/h2c/server.rs b/examples/src/h2c/server.rs
@@ -61,7 +61,7 @@ mod h2c {
 
     impl<S> Service<Request<Body>> for H2c<S>
     where
-        S: Service<Request<Body>, Response = Response<tonic::transport::AxumBoxBody>>
+        S: Service<Request<Body>, Response = Response<tonic::body::BoxBody>>
             + Clone
             + Send
             + 'static,

diff --git a/tonic/Cargo.toml b/tonic/Cargo.toml
@@ -33,7 +33,7 @@ tls-roots-common = ["tls"]
 tls-webpki-roots = ["tls-roots-common", "dep:webpki-roots"]
 transport = [
   "dep:async-stream",
-  "dep:axum",
+  "dep:matchit",
   "channel",
   "dep:h2",
   "dep:hyper",
@@ -73,7 +73,7 @@ hyper = {version = "0.14.26", features = ["full"], optional = true}
 hyper-timeout = {version = "0.4", optional = true}
 tokio-stream = "0.1"
 tower = {version = "0.4.7", default-features = false, features = ["balance", "buffer", "discover", "limit", "load", "make", "timeout", "util"], optional = true}
-axum = {version = "0.6.9", default_features = false, optional = true}
+matchit = {version="0.7.3", optional = true}
 
 # rustls
 async-stream = { version = "0.3", optional = true }

diff --git a/tonic/src/transport/mod.rs b/tonic/src/transport/mod.rs
@@ -106,7 +106,6 @@ pub use self::service::grpc_timeout::TimeoutExpired;
 pub use self::tls::Certificate;
 #[doc(inline)]
 pub use crate::server::NamedService;
-pub use axum::{body::BoxBody as AxumBoxBody, Router as AxumRouter};
 pub use hyper::{Body, Uri};
 
 pub(crate) use self::service::executor::Executor;

diff --git a/tonic/src/transport/server/mod.rs b/tonic/src/transport/server/mod.rs
@@ -601,9 +601,9 @@ impl<L> Router<L> {
         self
     }
 
-    /// Convert this tonic `Router` into an axum `Router` consuming the tonic one.
-    pub fn into_router(self) -> axum::Router {
-        self.routes.into_router()
+    /// Convert this tonic `Router` into [`Routes`] consuming the tonic one.
+    pub fn into_router(self) -> Routes {
+        self.routes
     }
 
     /// Consume this [`Server`] creating a future that will execute the server
@@ -624,7 +624,7 @@ impl<L> Router<L> {
             .map_err(super::Error::from_source)?;
         self.server
             .serve_with_shutdown::<_, _, future::Ready<()>, _, _, ResBody>(
-                self.routes.prepare(),
+                self.routes,
                 incoming,
                 None,
             )
@@ -653,7 +653,7 @@ impl<L> Router<L> {
         let incoming = TcpIncoming::new(addr, self.server.tcp_nodelay, self.server.tcp_keepalive)
             .map_err(super::Error::from_source)?;
         self.server
-            .serve_with_shutdown(self.routes.prepare(), incoming, Some(signal))
+            .serve_with_shutdown(self.routes, incoming, Some(signal))
             .await
     }
 
@@ -681,7 +681,7 @@ impl<L> Router<L> {
     {
         self.server
             .serve_with_shutdown::<_, _, future::Ready<()>, _, _, ResBody>(
-                self.routes.prepare(),
+                self.routes,
                 incoming,
                 None,
             )
@@ -715,7 +715,7 @@ impl<L> Router<L> {
         ResBody::Error: Into<crate::Error>,
     {
         self.server
-            .serve_with_shutdown(self.routes.prepare(), incoming, Some(signal))
+            .serve_with_shutdown(self.routes, incoming, Some(signal))
             .await
     }
 
@@ -729,7 +729,7 @@ impl<L> Router<L> {
         ResBody: http_body::Body<Data = Bytes> + Send + 'static,
         ResBody::Error: Into<crate::Error>,
     {
-        self.server.service_builder.service(self.routes.prepare())
+        self.server.service_builder.service(self.routes)
     }
 }
 

diff --git a/tonic/src/transport/service/router.rs b/tonic/src/transport/service/router.rs
@@ -1,24 +1,24 @@
-use crate::{
-    body::{boxed, BoxBody},
-    server::NamedService,
-};
+use crate::{body::BoxBody, server::NamedService, transport::BoxFuture};
 use http::{Request, Response};
 use hyper::Body;
-use pin_project::pin_project;
 use std::{
     convert::Infallible,
     fmt,
-    future::Future,
-    pin::Pin,
-    task::{ready, Context, Poll},
+    task::{Context, Poll},
 };
-use tower::ServiceExt;
+use tower::{util::BoxCloneService, ServiceExt};
 use tower_service::Service;
 
 /// A [`Service`] router.
-#[derive(Debug, Default, Clone)]
+#[derive(Default, Clone)]
 pub struct Routes {
-    router: axum::Router,
+    router: matchit::Router<BoxCloneService<Request<Body>, Response<BoxBody>, crate::Error>>,
+}
+
+impl fmt::Debug for Routes {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Routes").finish()
+    }
 }
 
 #[derive(Debug, Default, Clone)]
@@ -49,6 +49,7 @@ impl RoutesBuilder {
         self.routes.unwrap_or_default()
     }
 }
+
 impl Routes {
     /// Create a new routes with `svc` already added to it.
     pub fn new<S>(svc: S) -> Self
@@ -61,8 +62,7 @@ impl Routes {
         S::Future: Send + 'static,
         S::Error: Into<crate::Error> + Send,
     {
-        let router = axum::Router::new().fallback(unimplemented);
-        Self { router }.add_service(svc)
+        Self::default().add_service(svc)
     }
 
     /// Add a new service.
@@ -76,64 +76,38 @@ impl Routes {
         S::Future: Send + 'static,
         S::Error: Into<crate::Error> + Send,
     {
-        let svc = svc.map_response(|res| res.map(axum::body::boxed));
-        self.router = self
-            .router
-            .route_service(&format!("/{}/*rest", S::NAME), svc);
-        self
-    }
-
-    pub(crate) fn prepare(self) -> Self {
-        Self {
-            // this makes axum perform update some internals of the router that improves perf
-            // see https://docs.rs/axum/latest/axum/routing/struct.Router.html#a-note-about-performance
-            router: self.router.with_state(()),
-        }
-    }
-
-    /// Convert this `Routes` into an [`axum::Router`].
-    pub fn into_router(self) -> axum::Router {
+        let svc = svc.map_err(Into::into);
         self.router
+            .insert(format!("/{}/*rest", S::NAME), BoxCloneService::new(svc))
+            .unwrap_or_else(|e| panic!("failed to configurate routing: {e}"));
+        self
     }
 }
 
-async fn unimplemented() -> impl axum::response::IntoResponse {
-    let status = http::StatusCode::OK;
-    let headers = [("grpc-status", "12"), ("content-type", "application/grpc")];
-    (status, headers)
+async fn unimplemented() -> Result<Response<BoxBody>, crate::Error> {
+    let response = Response::builder()
+        .status(http::StatusCode::OK)
+        .header("grpc-status", "12")
+        .header("content-type", "application/grpc")
+        .body(crate::body::empty_body())
+        .unwrap();
+    Ok(response)
 }
 
 impl Service<Request<Body>> for Routes {
     type Response = Response<BoxBody>;
     type Error = crate::Error;
-    type Future = RoutesFuture;
+    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
 
     #[inline]
     fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         Poll::Ready(Ok(()))
     }
 
     fn call(&mut self, req: Request<Body>) -> Self::Future {
-        RoutesFuture(self.router.call(req))
-    }
-}
-
-#[pin_project]
-pub struct RoutesFuture(#[pin] axum::routing::future::RouteFuture<Body, Infallible>);
-
-impl fmt::Debug for RoutesFuture {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.debug_tuple("RoutesFuture").finish()
-    }
-}
-
-impl Future for RoutesFuture {
-    type Output = Result<Response<BoxBody>, crate::Error>;
-
-    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        match ready!(self.project().0.poll(cx)) {
-            Ok(res) => Ok(res.map(boxed)).into(),
-            Err(err) => match err {},
+        match self.router.at_mut(req.uri().path()) {
+            Ok(found) => found.value.call(req),
+            Err(_) => Box::pin(unimplemented()),
         }
     }
 }