Add RouteDsl::or to combine routes

src/lib.rs

@@ -128,6 +128,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
@@ -671,6 +673,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
 
 #![doc(html_root_url = "https://docs.rs/axum/0.1.2")]
 #![warn(

src/routing.rs

@@ -29,6 +29,8 @@ use tower::{
 };
 use tower_http::map_response_body::MapResponseBodyLayer;
 
+pub mod or;
+
 /// A filter that matches one or more HTTP methods.
 #[derive(Debug, Copy, Clone)]
 pub enum MethodFilter {
@@ -354,6 +356,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> {}
@@ -507,7 +543,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>;
@@ -516,8 +555,9 @@ 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: future::ok(res),
@@ -531,6 +571,16 @@ opaque_future! {
         future::Ready<Result<Response<BoxBody>, E>>;
 }
 
+/// 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 response 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>);
 

src/routing/or.rs

@@ -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);
+        }
+    }
+}

src/tests.rs

@@ -1,3 +1,5 @@
+#![allow(clippy::blacklisted_name)]
+
 use crate::{
     extract::RequestParts, handler::on, prelude::*, response::IntoResponse, routing::nest,
     routing::MethodFilter, service,
@@ -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() {