Skip to main content

satex_server/router/
mod.rs

1mod route;
2
3use futures::{Stream, StreamExt};
4pub use route::Builder as RouteBuilder;
5pub use route::Route;
6
7use bytes::Bytes;
8use futures::future::{BoxFuture, LocalBoxFuture};
9use http::request::Parts;
10use http::{Request, Response, StatusCode};
11use hyper::service::Service as HyperService;
12use satex_core::body::Body;
13use satex_core::extension::{RawUri, RouteId};
14use satex_core::util::ResponseExt;
15use satex_core::{BoxError, Error};
16use satex_matcher::RouteMatcher;
17use std::future::{poll_fn, ready, Ready};
18use std::pin::pin;
19use std::sync::Arc;
20use tokio::sync::RwLock;
21use tower::Service;
22use tracing::{debug, info};
23
24#[derive(Clone, Default)]
25pub struct Router {
26    routes: Vec<Route>,
27}
28
29impl Router {
30    pub fn new(routes: Vec<Route>) -> Self {
31        Self { routes }
32    }
33
34    pub fn into_static_service(self) -> MakeRouterService {
35        MakeRouterService::new(InternalRouter::Static(Arc::from(self.routes)))
36    }
37
38    pub fn into_dynamic_service<S, F>(self, events: S, f: F) -> MakeRouterService
39    where
40        S: Stream<Item=Event> + Send + 'static,
41        F: FnOnce(BoxFuture<'static, ()>),
42    {
43        let routes = Arc::new(RwLock::new(self.routes));
44        let make_service = MakeRouterService::new(InternalRouter::Dynamic(routes.clone()));
45
46        // 更新路由任务
47        let future = async move {
48            let mut events = pin!(events);
49            while let Some(event) = events.next().await {
50                match event {
51                    Event::Set(router) => {
52                        info!("refresh routes: {}", router.routes.len());
53                        *routes.write().await = router.routes;
54                    }
55                    Event::Clear => {
56                        routes.write().await.clear();
57                    }
58                }
59            }
60        };
61        f(Box::pin(future));
62
63        make_service
64    }
65}
66
67#[derive(Clone)]
68pub enum Event {
69    Set(Router),
70    Clear,
71}
72
73#[doc(hidden)]
74#[derive(Clone)]
75pub enum InternalRouter {
76    Static(Arc<[Route]>),
77    Dynamic(Arc<RwLock<Vec<Route>>>),
78}
79
80impl<ReqBody> HyperService<Request<ReqBody>> for InternalRouter
81where
82    ReqBody: http_body::Body<Data=Bytes> + Send + 'static,
83    ReqBody::Error: Into<BoxError>,
84{
85    type Response = Response<Body>;
86    type Error = Error;
87    type Future = LocalBoxFuture<'static, Result<Self::Response, Self::Error>>;
88
89    fn call(&self, request: Request<ReqBody>) -> Self::Future {
90        let router = self.clone();
91        Box::pin(async move {
92            let (mut parts, body) = request.into_parts();
93
94            // 如果是动态路由, 找到匹配的路由后就会释放读锁
95            let route = match &router {
96                InternalRouter::Static(routes) => find_route(routes, &mut parts).await,
97                InternalRouter::Dynamic(routes) => {
98                    find_route(routes.read().await.as_slice(), &mut parts).await
99                }
100            };
101
102            match route {
103                Ok(Some(mut route)) => {
104                    match poll_fn(|ctx| {
105                        <Route as Service<Request<ReqBody>>>::poll_ready(&mut route, ctx)
106                    })
107                        .await
108                    {
109                        Ok(_) => {
110                            // raw uri
111                            parts.extensions.insert(RawUri::new(parts.uri.clone()));
112                            // route id
113                            parts.extensions.insert(RouteId::new(route.id()));
114                            // poll ready
115                            poll_fn(|ctx| {
116                                <Route as Service<Request<ReqBody>>>::poll_ready(&mut route, ctx)
117                            })
118                                .await
119                                .map_err(Error::new)?;
120
121                            // call route
122                            route.call(Request::from_parts(parts, body)).await
123                        }
124                        Err(e) => Err(Error::new(e)),
125                    }
126                }
127                Ok(None) => Ok(Response::new(Body::empty()).with_status(StatusCode::NOT_FOUND)),
128                Err(e) => Err(e),
129            }
130        })
131    }
132}
133
134#[inline(always)]
135async fn find_route(routes: &[Route], parts: &mut Parts) -> Result<Option<Route>, Error> {
136    for route in routes {
137        match route.matches(parts).await {
138            Err(e) => {
139                return Err(e);
140            }
141            Ok(false) => {
142                debug!("Not matched route: {}", route.id());
143            }
144            Ok(true) => {
145                debug!("Matched route: {}", route.id());
146                return Ok(Some(route.clone()));
147            }
148        }
149    }
150    debug!("Not find any matched route!");
151    Ok(None)
152}
153
154#[derive(Clone)]
155pub struct MakeRouterService(InternalRouter);
156
157impl MakeRouterService {
158    fn new(router: InternalRouter) -> Self {
159        Self(router)
160    }
161}
162
163impl HyperService<()> for MakeRouterService {
164    type Response = InternalRouter;
165    type Error = ();
166    type Future = Ready<Result<Self::Response, Self::Error>>;
167
168    fn call(&self, _: ()) -> Self::Future {
169        ready(Ok(self.0.clone()))
170    }
171}