tonic/transport/service/
router.rs

1use futures_util::{
2    future::Either,
3    future::{MapErr, TryFutureExt},
4};
5use std::{
6    fmt,
7    sync::Arc,
8    task::{Context, Poll},
9};
10use tower_service::Service;
11
12#[doc(hidden)]
13#[derive(Debug)]
14pub struct Routes<A, B, Request> {
15    routes: Or<A, B, Request>,
16}
17
18impl<A, B, Request> Routes<A, B, Request> {
19    pub(crate) fn new(
20        predicate: impl Fn(&Request) -> bool + Send + Sync + 'static,
21        a: A,
22        b: B,
23    ) -> Self {
24        let routes = Or::new(predicate, a, b);
25        Self { routes }
26    }
27}
28
29impl<A, B, Request> Routes<A, B, Request> {
30    pub(crate) fn push<C>(
31        self,
32        predicate: impl Fn(&Request) -> bool + Send + Sync + 'static,
33        route: C,
34    ) -> Routes<C, Or<A, B, Request>, Request> {
35        let routes = Or::new(predicate, route, self.routes);
36        Routes { routes }
37    }
38}
39
40impl<A, B, Request> Service<Request> for Routes<A, B, Request>
41where
42    A: Service<Request>,
43    A::Future: Send + 'static,
44    A::Error: Into<crate::Error>,
45    B: Service<Request, Response = A::Response>,
46    B::Future: Send + 'static,
47    B::Error: Into<crate::Error>,
48{
49    type Response = A::Response;
50    type Error = crate::Error;
51    type Future = <Or<A, B, Request> as Service<Request>>::Future;
52
53    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
54        Ok(()).into()
55    }
56
57    fn call(&mut self, req: Request) -> Self::Future {
58        self.routes.call(req)
59    }
60}
61
62impl<A: Clone, B: Clone, Request> Clone for Routes<A, B, Request> {
63    fn clone(&self) -> Self {
64        Self {
65            routes: self.routes.clone(),
66        }
67    }
68}
69
70#[doc(hidden)]
71pub struct Or<A, B, Request> {
72    predicate: Arc<dyn Fn(&Request) -> bool + Send + Sync + 'static>,
73    a: A,
74    b: B,
75}
76
77impl<A, B, Request> Or<A, B, Request> {
78    pub(crate) fn new<F>(predicate: F, a: A, b: B) -> Self
79    where
80        F: Fn(&Request) -> bool + Send + Sync + 'static,
81    {
82        let predicate = Arc::new(predicate);
83        Self { predicate, a, b }
84    }
85}
86
87impl<A, B, Request> Service<Request> for Or<A, B, Request>
88where
89    A: Service<Request>,
90    A::Future: Send + 'static,
91    A::Error: Into<crate::Error>,
92    B: Service<Request, Response = A::Response>,
93    B::Future: Send + 'static,
94    B::Error: Into<crate::Error>,
95{
96    type Response = A::Response;
97    type Error = crate::Error;
98
99    #[allow(clippy::type_complexity)]
100    type Future = Either<
101        MapErr<A::Future, fn(A::Error) -> crate::Error>,
102        MapErr<B::Future, fn(B::Error) -> crate::Error>,
103    >;
104
105    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
106        Ok(()).into()
107    }
108
109    fn call(&mut self, req: Request) -> Self::Future {
110        if (self.predicate)(&req) {
111            Either::Left(self.a.call(req).map_err(|e| e.into()))
112        } else {
113            Either::Right(self.b.call(req).map_err(|e| e.into()))
114        }
115    }
116}
117
118impl<A: Clone, B: Clone, Request> Clone for Or<A, B, Request> {
119    fn clone(&self) -> Self {
120        Self {
121            predicate: self.predicate.clone(),
122            a: self.a.clone(),
123            b: self.b.clone(),
124        }
125    }
126}
127
128impl<A, B, Request> fmt::Debug for Or<A, B, Request> {
129    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
130        write!(f, "Or {{ .. }}")
131    }
132}