Skip to content

Commit

Permalink
Add RouteDsl::or to combine routes (#108)
Browse files Browse the repository at this point in the history
With this you'll be able to do:

```rust
let one = route("/foo", get(|| async { "foo" }))
    .route("/bar", get(|| async { "bar" }));

let two = route("/baz", get(|| async { "baz" }));

let app = one.or(two);
```

Fixes #101
  • Loading branch information
davidpdrsn authored Aug 7, 2021
1 parent b1e7a6a commit 045ec57
Show file tree
Hide file tree
Showing 6 changed files with 385 additions and 2 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

## Breaking changes

- Add `RoutingDsl::or` for combining routes. ([#108](https://github.com/tokio-rs/axum/pull/108))
- Ensure a `HandleError` service created from `axum::ServiceExt::handle_error`
_does not_ implement `RoutingDsl` as that could lead to confusing routing
behavior. ([#120](https://github.com/tokio-rs/axum/pull/120))
Expand Down
3 changes: 3 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,8 @@
//! Routes can also be dynamic like `/users/:id`. See [extractors](#extractors)
//! for more details.
//!
//! You can also define routes separately and merge them with [`RoutingDsl::or`].
//!
//! ## Precedence
//!
//! Note that routes are matched _bottom to top_ so routes that should have
Expand Down Expand Up @@ -662,6 +664,7 @@
//! [`IntoResponse`]: crate::response::IntoResponse
//! [`Timeout`]: tower::timeout::Timeout
//! [examples]: https://github.com/tokio-rs/axum/tree/main/examples
//! [`RoutingDsl::or`]: crate::routing::RoutingDsl::or
//! [`axum::Server`]: hyper::server::Server
#![warn(
Expand Down
53 changes: 51 additions & 2 deletions src/routing.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ use tower::{
use tower_http::map_response_body::MapResponseBodyLayer;

pub mod future;
pub mod or;

/// A filter that matches one or more HTTP methods.
#[derive(Debug, Copy, Clone)]
Expand Down Expand Up @@ -354,6 +355,40 @@ pub trait RoutingDsl: crate::sealed::Sealed + Sized {
{
IntoMakeServiceWithConnectInfo::new(self)
}

/// Merge two routers into one.
///
/// This is useful for breaking apps into smaller pieces and combining them
/// into one.
///
/// ```
/// use axum::prelude::*;
/// #
/// # async fn users_list() {}
/// # async fn users_show() {}
/// # async fn teams_list() {}
///
/// // define some routes separately
/// let user_routes = route("/users", get(users_list))
/// .route("/users/:id", get(users_show));
///
/// let team_routes = route("/teams", get(teams_list));
///
/// // combine them into one
/// let app = user_routes.or(team_routes);
/// # async {
/// # hyper::Server::bind(&"".parse().unwrap()).serve(app.into_make_service()).await.unwrap();
/// # };
/// ```
fn or<S>(self, other: S) -> or::Or<Self, S>
where
S: RoutingDsl,
{
or::Or {
first: self,
second: other,
}
}
}

impl<S, F> RoutingDsl for Route<S, F> {}
Expand Down Expand Up @@ -448,7 +483,10 @@ impl<E> RoutingDsl for EmptyRouter<E> {}

impl<E> crate::sealed::Sealed for EmptyRouter<E> {}

impl<B, E> Service<Request<B>> for EmptyRouter<E> {
impl<B, E> Service<Request<B>> for EmptyRouter<E>
where
B: Send + Sync + 'static,
{
type Response = Response<BoxBody>;
type Error = E;
type Future = EmptyRouterFuture<E>;
Expand All @@ -457,15 +495,26 @@ impl<B, E> Service<Request<B>> for EmptyRouter<E> {
Poll::Ready(Ok(()))
}

fn call(&mut self, _req: Request<B>) -> Self::Future {
fn call(&mut self, request: Request<B>) -> Self::Future {
let mut res = Response::new(crate::body::empty());
res.extensions_mut().insert(FromEmptyRouter { request });
*res.status_mut() = self.status;
EmptyRouterFuture {
future: futures_util::future::ok(res),
}
}
}

/// Response extension used by [`EmptyRouter`] to send the request back to [`Or`] so
/// the other service can be called.
///
/// Without this we would loose ownership of the request when calling the first
/// service in [`Or`]. We also wouldn't be able to identify if the response came
/// from [`EmptyRouter`] and therefore can be discarded in [`Or`].
struct FromEmptyRouter<B> {
request: Request<B>,
}

#[derive(Debug, Clone)]
pub(crate) struct PathPattern(Arc<Inner>);

Expand Down
124 changes: 124 additions & 0 deletions src/routing/or.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,124 @@
//! [`Or`] used to combine two services into one.
use super::{FromEmptyRouter, RoutingDsl};
use crate::body::BoxBody;
use futures_util::ready;
use http::{Request, Response};
use pin_project_lite::pin_project;
use std::{
future::Future,
pin::Pin,
task::{Context, Poll},
};
use tower::{util::Oneshot, Service, ServiceExt};

/// [`tower::Service`] that is the combination of two routers.
///
/// See [`RoutingDsl::or`] for more details.
///
/// [`RoutingDsl::or`]: super::RoutingDsl::or
#[derive(Debug, Clone, Copy)]
pub struct Or<A, B> {
pub(super) first: A,
pub(super) second: B,
}

impl<A, B> RoutingDsl for Or<A, B> {}

impl<A, B> crate::sealed::Sealed for Or<A, B> {}

#[allow(warnings)]
impl<A, B, ReqBody> Service<Request<ReqBody>> for Or<A, B>
where
A: Service<Request<ReqBody>, Response = Response<BoxBody>> + Clone,
B: Service<Request<ReqBody>, Response = Response<BoxBody>, Error = A::Error> + Clone,
ReqBody: Send + Sync + 'static,
A: Send + 'static,
B: Send + 'static,
A::Future: Send + 'static,
B::Future: Send + 'static,
{
type Response = Response<BoxBody>;
type Error = A::Error;
type Future = ResponseFuture<A, B, ReqBody>;

fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}

fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
ResponseFuture {
state: State::FirstFuture {
f: self.first.clone().oneshot(req),
},
second: Some(self.second.clone()),
}
}
}

pin_project! {
/// Response future for [`Or`].
pub struct ResponseFuture<A, B, ReqBody>
where
A: Service<Request<ReqBody>>,
B: Service<Request<ReqBody>>,
{
#[pin]
state: State<A, B, ReqBody>,
second: Option<B>,
}
}

pin_project! {
#[project = StateProj]
enum State<A, B, ReqBody>
where
A: Service<Request<ReqBody>>,
B: Service<Request<ReqBody>>,
{
FirstFuture { #[pin] f: Oneshot<A, Request<ReqBody>> },
SecondFuture {
#[pin]
f: Oneshot<B, Request<ReqBody>>,
}
}
}

impl<A, B, ReqBody> Future for ResponseFuture<A, B, ReqBody>
where
A: Service<Request<ReqBody>, Response = Response<BoxBody>>,
B: Service<Request<ReqBody>, Response = Response<BoxBody>, Error = A::Error>,
ReqBody: Send + Sync + 'static,
{
type Output = Result<Response<BoxBody>, A::Error>;

fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
loop {
let mut this = self.as_mut().project();

let new_state = match this.state.as_mut().project() {
StateProj::FirstFuture { f } => {
let mut response = ready!(f.poll(cx)?);

let req = if let Some(ext) = response
.extensions_mut()
.remove::<FromEmptyRouter<ReqBody>>()
{
ext.request
} else {
return Poll::Ready(Ok(response));
};

let second = this.second.take().expect("future polled after completion");

State::SecondFuture {
f: second.oneshot(req),
}
}
StateProj::SecondFuture { f } => return f.poll(cx),
};

this.state.set(new_state);
}
}
}
3 changes: 3 additions & 0 deletions src/tests.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
#![allow(clippy::blacklisted_name)]

use crate::{
extract::RequestParts, handler::on, prelude::*, routing::nest, routing::MethodFilter, service,
};
Expand All @@ -18,6 +20,7 @@ use tower::{make::Shared, service_fn, BoxError, Service, ServiceBuilder};
use tower_http::{compression::CompressionLayer, trace::TraceLayer};

mod nest;
mod or;

#[tokio::test]
async fn hello_world() {
Expand Down
Loading

0 comments on commit 045ec57

Please sign in to comment.