Use a new prost ServiceGenerator finalizer to fix duplicate client and server definitions

Cargo.toml

@@ -5,6 +5,7 @@ members = [
     "tonic-examples",
     "tonic-interop",
 
+    "tests/complex_package_layout",
     "tests/included_service",
     "tests/same_name",
     "tests/wellknown",

tests/complex_package_layout/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "complex_package_layout"
+version = "0.1.0"
+authors = ["Jay Oster <[email protected]>"]
+edition = "2018"
+publish = false
+license = "MIT"
+
+[dependencies]
+bytes = "0.4"
+prost = "0.5"
+tonic = { path = "../../tonic" }
+
+[build-dependencies]
+tonic-build = { path = "../../tonic-build" }

tests/complex_package_layout/build.rs

@@ -0,0 +1,9 @@
+fn main() {
+    tonic_build::configure()
+        .compile(&["proto/hello.proto", "proto/goodbye.proto"], &["proto"])
+        .unwrap();
+
+    println!("cargo:rerun-if-changed=proto/hello.proto");
+    println!("cargo:rerun-if-changed=proto/goodbye.proto");
+    println!("cargo:rerun-if-changed=proto/types.proto");
+}

tests/complex_package_layout/proto/goodbye.proto

@@ -0,0 +1,9 @@
+syntax = "proto3";
+
+import "types.proto";
+
+package helloworld;
+
+service Farewell {
+  rpc Goodbye (Message) returns (Response) {}
+}

tests/complex_package_layout/proto/hello.proto

@@ -0,0 +1,9 @@
+syntax = "proto3";
+
+import "types.proto";
+
+package helloworld;
+
+service Greeting {
+  rpc Hello (Message) returns (Response) {}
+}

tests/complex_package_layout/proto/types.proto

@@ -0,0 +1,11 @@
+syntax = "proto3";
+
+package helloworld;
+
+message Message {
+  string say = 1;
+}
+
+message Response {
+  string say = 1;
+}

tests/complex_package_layout/src/lib.rs

@@ -0,0 +1,8 @@
+mod pb {
+    tonic::include_proto!("helloworld");
+}
+
+// Ensure that multiple services defined for a single package, spread across multiple `.proto`
+// files, are all available in the generated Rust module.
+type _Test1 = dyn pb::server::Greeting;
+type _Test2 = dyn pb::server::Farewell;

tonic-build/Cargo.toml

@@ -16,7 +16,7 @@ keywords = ["rpc", "grpc", "async", "codegen", "protobuf"]
 
 
 [dependencies]
-prost-build = "0.5"
+prost-build = { git = "https://github.com/parasyte/prost.git", branch = "feature/defer-service-generator-finalize" }
 syn = "1.0"
 quote = "1.0"
 proc-macro2 = "1.0"

tonic-build/src/lib.rs

@@ -61,6 +61,7 @@ use proc_macro2::{Delimiter, Group, Ident, Literal, Punct, Spacing, Span, TokenS
 use prost_build::{Config, Method};
 use quote::{ToTokens, TokenStreamExt};
 
+use std::collections::HashMap;
 #[cfg(feature = "rustfmt")]
 use std::process::Command;
 use std::{
@@ -217,16 +218,16 @@ fn fmt(out_dir: &str) {
 
 struct ServiceGenerator {
     builder: Builder,
-    clients: TokenStream,
-    servers: TokenStream,
+    clients: HashMap<String, TokenStream>,
+    servers: HashMap<String, TokenStream>,
 }
 
 impl ServiceGenerator {
     fn new(builder: Builder) -> Self {
         ServiceGenerator {
             builder,
-            clients: TokenStream::default(),
-            servers: TokenStream::default(),
+            clients: HashMap::default(),
+            servers: HashMap::default(),
         }
     }
 }
@@ -237,18 +238,25 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
 
         if self.builder.build_server {
             let server = server::generate(&service, path);
-            self.servers.extend(server);
+            self.servers
+                .entry(service.package.clone())
+                .or_insert_with(TokenStream::default)
+                .extend(server);
         }
 
         if self.builder.build_client {
             let client = client::generate(&service, path);
-            self.clients.extend(client);
+            self.clients
+                .entry(service.package)
+                .or_insert_with(TokenStream::default)
+                .extend(client);
         }
     }
 
-    fn finalize(&mut self, buf: &mut String) {
-        if self.builder.build_client && !self.clients.is_empty() {
-            let clients = &self.clients;
+    fn finalize_package(&mut self, package: &str, buf: &mut String) {
+        let clients = self.clients.get(package);
+        if self.builder.build_client && clients.is_some() {
+            let clients = clients.unwrap();
 
             let client_service = quote::quote! {
                 /// Generated client implementations.
@@ -263,11 +271,13 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
             let code = format!("{}", client_service);
             buf.push_str(&code);
 
-            self.clients = TokenStream::default();
+            self.clients
+                .insert(package.to_string(), TokenStream::default());
         }
 
-        if self.builder.build_server && !self.servers.is_empty() {
-            let servers = &self.servers;
+        let servers = self.servers.get(package);
+        if self.builder.build_server && servers.is_some() {
+            let servers = servers.unwrap();
 
             let server_service = quote::quote! {
                 /// Generated server implementations.
@@ -282,7 +292,8 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
             let code = format!("{}", server_service);
             buf.push_str(&code);
 
-            self.servers = TokenStream::default();
+            self.servers
+                .insert(package.to_string(), TokenStream::default());
         }
     }
 }