Support for adding services dynamically. (#1408)

* Support for adding services dynamically.

* Fix typo in examples/Cargo.toml

Co-authored-by: Andrew Hickman <andrew.hickman1@sky.com>

* PR feedback: add test for RoutesBuilder

* Remove let else from the test to fix the failing build.

* Fix test after removal of futures-util crate.

* run cargo fmt

---------

Co-authored-by: tomek.sroka <tomek.sroka@cloudkitchens.com>
Co-authored-by: Andrew Hickman <andrew.hickman1@sky.com>
Co-authored-by: Lucio Franco <luciofranco14@gmail.com>

examples/Cargo.toml

@@ -253,6 +253,10 @@ name = "richer-error-server-vec"
 path = "src/richer-error/server_vec.rs"
 required-features = ["types"]
 
+[[bin]]
+name = "dynamic-server"
+path = "src/dynamic/server.rs"
+
 [[bin]]
 name = "h2c-server"
 path = "src/h2c/server.rs"

examples/src/dynamic/server.rs

@@ -0,0 +1,94 @@
+use std::env;
+use tonic::{transport::server::RoutesBuilder, transport::Server, Request, Response, Status};
+
+use hello_world::greeter_server::{Greeter, GreeterServer};
+use hello_world::{HelloReply, HelloRequest};
+
+use echo::echo_server::{Echo, EchoServer};
+use echo::{EchoRequest, EchoResponse};
+
+pub mod hello_world {
+    tonic::include_proto!("helloworld");
+}
+
+pub mod echo {
+    tonic::include_proto!("grpc.examples.unaryecho");
+}
+
+type EchoResult<T> = Result<Response<T>, Status>;
+
+#[derive(Default)]
+pub struct MyEcho {}
+
+#[tonic::async_trait]
+impl Echo for MyEcho {
+    async fn unary_echo(&self, request: Request<EchoRequest>) -> EchoResult<EchoResponse> {
+        println!("Got an echo request from {:?}", request.remote_addr());
+
+        let message = format!("you said: {}", request.into_inner().message);
+
+        Ok(Response::new(EchoResponse { message }))
+    }
+}
+
+fn init_echo(args: &[String], builder: &mut RoutesBuilder) {
+    let enabled = args
+        .into_iter()
+        .find(|arg| arg.as_str() == "echo")
+        .is_some();
+    if enabled {
+        println!("Adding Echo service...");
+        let svc = EchoServer::new(MyEcho::default());
+        builder.add_service(svc);
+    }
+}
+
+#[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 greet request from {:?}", request.remote_addr());
+
+        let reply = hello_world::HelloReply {
+            message: format!("Hello {}!", request.into_inner().name),
+        };
+        Ok(Response::new(reply))
+    }
+}
+
+fn init_greeter(args: &[String], builder: &mut RoutesBuilder) {
+    let enabled = args
+        .into_iter()
+        .find(|arg| arg.as_str() == "greeter")
+        .is_some();
+
+    if enabled {
+        println!("Adding Greeter service...");
+        let svc = GreeterServer::new(MyGreeter::default());
+        builder.add_service(svc);
+    }
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let args: Vec<String> = env::args().collect();
+    let mut routes_builder = RoutesBuilder::default();
+    init_greeter(&args, &mut routes_builder);
+    init_echo(&args, &mut routes_builder);
+
+    let addr = "[::1]:50051".parse().unwrap();
+
+    println!("Grpc server listening on {}", addr);
+
+    Server::builder()
+        .add_routes(routes_builder.routes())
+        .serve(addr)
+        .await?;
+
+    Ok(())
+}

tests/integration_tests/tests/routes_builder.rs

@@ -0,0 +1,105 @@
+use std::time::Duration;
+
+use tokio::sync::oneshot;
+use tokio_stream::StreamExt;
+
+use integration_tests::pb::{
+    test1_client, test1_server, test_client, test_server, Input, Input1, Output, Output1,
+};
+use tonic::codegen::BoxStream;
+use tonic::transport::server::RoutesBuilder;
+use tonic::{
+    transport::{Endpoint, Server},
+    Request, Response, Status,
+};
+
+#[tokio::test]
+async fn multiple_service_using_routes_builder() {
+    struct Svc1;
+
+    #[tonic::async_trait]
+    impl test_server::Test for Svc1 {
+        async fn unary_call(&self, _req: Request<Input>) -> Result<Response<Output>, Status> {
+            Ok(Response::new(Output {}))
+        }
+    }
+
+    struct Svc2;
+
+    #[tonic::async_trait]
+    impl test1_server::Test1 for Svc2 {
+        async fn unary_call(&self, request: Request<Input1>) -> Result<Response<Output1>, Status> {
+            Ok(Response::new(Output1 {
+                buf: request.into_inner().buf,
+            }))
+        }
+
+        type StreamCallStream = BoxStream<Output1>;
+
+        async fn stream_call(
+            &self,
+            request: Request<Input1>,
+        ) -> Result<Response<Self::StreamCallStream>, Status> {
+            let output = Output1 {
+                buf: request.into_inner().buf,
+            };
+            let stream = tokio_stream::iter(vec![Ok(output)]);
+
+            Ok(Response::new(Box::pin(stream)))
+        }
+    }
+
+    let svc1 = test_server::TestServer::new(Svc1);
+    let svc2 = test1_server::Test1Server::new(Svc2);
+
+    let (tx, rx) = oneshot::channel::<()>();
+    let mut routes_builder = RoutesBuilder::default();
+    routes_builder.add_service(svc1).add_service(svc2);
+
+    let jh = tokio::spawn(async move {
+        Server::builder()
+            .add_routes(routes_builder.routes())
+            .serve_with_shutdown("127.0.0.1:1400".parse().unwrap(), async { drop(rx.await) })
+            .await
+            .unwrap();
+    });
+
+    tokio::time::sleep(Duration::from_millis(100)).await;
+
+    let channel = Endpoint::from_static("http://127.0.0.1:1400")
+        .connect()
+        .await
+        .unwrap();
+
+    let mut client1 = test_client::TestClient::new(channel.clone());
+    let mut client2 = test1_client::Test1Client::new(channel);
+
+    client1.unary_call(Input {}).await.unwrap();
+
+    let resp2 = client2
+        .unary_call(Input1 {
+            buf: b"hello".to_vec(),
+        })
+        .await
+        .unwrap()
+        .into_inner();
+    assert_eq!(&resp2.buf, b"hello");
+    let mut stream_response = client2
+        .stream_call(Input1 {
+            buf: b"world".to_vec(),
+        })
+        .await
+        .unwrap()
+        .into_inner();
+    let first = match stream_response.next().await {
+        Some(Ok(first)) => first,
+        _ => panic!("expected one non-error item in the stream call response"),
+    };
+
+    assert_eq!(&first.buf, b"world");
+    assert!(stream_response.next().await.is_none());
+
+    tx.send(()).unwrap();
+
+    jh.await.unwrap();
+}

tonic/src/transport/server/mod.rs

@@ -10,6 +10,8 @@ mod tls;
 mod unix;
 
 pub use super::service::Routes;
+pub use super::service::RoutesBuilder;
+
 pub use crate::server::NamedService;
 pub use conn::{Connected, TcpConnectInfo};
 #[cfg(feature = "tls")]
@@ -391,6 +393,17 @@ impl<L> Server<L> {
         Router::new(self.clone(), routes)
     }
 
+    /// Create a router with given [`Routes`].
+    ///
+    /// This will clone the `Server` builder and create a router that will
+    /// route around different services that were already added to the provided `routes`.
+    pub fn add_routes(&mut self, routes: Routes) -> Router<L>
+    where
+        L: Clone,
+    {
+        Router::new(self.clone(), routes)
+    }
+
     /// Set the [Tower] [`Layer`] all services will be wrapped in.
     ///
     /// This enables using middleware from the [Tower ecosystem][eco].

tonic/src/transport/service/mod.rs

@@ -23,3 +23,4 @@ pub(crate) use self::tls::{TlsAcceptor, TlsConnector};
 pub(crate) use self::user_agent::UserAgent;
 
 pub use self::router::Routes;
+pub use self::router::RoutesBuilder;

tonic/src/transport/service/router.rs

@@ -21,8 +21,37 @@ pub struct Routes {
     router: axum::Router,
 }
 
+#[derive(Debug, Default, Clone)]
+/// Allows adding new services to routes by passing a mutable reference to this builder.
+pub struct RoutesBuilder {
+    routes: Option<Routes>,
+}
+
+impl RoutesBuilder {
+    /// Add a new service.
+    pub fn add_service<S>(&mut self, svc: S) -> &mut Self
+    where
+        S: Service<Request<Body>, Response = Response<BoxBody>, Error = Infallible>
+            + NamedService
+            + Clone
+            + Send
+            + 'static,
+        S::Future: Send + 'static,
+        S::Error: Into<crate::Error> + Send,
+    {
+        let routes = self.routes.take().unwrap_or_default();
+        self.routes.replace(routes.add_service(svc));
+        self
+    }
+
+    /// Returns the routes with added services or empty [`Routes`] if no service was added
+    pub fn routes(self) -> Routes {
+        self.routes.unwrap_or_default()
+    }
+}
 impl Routes {
-    pub(crate) fn new<S>(svc: S) -> Self
+    /// Create a new routes with `svc` already added to it.
+    pub fn new<S>(svc: S) -> Self
     where
         S: Service<Request<Body>, Response = Response<BoxBody>, Error = Infallible>
             + NamedService
@@ -36,7 +65,8 @@ impl Routes {
         Self { router }.add_service(svc)
     }
 
-    pub(crate) fn add_service<S>(mut self, svc: S) -> Self
+    /// Add a new service.
+    pub fn add_service<S>(mut self, svc: S) -> Self
     where
         S: Service<Request<Body>, Response = Response<BoxBody>, Error = Infallible>
             + NamedService