itsscb/axum
1
0
Fork 0
mirror of https://github.com/tokio-rs/axum.git synced 2025-11-04 07:23:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

Adding layers to the whole thing

Browse Source
This commit is contained in:
David Pedersen 2021-06-01 21:15:48 +02:00
parent 8ee3119fb0
commit 0d7e1e74c4
3 changed files with 69 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
src/lib.rs
View File

@ -41,6 +41,10 @@ pub struct App<R> {
} }
impl<R> App<R> { impl<R> App<R> {
fn new(service_tree: R) -> Self {
Self { service_tree }
}
pub fn at(self, route_spec: &str) -> RouteAt<R> { pub fn at(self, route_spec: &str) -> RouteAt<R> {
self.at_bytes(Bytes::copy_from_slice(route_spec.as_bytes())) self.at_bytes(Bytes::copy_from_slice(route_spec.as_bytes()))
} }
@ -53,19 +57,9 @@ impl<R> App<R> {
} }
} }
#[derive(Clone)]
pub struct IntoService<R> { pub struct IntoService<R> {
app: App<R>, service_tree: R
}
impl<R> Clone for IntoService<R>
where
R: Clone,
{
fn clone(&self) -> Self {
Self {
app: self.app.clone(),
}
}
} }
impl<R, B, T> Service<T> for IntoService<R> impl<R, B, T> Service<T> for IntoService<R>
@ -78,14 +72,14 @@ where
type Future = R::Future; type Future = R::Future;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
match ready!(self.app.service_tree.poll_ready(cx)) { match ready!(self.service_tree.poll_ready(cx)) {
Ok(_) => Poll::Ready(Ok(())), Ok(_) => Poll::Ready(Ok(())),
Err(err) => match err {}, Err(err) => match err {},
} }
} }
fn call(&mut self, req: T) -> Self::Future { fn call(&mut self, req: T) -> Self::Future {
self.app.service_tree.call(req) self.service_tree.call(req)
} }
} }

51
src/routing.rs
View File

@ -1,7 +1,8 @@
use crate::{ use crate::{
body::{Body, BoxBody}, body::{Body, BoxBody},
handler::{Handler, HandlerSvc}, handler::{Handler, HandlerSvc},
App, IntoService, ResultExt, response::IntoResponse,
App, HandleError, IntoService, ResultExt,
}; };
use bytes::Bytes; use bytes::Bytes;
use futures_util::{future, ready}; use futures_util::{future, ready};
@ -17,7 +18,7 @@ use std::{
use tower::{ use tower::{
buffer::{Buffer, BufferLayer}, buffer::{Buffer, BufferLayer},
util::BoxService, util::BoxService,
BoxError, Service, ServiceBuilder, BoxError, Layer, Service, ServiceBuilder,
}; };
#[derive(Clone, Copy)] #[derive(Clone, Copy)]
@ -152,6 +153,13 @@ where
} }
impl<R> RouteBuilder<R> { impl<R> RouteBuilder<R> {
fn new(app: App<R>, route_spec: impl Into<Bytes>) -> Self {
Self {
app,
route_spec: route_spec.into(),
}
}
pub fn at(self, route_spec: &str) -> RouteAt<R> { pub fn at(self, route_spec: &str) -> RouteAt<R> {
self.app.at(route_spec) self.app.at(route_spec)
} }
@ -167,11 +175,32 @@ impl<R> RouteBuilder<R> {
define_route_at_methods!(RouteBuilder: trace, trace_service, TRACE); define_route_at_methods!(RouteBuilder: trace, trace_service, TRACE);
pub fn into_service(self) -> IntoService<R> { pub fn into_service(self) -> IntoService<R> {
IntoService { app: self.app } IntoService {
service_tree: self.app.service_tree,
}
} }
// TODO(david): Add `layer` method here that applies a `tower::Layer` inside the service tree pub fn layer<L>(self, layer: L) -> RouteBuilder<L::Service>
// that way we get to map errors where
L: Layer<R>,
{
let layered = layer.layer(self.app.service_tree);
let app = App::new(layered);
RouteBuilder::new(app, self.route_spec)
}
pub fn handle_error<F, B, Res>(self, f: F) -> RouteBuilder<HandleError<R, F, R::Error>>
where
R: Service<Request<Body>, Response = Response<B>>,
F: FnOnce(R::Error) -> Res,
Res: IntoResponse<Body>,
B: http_body::Body<Data = Bytes> + Send + Sync + 'static,
B::Error: Into<BoxError> + Send + Sync + 'static,
{
let svc = HandleError::new(self.app.service_tree, f);
let app = App::new(svc);
RouteBuilder::new(app, self.route_spec)
}
pub fn boxed<B>(self) -> RouteBuilder<BoxServiceTree<B>> pub fn boxed<B>(self) -> RouteBuilder<BoxServiceTree<B>>
where where
@ -184,17 +213,11 @@ impl<R> RouteBuilder<R> {
.layer(BoxService::layer()) .layer(BoxService::layer())
.service(self.app.service_tree); .service(self.app.service_tree);
let app = App { let app = App::new(BoxServiceTree {
service_tree: BoxServiceTree {
inner: svc, inner: svc,
poll_ready_error: None, poll_ready_error: None,
}, });
}; RouteBuilder::new(app, self.route_spec)
RouteBuilder {
app,
route_spec: self.route_spec,
}
} }
} }

23
src/tests.rs
View File

@ -375,7 +375,28 @@ async fn handling_errors_from_layered_single_routes() {
assert_eq!(res.status(), StatusCode::INTERNAL_SERVER_ERROR); assert_eq!(res.status(), StatusCode::INTERNAL_SERVER_ERROR);
} }
// TODO(david): .layer() on RouteBuilder #[tokio::test]
async fn layer_on_whole_router() {
use tower::timeout::TimeoutLayer;
async fn handle(_req: Request<Body>) -> &'static str {
tokio::time::sleep(Duration::from_secs(10)).await;
""
}
let app = app()
.at("/")
.get(handle)
.layer(TimeoutLayer::new(Duration::from_millis(100)))
.handle_error(|_err: BoxError| StatusCode::INTERNAL_SERVER_ERROR)
.into_service();
let addr = run_in_background(app).await;
let res = reqwest::get(format!("http://{}", addr)).await.unwrap();
assert_eq!(res.status(), StatusCode::INTERNAL_SERVER_ERROR);
}
// TODO(david): composing two apps // TODO(david): composing two apps
// TODO(david): composing two apps with one at a "sub path" // TODO(david): composing two apps with one at a "sub path"
// TODO(david): composing two boxed apps // TODO(david): composing two boxed apps