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

hyperium · Jul 10, 2020 · 2b997b0 · 2b997b0
1 parent 034c850
commit 2b997b0
Show file tree

Hide file tree

Showing 3 changed files with 127 additions and 1 deletion.
diff --git a/examples/Cargo.toml b/examples/Cargo.toml
@@ -122,6 +122,10 @@ path = "src/health/server.rs"
 name = "autoreload-server"
 path = "src/autoreload/server.rs"
 
+[[bin]]
+name = "optional-server"
+path = "src/optional/server.rs"
+
 [dependencies]
 tonic = { path = "../tonic", features = ["tls"] }
 prost = "0.6"

diff --git a/examples/src/optional/server.rs b/examples/src/optional/server.rs
@@ -0,0 +1,53 @@
+use std::env;
+use tonic::{transport::Server, Request, Response, Status};
+
+use hello_world::greeter_server::{Greeter, GreeterServer};
+use hello_world::{HelloReply, HelloRequest};
+
+pub mod hello_world {
+ tonic::include_proto!("helloworld");
+}
+
+#[derive(Default)]
+pub struct MyGreeter {}
+
+#[tonic::async_trait]
+impl Greeter for MyGreeter {
+ async fn say_hello(
+ &self,
+ request: Request<HelloRequest>,
+ ) -> Result<Response<HelloReply>, Status> {
+ println!("Got a request from {:?}", request.remote_addr());
+
+ let reply = hello_world::HelloReply {
+ message: format!("Hello {}!", request.into_inner().name),
+ };
+ Ok(Response::new(reply))
+ }
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+ let args: Vec<String> = env::args().collect();
+ let enabled = args.get(1) == Some(&"enable".to_string());
+
+ let addr = "[::1]:50051".parse().unwrap();
+ let greeter = MyGreeter::default();
+
+ let optional_service = if enabled {
+ println!("MyGreeter enabled");
+ Some(GreeterServer::new(greeter))
+ } else {
+ println!("MyGreeter disabled");
+ None
+ };
+
+ println!("GreeterServer listening on {}", addr);
+
+ Server::builder()
+ .add_optional_service(optional_service)
+ .serve(addr)
+ .await?;
+
+ Ok(())
+}
diff --git a/tonic/src/transport/server/mod.rs b/tonic/src/transport/server/mod.rs
@@ -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.
  ///