feat(transport): Add optional service methods (#275)

2020-07-10 12:39:52 -04:00
parent 034c850254
commit 2b997b0c5f
3 changed files with 127 additions and 1 deletions
@@ -41,7 +41,8 @@ use std::{
 };
 use tokio::io::{AsyncRead, AsyncWrite};
 use tower::{
-    limit::concurrency::ConcurrencyLimitLayer, timeout::TimeoutLayer, Service, ServiceBuilder,
+    limit::concurrency::ConcurrencyLimitLayer, timeout::TimeoutLayer, util::Either, Service,
+    ServiceBuilder,
 };
 use tracing_futures::{Instrument, Instrumented};

@@ -86,6 +87,10 @@ pub trait NamedService {
    const NAME: &'static str;
 }

+impl<S: NamedService, T> NamedService for Either<S, T> {
+    const NAME: &'static str = S::NAME;
+}
+
 impl Server {
    /// Create a new server builder that can configure a [`Server`].
    pub fn builder() -> Self {
@@ -231,6 +236,34 @@ impl Server {
        Router::new(self.clone(), svc)
    }

+    /// Create a router with the optional `S` typed service as the first service.
+    ///
+    /// This will clone the `Server` builder and create a router that will
+    /// route around different services.
+    ///
+    /// # Note
+    /// Even when the argument given is `None` this will capture *all* requests to this service name.
+    /// As a result, one cannot use this to toggle between two indentically named implementations.
+    pub fn add_optional_service<S>(
+        &mut self,
+        svc: Option<S>,
+    ) -> Router<Either<S, Unimplemented>, Unimplemented>
+    where
+        S: Service<Request<Body>, Response = Response<BoxBody>>
+            + NamedService
+            + Clone
+            + Send
+            + 'static,
+        S::Future: Send + 'static,
+        S::Error: Into<crate::Error> + Send,
+    {
+        let svc = match svc {
+            Some(some) => Either::A(some),
+            None => Either::B(Unimplemented::default()),
+        };
+        Router::new(self.clone(), svc)
+    }
+
    pub(crate) async fn serve_with_shutdown<S, I, F, IO, IE>(
        self,
        svc: S,
@@ -342,6 +375,42 @@ where
        Router { server, routes }
    }

+    /// Add a new optional service to this router.
+    ///
+    /// # Note
+    /// Even when the argument given is `None` this will capture *all* requests to this service name.
+    /// As a result, one cannot use this to toggle between two indentically named implementations.
+    pub fn add_optional_service<S>(
+        self,
+        svc: Option<S>,
+    ) -> Router<Either<S, Unimplemented>, Or<A, B, Request<Body>>>
+    where
+        S: Service<Request<Body>, Response = Response<BoxBody>>
+            + NamedService
+            + Clone
+            + Send
+            + 'static,
+        S::Future: Send + 'static,
+        S::Error: Into<crate::Error> + Send,
+    {
+        let Self { routes, server } = self;
+
+        let svc_name = <S as NamedService>::NAME;
+        let svc_route = format!("/{}", svc_name);
+        let pred = move |req: &Request<Body>| {
+            let path = req.uri().path();
+
+            path.starts_with(&svc_route)
+        };
+        let svc = match svc {
+            Some(some) => Either::A(some),
+            None => Either::B(Unimplemented::default()),
+        };
+        let routes = routes.push(pred, svc);
+
+        Router { server, routes }
+    }
+
    /// Consume this [`Server`] creating a future that will execute the server
    /// on [`tokio`]'s default executor.
    ///