Nest server and client in their own mods

tonic-build/src/lib.rs

@@ -1,5 +1,6 @@
+use proc_macro2::TokenStream;
 use prost_build::Config;
-use std::{io, path, process::Command};
+use std::{io, path, path::Path, process::Command};
 
 mod client;
 mod service;
@@ -9,13 +10,25 @@ where
     P: AsRef<path::Path>,
 {
     let out_dir = std::env::var("OUT_DIR").unwrap();
+    compile_protos_with_out_dir(protos, includes, package, out_dir.as_str())
+}
+
+pub fn compile_protos_with_out_dir<P: AsRef<Path>>(
+    protos: &[P],
+    includes: &[P],
+    package: &str,
+    out_dir: impl AsRef<Path>,
+) -> io::Result<()> {
     let mut config = Config::new();
 
-    config.service_generator(Box::new(ServiceGenerator {}));
+    config.service_generator(Box::new(ServiceGenerator::default()));
-    config.out_dir(&out_dir);
+    config.out_dir(out_dir.as_ref());
     config.compile_protos(protos, includes)?;
 
-    fmt(&out_dir, &format!("{}.rs", package));
+    fmt(
+        out_dir.as_ref().to_str().expect("Execpted utf8 out_dir"),
+        &format!("{}.rs", package),
+    );
 
     Ok(())
 }
@@ -34,17 +47,44 @@ fn fmt(out_dir: &str, file: &str) {
     assert!(out.status.success());
 }
 
-pub struct ServiceGenerator {}
+#[derive(Default)]
+pub struct ServiceGenerator {
+    clients: TokenStream,
+    servers: TokenStream,
+}
 
 impl prost_build::ServiceGenerator for ServiceGenerator {
-    fn generate(&mut self, service: prost_build::Service, buf: &mut String) {
+    fn generate(&mut self, service: prost_build::Service, _buf: &mut String) {
-        let path = "self";
+        let path = "super";
+
         let server = service::generate(&service, path);
-        let code = format!("{}", server);
+        self.servers.extend(server);
-        buf.push_str(&code);
 
         let client = client::generate(&service, path);
-        let code = format!("{}", client);
+        self.clients.extend(client);
-        buf.push_str(&code);
+    }
+
+    fn finalize(&mut self, buf: &mut String) {
+        if !self.clients.is_empty() && !self.servers.is_empty() {
+            let clients = &self.clients;
+            let servers = &self.servers;
+
+            let service = quote::quote! {
+                pub mod client {
+                    #![allow(unused_variables, dead_code, missing_docs)]
+
+                    #clients
+                }
+
+                pub mod server {
+                    #![allow(unused_variables, dead_code, missing_docs)]
+
+                    #servers
+                }
+            };
+
+            let code = format!("{}", service);
+            buf.push_str(&code);
+        }
     }
 }

tonic-build/src/service.rs

@@ -321,12 +321,10 @@ fn generate_streaming(
 
     let response_stream = quote::format_ident!("{}Stream", method.proto_name);
 
-    // TODO: parse response stream type, if it is a concrete type then use that
-    // as the ResponseStream type, if it is a impl Trait then we need to box.
     quote! {
-        struct #service_ident<T: #server_trait >(pub std::sync::Arc<T>);
+        struct #service_ident<T: #server_trait>(pub std::sync::Arc<T>);
 
-        impl<T: #server_trait > tonic::server::StreamingService<#request> for #service_ident <T>
+        impl<T: #server_trait> tonic::server::StreamingService<#request> for #service_ident <T>
         {
             type Response = #response;
             type ResponseStream = T::#response_stream;

tonic-examples/src/helloworld/client.rs

@@ -4,15 +4,17 @@ pub mod hello_world {
     include!(concat!(env!("OUT_DIR"), "/helloworld.rs"));
 }
 
+use hello_world::{client::GreeterClient, HelloRequest};
+
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let origin = vec![http::Uri::from_static("http://[::1]:50051").into()];
 
     let svc = Channel::builder().balance_list(origin)?;
 
-    let mut client = hello_world::GreeterClient::new(svc);
+    let mut client = GreeterClient::new(svc);
 
-    let request = tonic::Request::new(hello_world::HelloRequest {
+    let request = tonic::Request::new(HelloRequest {
         name: "hello".into(),
     });
 

tonic-examples/src/helloworld/server.rs

@@ -4,17 +4,22 @@ pub mod hello_world {
     include!(concat!(env!("OUT_DIR"), "/helloworld.rs"));
 }
 
+use hello_world::{
+    server::{Greeter, GreeterServer},
+    HelloReply, HelloRequest,
+};
+
 #[derive(Default)]
 pub struct MyGreeter {
     data: String,
 }
 
 #[tonic::async_trait]
-impl hello_world::Greeter for MyGreeter {
+impl Greeter for MyGreeter {
     async fn say_hello(
         &self,
-        request: Request<hello_world::HelloRequest>,
+        request: Request<HelloRequest>,
-    ) -> Result<Response<hello_world::HelloReply>, Status> {
+    ) -> Result<Response<HelloReply>, Status> {
         println!("Got a request: {:?}", request);
 
         let string = &self.data;
@@ -34,7 +39,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let greeter = MyGreeter::default();
 
     Server::builder()
-        .serve(addr, hello_world::GreeterServer::new(greeter))
+        .serve(addr, GreeterServer::new(greeter))
         .await?;
 
     Ok(())

tonic-examples/src/routeguide/client.rs

@@ -8,12 +8,14 @@ mod route_guide {
     include!(concat!(env!("OUT_DIR"), "/routeguide.rs"));
 }
 
+use route_guide::client::RouteGuideClient;
+
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let origin = http::Uri::from_static("http://[::1]:10000");
 
     let svc = Channel::builder().build(origin)?;
-    let mut client = route_guide::RouteGuideClient::new(svc);
+    let mut client = RouteGuideClient::new(svc);
 
     let start = Instant::now();
 

tonic-examples/src/routeguide/server.rs

@@ -14,7 +14,7 @@ pub mod routeguide {
     include!(concat!(env!("OUT_DIR"), "/routeguide.rs"));
 }
 
-use routeguide::{Feature, Point, Rectangle, RouteNote, RouteSummary};
+use routeguide::{server, Feature, Point, Rectangle, RouteNote, RouteSummary};
 
 #[derive(Debug)]
 pub struct RouteGuide {
@@ -28,7 +28,7 @@ struct State {
 }
 
 #[tonic::async_trait]
-impl routeguide::RouteGuide for RouteGuide {
+impl server::RouteGuide for RouteGuide {
     async fn get_feature(&self, request: Request<Point>) -> Result<Response<Feature>, Status> {
         println!("GetFeature = {:?}", request);
 
@@ -168,7 +168,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         },
     };
 
-    let svc = routeguide::RouteGuideServer::new(route_guide);
+    let svc = server::RouteGuideServer::new(route_guide);
 
     Server::builder().serve(addr, svc).await?;
 

tonic-interop/src/client.rs

@@ -1,4 +1,4 @@
-use crate::{pb::*, test_assert, TestAssertion};
+use crate::{pb::client::*, pb::*, test_assert, TestAssertion};
 use futures_util::{future, stream, SinkExt, StreamExt};
 use tokio::sync::mpsc;
 use tonic::transport::Channel;

tonic-interop/src/server.rs

@@ -1,9 +1,9 @@
-use crate::pb::*;
+use crate::pb::{self, *};
 use std::pin::Pin;
 use tonic::{Code, Request, Response, Status};
 
-pub fn create() -> TestServiceServer<TestService> {
+pub fn create() -> pb::server::TestServiceServer<TestService> {
-    TestServiceServer::new(TestService {
+    server::TestServiceServer::new(TestService {
         data: String::new(),
     })
 }
@@ -19,7 +19,7 @@ type Stream<T> =
     Pin<Box<dyn futures_core::Stream<Item = std::result::Result<T, Status>> + Send + 'static>>;
 
 #[tonic::async_trait]
-impl crate::pb::TestService for TestService {
+impl pb::server::TestService for TestService {
     async fn empty_call(&self, _request: Request<Empty>) -> Result<Empty> {
         println!("empty_call");
         Ok(Response::new(Empty {}))