tonic-build: option to emit Arc<Self> as server receiver (#1352)

This implements an option to emit `Arc<Self>` instead of `&self` in
server traits. This enables implementor to reference `Self` data for
indefinite duration, which may be useful for streaming requests.

Fixes: #1351

Cargo.toml

@@ -23,5 +23,6 @@ members = [
   "tests/compression",
   "tonic-web/tests/integration",
   "tests/service_named_result",
+  "tests/use_arc_self",
 ]
 resolver = "2"

tests/use_arc_self/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+authors = ["Aurimas Blažulionis <aurimas@chorus.one>"]
+edition = "2021"
+license = "MIT"
+name = "use_arc_self"
+publish = false
+version = "0.1.0"
+
+[dependencies]
+futures = "0.3"
+prost = "0.11"
+tonic = {path = "../../tonic", features = ["gzip"]}
+
+[build-dependencies]
+tonic-build = {path = "../../tonic-build" }
+
+[package.metadata.cargo-machete]
+ignored = ["prost"]

tests/use_arc_self/LICENSE

@@ -0,0 +1,19 @@
+Copyright (c) 2020 Lucio Franco
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

tests/use_arc_self/build.rs

@@ -0,0 +1,6 @@
+fn main() {
+    tonic_build::configure()
+        .use_arc_self(true)
+        .compile(&["proto/test.proto"], &["proto"])
+        .unwrap();
+}

tests/use_arc_self/proto/test.proto

@@ -0,0 +1,12 @@
+syntax = "proto3";
+
+package test;
+
+service Test {
+  rpc TestRequest(SomeData) returns (SomeData);
+}
+
+message SomeData {
+  // include a bunch of data so there actually is something to compress
+  bytes data = 1;
+}

tests/use_arc_self/src/lib.rs

@@ -0,0 +1,20 @@
+#![allow(unused_imports)]
+
+use futures::{Stream, StreamExt};
+use std::sync::Arc;
+use tonic::{Request, Response, Status};
+
+tonic::include_proto!("test");
+
+#[derive(Debug, Default)]
+struct Svc;
+
+#[tonic::async_trait]
+impl test_server::Test for Svc {
+    async fn test_request(
+        self: Arc<Self>,
+        req: Request<SomeData>,
+    ) -> Result<Response<SomeData>, Status> {
+        Ok(Response::new(req.into_inner()))
+    }
+}

tonic-build/src/code_gen.rs

@@ -12,6 +12,7 @@ pub struct CodeGenBuilder {
     attributes: Attributes,
     build_transport: bool,
     disable_comments: HashSet<String>,
+    use_arc_self: bool,
 }
 
 impl CodeGenBuilder {
@@ -57,6 +58,12 @@ impl CodeGenBuilder {
         self
     }
 
+    /// Emit `Arc<Self>` instead of `&self` in service trait.
+    pub fn use_arc_self(&mut self, enable: bool) -> &mut Self {
+        self.use_arc_self = enable;
+        self
+    }
+
     /// Generate client code based on `Service`.
     ///
     /// This takes some `Service` and will generate a `TokenStream` that contains
@@ -85,6 +92,7 @@ impl CodeGenBuilder {
             self.compile_well_known_types,
             &self.attributes,
             &self.disable_comments,
+            self.use_arc_self,
         )
     }
 }
@@ -97,6 +105,7 @@ impl Default for CodeGenBuilder {
             attributes: Attributes::default(),
             build_transport: true,
             disable_comments: HashSet::default(),
+            use_arc_self: false,
         }
     }
 }

tonic-build/src/prost.rs

@@ -37,6 +37,7 @@ pub fn configure() -> Builder {
         include_file: None,
         emit_rerun_if_changed: std::env::var_os("CARGO").is_some(),
         disable_comments: HashSet::default(),
+        use_arc_self: false,
     }
 }
 
@@ -170,6 +171,7 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
                 .compile_well_known_types(self.builder.compile_well_known_types)
                 .attributes(self.builder.server_attributes.clone())
                 .disable_comments(self.builder.disable_comments.clone())
+                .use_arc_self(self.builder.use_arc_self)
                 .generate_server(&service, &self.builder.proto_path);
 
             self.servers.extend(server);
@@ -242,6 +244,7 @@ pub struct Builder {
     pub(crate) include_file: Option<PathBuf>,
     pub(crate) emit_rerun_if_changed: bool,
     pub(crate) disable_comments: HashSet<String>,
+    pub(crate) use_arc_self: bool,
 
     out_dir: Option<PathBuf>,
 }
@@ -411,6 +414,12 @@ impl Builder {
         self
     }
 
+    /// Emit `Arc<Self>` receiver type in server traits instead of `&self`.
+    pub fn use_arc_self(mut self, enable: bool) -> Self {
+        self.use_arc_self = enable;
+        self
+    }
+
     /// Emits GRPC endpoints with no attached package. Effectively ignores protofile package declaration from grpc context.
     ///
     /// This effectively sets prost's exported package to an empty string.

tonic-build/src/server.rs

@@ -28,6 +28,7 @@ pub fn generate<T: Service>(
         compile_well_known_types,
         attributes,
         &HashSet::default(),
+        false,
     )
 }
 
@@ -38,8 +39,15 @@ pub(crate) fn generate_internal<T: Service>(
     compile_well_known_types: bool,
     attributes: &Attributes,
     disable_comments: &HashSet<String>,
+    use_arc_self: bool,
 ) -> TokenStream {
-    let methods = generate_methods(service, emit_package, proto_path, compile_well_known_types);
+    let methods = generate_methods(
+        service,
+        emit_package,
+        proto_path,
+        compile_well_known_types,
+        use_arc_self,
+    );
 
     let server_service = quote::format_ident!("{}Server", service.name());
     let server_trait = quote::format_ident!("{}", service.name());
@@ -51,6 +59,7 @@ pub(crate) fn generate_internal<T: Service>(
         compile_well_known_types,
         server_trait.clone(),
         disable_comments,
+        use_arc_self,
     );
     let package = if emit_package { service.package() } else { "" };
     // Transport based implementations
@@ -227,6 +236,7 @@ fn generate_trait<T: Service>(
     compile_well_known_types: bool,
     server_trait: Ident,
     disable_comments: &HashSet<String>,
+    use_arc_self: bool,
 ) -> TokenStream {
     let methods = generate_trait_methods(
         service,
@@ -234,6 +244,7 @@ fn generate_trait<T: Service>(
         proto_path,
         compile_well_known_types,
         disable_comments,
+        use_arc_self,
     );
     let trait_doc = generate_doc_comment(format!(
         " Generated trait containing gRPC methods that should be implemented for use with {}Server.",
@@ -255,6 +266,7 @@ fn generate_trait_methods<T: Service>(
     proto_path: &str,
     compile_well_known_types: bool,
     disable_comments: &HashSet<String>,
+    use_arc_self: bool,
 ) -> TokenStream {
     let mut stream = TokenStream::new();
 
@@ -271,18 +283,24 @@ fn generate_trait_methods<T: Service>(
                 generate_doc_comments(method.comment())
             };
 
+        let self_param = if use_arc_self {
+            quote!(self: std::sync::Arc<Self>)
+        } else {
+            quote!(&self)
+        };
+
         let method = match (method.client_streaming(), method.server_streaming()) {
             (false, false) => {
                 quote! {
                     #method_doc
-                    async fn #name(&self, request: tonic::Request<#req_message>)
+                    async fn #name(#self_param, request: tonic::Request<#req_message>)
                         -> std::result::Result<tonic::Response<#res_message>, tonic::Status>;
                 }
             }
             (true, false) => {
                 quote! {
                     #method_doc
-                    async fn #name(&self, request: tonic::Request<tonic::Streaming<#req_message>>)
+                    async fn #name(#self_param, request: tonic::Request<tonic::Streaming<#req_message>>)
                         -> std::result::Result<tonic::Response<#res_message>, tonic::Status>;
                 }
             }
@@ -298,7 +316,7 @@ fn generate_trait_methods<T: Service>(
                     type #stream: futures_core::Stream<Item = std::result::Result<#res_message, tonic::Status>> + Send + 'static;
 
                     #method_doc
-                    async fn #name(&self, request: tonic::Request<#req_message>)
+                    async fn #name(#self_param, request: tonic::Request<#req_message>)
                         -> std::result::Result<tonic::Response<Self::#stream>, tonic::Status>;
                 }
             }
@@ -314,7 +332,7 @@ fn generate_trait_methods<T: Service>(
                     type #stream: futures_core::Stream<Item = std::result::Result<#res_message, tonic::Status>> + Send + 'static;
 
                     #method_doc
-                    async fn #name(&self, request: tonic::Request<tonic::Streaming<#req_message>>)
+                    async fn #name(#self_param, request: tonic::Request<tonic::Streaming<#req_message>>)
                         -> std::result::Result<tonic::Response<Self::#stream>, tonic::Status>;
                 }
             }
@@ -345,6 +363,7 @@ fn generate_methods<T: Service>(
     emit_package: bool,
     proto_path: &str,
     compile_well_known_types: bool,
+    use_arc_self: bool,
 ) -> TokenStream {
     let mut stream = TokenStream::new();
 
@@ -361,6 +380,7 @@ fn generate_methods<T: Service>(
                 compile_well_known_types,
                 ident,
                 server_trait,
+                use_arc_self,
             ),
 
             (false, true) => generate_server_streaming(
@@ -369,6 +389,7 @@ fn generate_methods<T: Service>(
                 compile_well_known_types,
                 ident.clone(),
                 server_trait,
+                use_arc_self,
             ),
             (true, false) => generate_client_streaming(
                 method,
@@ -376,6 +397,7 @@ fn generate_methods<T: Service>(
                 compile_well_known_types,
                 ident.clone(),
                 server_trait,
+                use_arc_self,
             ),
 
             (true, true) => generate_streaming(
@@ -384,6 +406,7 @@ fn generate_methods<T: Service>(
                 compile_well_known_types,
                 ident.clone(),
                 server_trait,
+                use_arc_self,
             ),
         };
 
@@ -404,6 +427,7 @@ fn generate_unary<T: Method>(
     compile_well_known_types: bool,
     method_ident: Ident,
     server_trait: Ident,
+    use_arc_self: bool,
 ) -> TokenStream {
     let codec_name = syn::parse_str::<syn::Path>(method.codec_path()).unwrap();
 
@@ -411,6 +435,12 @@ fn generate_unary<T: Method>(
 
     let (request, response) = method.request_response_name(proto_path, compile_well_known_types);
 
+    let inner_arg = if use_arc_self {
+        quote!(inner)
+    } else {
+        quote!(&inner)
+    };
+
     quote! {
         #[allow(non_camel_case_types)]
         struct #service_ident<T: #server_trait >(pub Arc<T>);
@@ -422,7 +452,7 @@ fn generate_unary<T: Method>(
             fn call(&mut self, request: tonic::Request<#request>) -> Self::Future {
                 let inner = Arc::clone(&self.0);
                 let fut = async move {
-                    (*inner).#method_ident(request).await
+                    <T as #server_trait>::#method_ident(#inner_arg, request).await
                 };
                 Box::pin(fut)
             }
@@ -456,6 +486,7 @@ fn generate_server_streaming<T: Method>(
     compile_well_known_types: bool,
     method_ident: Ident,
     server_trait: Ident,
+    use_arc_self: bool,
 ) -> TokenStream {
     let codec_name = syn::parse_str::<syn::Path>(method.codec_path()).unwrap();
 
@@ -465,6 +496,12 @@ fn generate_server_streaming<T: Method>(
 
     let response_stream = quote::format_ident!("{}Stream", method.identifier());
 
+    let inner_arg = if use_arc_self {
+        quote!(inner)
+    } else {
+        quote!(&inner)
+    };
+
     quote! {
         #[allow(non_camel_case_types)]
         struct #service_ident<T: #server_trait >(pub Arc<T>);
@@ -477,7 +514,7 @@ fn generate_server_streaming<T: Method>(
             fn call(&mut self, request: tonic::Request<#request>) -> Self::Future {
                 let inner = Arc::clone(&self.0);
                 let fut = async move {
-                    (*inner).#method_ident(request).await
+                    <T as #server_trait>::#method_ident(#inner_arg, request).await
                 };
                 Box::pin(fut)
             }
@@ -511,12 +548,19 @@ fn generate_client_streaming<T: Method>(
     compile_well_known_types: bool,
     method_ident: Ident,
     server_trait: Ident,
+    use_arc_self: bool,
 ) -> TokenStream {
     let service_ident = quote::format_ident!("{}Svc", method.identifier());
 
     let (request, response) = method.request_response_name(proto_path, compile_well_known_types);
     let codec_name = syn::parse_str::<syn::Path>(method.codec_path()).unwrap();
 
+    let inner_arg = if use_arc_self {
+        quote!(inner)
+    } else {
+        quote!(&inner)
+    };
+
     quote! {
         #[allow(non_camel_case_types)]
         struct #service_ident<T: #server_trait >(pub Arc<T>);
@@ -529,8 +573,7 @@ fn generate_client_streaming<T: Method>(
             fn call(&mut self, request: tonic::Request<tonic::Streaming<#request>>) -> Self::Future {
                 let inner = Arc::clone(&self.0);
                 let fut = async move {
-                    (*inner).#method_ident(request).await
+                    <T as #server_trait>::#method_ident(#inner_arg, request).await
-
                 };
                 Box::pin(fut)
             }
@@ -564,6 +607,7 @@ fn generate_streaming<T: Method>(
     compile_well_known_types: bool,
     method_ident: Ident,
     server_trait: Ident,
+    use_arc_self: bool,
 ) -> TokenStream {
     let codec_name = syn::parse_str::<syn::Path>(method.codec_path()).unwrap();
 
@@ -573,6 +617,12 @@ fn generate_streaming<T: Method>(
 
     let response_stream = quote::format_ident!("{}Stream", method.identifier());
 
+    let inner_arg = if use_arc_self {
+        quote!(inner)
+    } else {
+        quote!(&inner)
+    };
+
     quote! {
         #[allow(non_camel_case_types)]
         struct #service_ident<T: #server_trait>(pub Arc<T>);
@@ -586,7 +636,7 @@ fn generate_streaming<T: Method>(
             fn call(&mut self, request: tonic::Request<tonic::Streaming<#request>>) -> Self::Future {
                 let inner = Arc::clone(&self.0);
                 let fut = async move {
-                    (*inner).#method_ident(request).await
+                    <T as #server_trait>::#method_ident(#inner_arg, request).await
                 };
                 Box::pin(fut)
             }

tonic-health/src/generated/grpc.health.v1.rs

@@ -337,7 +337,9 @@ pub mod health_server {
                             request: tonic::Request<super::HealthCheckRequest>,
                         ) -> Self::Future {
                             let inner = Arc::clone(&self.0);
-                            let fut = async move { (*inner).check(request).await };
+                            let fut = async move {
+                                <T as Health>::check(&inner, request).await
+                            };
                             Box::pin(fut)
                         }
                     }
@@ -382,7 +384,9 @@ pub mod health_server {
                             request: tonic::Request<super::HealthCheckRequest>,
                         ) -> Self::Future {
                             let inner = Arc::clone(&self.0);
-                            let fut = async move { (*inner).watch(request).await };
+                            let fut = async move {
+                                <T as Health>::watch(&inner, request).await
+                            };
                             Box::pin(fut)
                         }
                     }

tonic-reflection/src/generated/grpc.reflection.v1alpha.rs

@@ -389,7 +389,11 @@ pub mod server_reflection_server {
                         ) -> Self::Future {
                             let inner = Arc::clone(&self.0);
                             let fut = async move {
-                                (*inner).server_reflection_info(request).await
+                                <T as ServerReflection>::server_reflection_info(
+                                        &inner,
+                                        request,
+                                    )
+                                    .await
                             };
                             Box::pin(fut)
                         }