feat(transport): Add Router::into_service (#419)

Co-authored-by: Danny Hua <[email protected]>
2020-08-20 15:11:31 -04:00
parent 90858926b6
commit 37f6733f85
4 changed files with 292 additions and 1 deletions
@@ -25,7 +25,7 @@ use super::service::{Or, Routes, ServerIo, ServiceBuilderExt};
 use crate::{body::BoxBody, request::ConnectionInfo};
 use futures_core::Stream;
 use futures_util::{
-    future::{self, MapErr},
+    future::{self, Either as FutureEither, MapErr},
    TryFutureExt,
 };
 use http::{HeaderMap, Request, Response};
@@ -78,6 +78,42 @@ pub struct Router<A, B> {
    routes: Routes<A, B, Request<Body>>,
 }

+/// A service that is produced from a Tonic `Router`.
+///
+/// This service implementation will route between multiple Tonic
+/// gRPC endpoints and can be consumed with the rest of the `tower`
+/// ecosystem.
+#[derive(Debug)]
+pub struct RouterService<A, B> {
+    router: Router<A, B>,
+}
+
+impl<A, B> Service<Request<Body>> for RouterService<A, B>
+where
+    A: Service<Request<Body>, Response = Response<BoxBody>> + Clone + Send + 'static,
+    A::Future: Send + 'static,
+    A::Error: Into<crate::Error> + Send,
+    B: Service<Request<Body>, Response = Response<BoxBody>> + Clone + Send + 'static,
+    B::Future: Send + 'static,
+    B::Error: Into<crate::Error> + Send,
+{
+    type Response = Response<BoxBody>;
+    type Future = FutureEither<
+        MapErr<A::Future, fn(A::Error) -> crate::Error>,
+        MapErr<B::Future, fn(B::Error) -> crate::Error>,
+    >;
+    type Error = crate::Error;
+
+    fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        Poll::Ready(Ok(()))
+    }
+
+    #[inline]
+    fn call(&mut self, req: Request<Body>) -> Self::Future {
+        self.router.routes.call(req)
+    }
+}
+
 /// A trait to provide a static reference to the service's
 /// name. This is used for routing service's within the router.
 pub trait NamedService {
@@ -476,6 +512,11 @@ where
            .serve_with_shutdown(self.routes, incoming, Some(signal))
            .await
    }
+
+    /// Create a tower service out of a router.
+    pub fn into_service(self) -> RouterService<A, B> {
+        RouterService { router: self }
+    }
 }

 impl fmt::Debug for Server {