chore: update to axum 0.6.0 (#1155)

hyperium · Nov 28, 2022 · 542c5b3 · 542c5b3
1 parent b65b52b
commit 542c5b3
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 9 deletions.
diff --git a/tonic/Cargo.toml b/tonic/Cargo.toml
@@ -81,7 +81,7 @@ tokio = {version = "1.0.1", features = ["net"], 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}
 tracing-futures = {version = "0.2", optional = true}
-axum = {version = "0.5.16", default_features = false, optional = true}
+axum = {version = "0.6", default_features = false, optional = true}
 
 # rustls
 rustls-pemfile = { version = "1.0", optional = true }

diff --git a/tonic/src/transport/server/mod.rs b/tonic/src/transport/server/mod.rs
@@ -589,7 +589,7 @@ impl<L> Router<L> {
             .map_err(super::Error::from_source)?;
         self.server
             .serve_with_shutdown::<_, _, future::Ready<()>, _, _, ResBody>(
-                self.routes,
+                self.routes.prepare(),
                 incoming,
                 None,
             )
@@ -618,7 +618,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, incoming, Some(signal))
+            .serve_with_shutdown(self.routes.prepare(), incoming, Some(signal))
             .await
     }
 
@@ -644,7 +644,7 @@ impl<L> Router<L> {
     {
         self.server
             .serve_with_shutdown::<_, _, future::Ready<()>, _, _, ResBody>(
-                self.routes,
+                self.routes.prepare(),
                 incoming,
                 None,
             )
@@ -676,7 +676,7 @@ impl<L> Router<L> {
         ResBody::Error: Into<crate::Error>,
     {
         self.server
-            .serve_with_shutdown(self.routes, incoming, Some(signal))
+            .serve_with_shutdown(self.routes.prepare(), incoming, Some(signal))
             .await
     }
 
@@ -690,7 +690,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)
+        self.server.service_builder.service(self.routes.prepare())
     }
 }
 

diff --git a/tonic/src/transport/service/router.rs b/tonic/src/transport/service/router.rs
@@ -2,7 +2,6 @@ use crate::{
     body::{boxed, BoxBody},
     transport::NamedService,
 };
-use axum::handler::Handler;
 use http::{Request, Response};
 use hyper::Body;
 use pin_project::pin_project;
@@ -33,7 +32,7 @@ impl Routes {
         S::Future: Send + 'static,
         S::Error: Into<crate::Error> + Send,
     {
-        let router = axum::Router::new().fallback(unimplemented.into_service());
+        let router = axum::Router::new().fallback(unimplemented);
         Self { router }.add_service(svc)
     }
 
@@ -48,9 +47,19 @@ impl Routes {
         S::Error: Into<crate::Error> + Send,
     {
         let svc = svc.map_response(|res| res.map(axum::body::boxed));
-        self.router = self.router.route(&format!("/{}/*rest", S::NAME), svc);
+        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(()),
+        }
+    }
 }
 
 async fn unimplemented() -> impl axum::response::IntoResponse {