tonic/service/
router.rs

1use crate::{body::Body, server::NamedService, Status};
2use http::{Request, Response};
3use std::{
4    convert::Infallible,
5    fmt,
6    future::Future,
7    pin::Pin,
8    task::{Context, Poll},
9};
10use tower::{Service, ServiceExt};
11
12/// A [`Service`] router.
13#[derive(Debug, Clone)]
14pub struct Routes {
15    router: axum::Router,
16}
17
18#[derive(Debug, Default, Clone)]
19/// Allows adding new services to routes by passing a mutable reference to this builder.
20pub struct RoutesBuilder {
21    routes: Option<Routes>,
22}
23
24impl RoutesBuilder {
25    /// Add a new service.
26    pub fn add_service<S>(&mut self, svc: S) -> &mut Self
27    where
28        S: Service<Request<Body>, Error = Infallible>
29            + NamedService
30            + Clone
31            + Send
32            + Sync
33            + 'static,
34        S::Response: axum::response::IntoResponse,
35        S::Future: Send + 'static,
36    {
37        let routes = self.routes.take().unwrap_or_default();
38        self.routes.replace(routes.add_service(svc));
39        self
40    }
41
42    /// Returns the routes with added services or empty [`Routes`] if no service was added
43    pub fn routes(self) -> Routes {
44        self.routes.unwrap_or_default()
45    }
46}
47
48impl Default for Routes {
49    fn default() -> Self {
50        Self {
51            router: axum::Router::new().fallback(unimplemented),
52        }
53    }
54}
55
56impl Routes {
57    /// Create a new routes with `svc` already added to it.
58    pub fn new<S>(svc: S) -> Self
59    where
60        S: Service<Request<Body>, Error = Infallible>
61            + NamedService
62            + Clone
63            + Send
64            + Sync
65            + 'static,
66        S::Response: axum::response::IntoResponse,
67        S::Future: Send + 'static,
68    {
69        Self::default().add_service(svc)
70    }
71
72    /// Create a new empty builder.
73    pub fn builder() -> RoutesBuilder {
74        RoutesBuilder::default()
75    }
76
77    /// Add a new service.
78    pub fn add_service<S>(mut self, svc: S) -> Self
79    where
80        S: Service<Request<Body>, Error = Infallible>
81            + NamedService
82            + Clone
83            + Send
84            + Sync
85            + 'static,
86        S::Response: axum::response::IntoResponse,
87        S::Future: Send + 'static,
88    {
89        self.router = self.router.route_service(
90            &format!("/{}/{{*rest}}", S::NAME),
91            svc.map_request(|req: Request<axum::body::Body>| req.map(Body::new)),
92        );
93        self
94    }
95
96    /// This makes axum perform update some internals of the router that improves perf.
97    ///
98    /// See <https://docs.rs/axum/latest/axum/routing/struct.Router.html#a-note-about-performance>
99    pub fn prepare(self) -> Self {
100        Self {
101            router: self.router.with_state(()),
102        }
103    }
104
105    /// Convert this `Routes` into an [`axum::Router`].
106    pub fn into_axum_router(self) -> axum::Router {
107        self.router
108    }
109
110    /// Get a mutable reference to the [`axum::Router`].
111    pub fn axum_router_mut(&mut self) -> &mut axum::Router {
112        &mut self.router
113    }
114}
115
116impl From<Routes> for RoutesBuilder {
117    fn from(routes: Routes) -> Self {
118        Self {
119            routes: Some(routes),
120        }
121    }
122}
123
124impl From<axum::Router> for RoutesBuilder {
125    fn from(router: axum::Router) -> Self {
126        Self {
127            routes: Some(router.into()),
128        }
129    }
130}
131
132impl From<axum::Router> for Routes {
133    fn from(router: axum::Router) -> Self {
134        Self { router }
135    }
136}
137
138async fn unimplemented() -> Response<Body> {
139    let (parts, ()) = Status::unimplemented("").into_http::<()>().into_parts();
140    Response::from_parts(parts, Body::empty())
141}
142
143impl<B> Service<Request<B>> for Routes
144where
145    B: http_body::Body<Data = bytes::Bytes> + Send + 'static,
146    B::Error: Into<crate::BoxError>,
147{
148    type Response = Response<Body>;
149    type Error = Infallible;
150    type Future = RoutesFuture;
151
152    #[inline]
153    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
154        Service::<Request<B>>::poll_ready(&mut self.router, cx)
155    }
156
157    fn call(&mut self, req: Request<B>) -> Self::Future {
158        RoutesFuture(self.router.call(req))
159    }
160}
161
162pub struct RoutesFuture(axum::routing::future::RouteFuture<Infallible>);
163
164impl fmt::Debug for RoutesFuture {
165    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
166        f.debug_tuple("RoutesFuture").finish()
167    }
168}
169
170impl Future for RoutesFuture {
171    type Output = Result<Response<Body>, Infallible>;
172
173    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
174        Pin::new(&mut self.as_mut().0)
175            .poll(cx)
176            .map_ok(|res| res.map(Body::new))
177    }
178}
tonic/service/router.rs

tonic/service/
router.rs
