mirror of
https://github.com/tokio-rs/axum.git
synced 2024-10-23 17:36:39 +02:00
Adding layers to the whole thing
This commit is contained in:
David Pedersen
parent
8ee3119fb0
commit
0d7e1e74c4
3 changed files with 69 additions and 31 deletions
22
src/lib.rs
22
src/lib.rs
|
|
@ -41,6 +41,10 @@ pub struct App<R> {
|
|||
}
|
||||
|
||||
impl<R> App<R> {
|
||||
fn new(service_tree: R) -> Self {
|
||||
Self { service_tree }
|
||||
}
|
||||
|
||||
pub fn at(self, route_spec: &str) -> RouteAt<R> {
|
||||
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> {
|
||||
app: App<R>,
|
||||
}
|
||||
|
||||
impl<R> Clone for IntoService<R>
|
||||
where
|
||||
R: Clone,
|
||||
{
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
app: self.app.clone(),
|
||||
}
|
||||
}
|
||||
service_tree: R
|
||||
}
|
||||
|
||||
impl<R, B, T> Service<T> for IntoService<R>
|
||||
|
|
@ -78,14 +72,14 @@ where
|
|||
type Future = R::Future;
|
||||
|
||||
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(())),
|
||||
Err(err) => match err {},
|
||||
}
|
||||
}
|
||||
|
||||
fn call(&mut self, req: T) -> Self::Future {
|
||||
self.app.service_tree.call(req)
|
||||
self.service_tree.call(req)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,8 @@
|
|||
use crate::{
|
||||
body::{Body, BoxBody},
|
||||
handler::{Handler, HandlerSvc},
|
||||
App, IntoService, ResultExt,
|
||||
response::IntoResponse,
|
||||
App, HandleError, IntoService, ResultExt,
|
||||
};
|
||||
use bytes::Bytes;
|
||||
use futures_util::{future, ready};
|
||||
|
|
@ -17,7 +18,7 @@ use std::{
|
|||
use tower::{
|
||||
buffer::{Buffer, BufferLayer},
|
||||
util::BoxService,
|
||||
BoxError, Service, ServiceBuilder,
|
||||
BoxError, Layer, Service, ServiceBuilder,
|
||||
};
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
|
|
@ -152,6 +153,13 @@ where
|
|||
}
|
||||
|
||||
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> {
|
||||
self.app.at(route_spec)
|
||||
}
|
||||
|
|
@ -167,11 +175,32 @@ impl<R> RouteBuilder<R> {
|
|||
define_route_at_methods!(RouteBuilder: trace, trace_service, TRACE);
|
||||
|
||||
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
|
||||
// that way we get to map errors
|
||||
pub fn layer<L>(self, layer: L) -> RouteBuilder<L::Service>
|
||||
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>>
|
||||
where
|
||||
|
|
@ -184,17 +213,11 @@ impl<R> RouteBuilder<R> {
|
|||
.layer(BoxService::layer())
|
||||
.service(self.app.service_tree);
|
||||
|
||||
let app = App {
|
||||
service_tree: BoxServiceTree {
|
||||
inner: svc,
|
||||
poll_ready_error: None,
|
||||
},
|
||||
};
|
||||
|
||||
RouteBuilder {
|
||||
app,
|
||||
route_spec: self.route_spec,
|
||||
}
|
||||
let app = App::new(BoxServiceTree {
|
||||
inner: svc,
|
||||
poll_ready_error: None,
|
||||
});
|
||||
RouteBuilder::new(app, self.route_spec)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
23
src/tests.rs
23
src/tests.rs
|
|
@ -375,7 +375,28 @@ async fn handling_errors_from_layered_single_routes() {
|
|||
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 with one at a "sub path"
|
||||
// TODO(david): composing two boxed apps
|
||||
|
|
|
|||
Loading…
Reference in a new issue