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

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)
             }