feat(build): Custom codecs for generated code (#1599)

* feat(tonic): Custom codecs for generated code

Broadly, this change does 2 things:
1. Allow the built-in Prost codec to have its buffer sizes customized
2. Allow users to specify custom codecs on the tonic_build::prost::Builder

The Prost codec is convenient, and handles any normal use case. However,
the buffer sizes today are too large in some cases - and they may grow too
aggressively. By exposing BufferSettings, users can make a small custom
codec with their own BufferSettings to control their memory usage - or give
enormous buffers to rpc's, as their use case requires.

While one can define a custom service and methods with a custom codec today
explicitly in Rust, the code generator does not have a means to supply a
custom codec. I've reached for .codec... on the tonic_build::prost::Builder
many times and keep forgetting it's not there. This change adds .codec_path
to the Builder, so people can simply add their custom buffer codec or even
their own full top level codec without reaching for manual service definition.

* replace threadlocal with service wrapper

* pull back ProstEn/Decoder, clean up other comments

* clippy and fmt

* feedback, clean up straggler changes

examples/Cargo.toml

@@ -276,6 +276,14 @@ required-features = ["cancellation"]
 name = "cancellation-client"
 path = "src/cancellation/client.rs"
 
+[[bin]]
+name = "codec-buffers-server"
+path = "src/codec_buffers/server.rs"
+
+[[bin]]
+name = "codec-buffers-client"
+path = "src/codec_buffers/client.rs"
+
 
 [features]
 gcp = ["dep:prost-types", "tonic/tls"]

examples/build.rs

@@ -33,6 +33,14 @@ fn main() {
         .unwrap();
 
     build_json_codec_service();
+
+    let smallbuff_copy = out_dir.join("smallbuf");
+    let _ = std::fs::create_dir(smallbuff_copy.clone()); // This will panic below if the directory failed to create
+    tonic_build::configure()
+        .out_dir(smallbuff_copy)
+        .codec_path("crate::common::SmallBufferCodec")
+        .compile(&["proto/helloworld/helloworld.proto"], &["proto"])
+        .unwrap();
 }
 
 // Manually define the json.helloworld.Greeter service which used a custom JsonCodec to use json

examples/src/codec_buffers/client.rs

@@ -0,0 +1,30 @@
+//! A HelloWorld example that uses a custom codec instead of the default Prost codec.
+//!
+//! Generated code is the output of codegen as defined in the `examples/build.rs` file.
+//! The generation is the one with .codec_path("crate::common::SmallBufferCodec")
+//! The generated code assumes that a module `crate::common` exists which defines
+//! `SmallBufferCodec`, and `SmallBufferCodec` must have a Default implementation.
+
+pub mod common;
+
+pub mod small_buf {
+    include!(concat!(env!("OUT_DIR"), "/smallbuf/helloworld.rs"));
+}
+use small_buf::greeter_client::GreeterClient;
+
+use crate::small_buf::HelloRequest;
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let mut client = GreeterClient::connect("http://[::1]:50051").await?;
+
+    let request = tonic::Request::new(HelloRequest {
+        name: "Tonic".into(),
+    });
+
+    let response = client.say_hello(request).await?;
+
+    println!("RESPONSE={:?}", response);
+
+    Ok(())
+}

examples/src/codec_buffers/common.rs

@@ -0,0 +1,41 @@
+//! This module defines a common encoder with small buffers. This is useful
+//! when you have many concurrent RPC's, and not a huge volume of data per
+//! rpc normally.
+//!
+//! Note that you can customize your codecs per call to the code generator's
+//! compile function. This lets you group services by their codec needs.
+//!
+//! While this codec demonstrates customizing the built-in Prost codec, you
+//! can use this to implement other codecs as well, as long as they have a
+//! Default implementation.
+
+use std::marker::PhantomData;
+
+use prost::Message;
+use tonic::codec::{BufferSettings, Codec, ProstCodec};
+
+#[derive(Debug, Clone, Copy, Default)]
+pub struct SmallBufferCodec<T, U>(PhantomData<(T, U)>);
+
+impl<T, U> Codec for SmallBufferCodec<T, U>
+where
+    T: Message + Send + 'static,
+    U: Message + Default + Send + 'static,
+{
+    type Encode = T;
+    type Decode = U;
+
+    type Encoder = <ProstCodec<T, U> as Codec>::Encoder;
+    type Decoder = <ProstCodec<T, U> as Codec>::Decoder;
+
+    fn encoder(&mut self) -> Self::Encoder {
+        // Here, we will just customize the prost codec's internal buffer settings.
+        // You can of course implement a complete Codec, Encoder, and Decoder if
+        // you wish!
+        ProstCodec::<T, U>::raw_encoder(BufferSettings::new(512, 4096))
+    }
+
+    fn decoder(&mut self) -> Self::Decoder {
+        ProstCodec::<T, U>::raw_decoder(BufferSettings::new(512, 4096))
+    }
+}

examples/src/codec_buffers/server.rs

@@ -0,0 +1,51 @@
+//! A HelloWorld example that uses a custom codec instead of the default Prost codec.
+//!
+//! Generated code is the output of codegen as defined in the `examples/build.rs` file.
+//! The generation is the one with .codec_path("crate::common::SmallBufferCodec")
+//! The generated code assumes that a module `crate::common` exists which defines
+//! `SmallBufferCodec`, and `SmallBufferCodec` must have a Default implementation.
+
+use tonic::{transport::Server, Request, Response, Status};
+
+pub mod common;
+
+pub mod small_buf {
+    include!(concat!(env!("OUT_DIR"), "/smallbuf/helloworld.rs"));
+}
+use small_buf::{
+    greeter_server::{Greeter, GreeterServer},
+    HelloReply, HelloRequest,
+};
+
+#[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 request from {:?}", request.remote_addr());
+
+        let reply = HelloReply {
+            message: format!("Hello {}!", request.into_inner().name),
+        };
+        Ok(Response::new(reply))
+    }
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let addr = "[::1]:50051".parse().unwrap();
+    let greeter = MyGreeter::default();
+
+    println!("GreeterServer listening on {}", addr);
+
+    Server::builder()
+        .add_service(GreeterServer::new(greeter))
+        .serve(addr)
+        .await?;
+
+    Ok(())
+}

tonic-build/src/compile_settings.rs

@@ -0,0 +1,14 @@
+#[derive(Debug, Clone)]
+pub(crate) struct CompileSettings {
+    #[cfg(feature = "prost")]
+    pub(crate) codec_path: String,
+}
+
+impl Default for CompileSettings {
+    fn default() -> Self {
+        Self {
+            #[cfg(feature = "prost")]
+            codec_path: "tonic::codec::ProstCodec".to_string(),
+        }
+    }
+}

tonic-build/src/lib.rs

@@ -97,6 +97,8 @@ pub mod server;
 mod code_gen;
 pub use code_gen::CodeGenBuilder;
 
+mod compile_settings;
+
 /// Service generation trait.
 ///
 /// This trait can be implemented and consumed

tonic-build/src/prost.rs

@@ -1,4 +1,4 @@
-use crate::code_gen::CodeGenBuilder;
+use crate::{code_gen::CodeGenBuilder, compile_settings::CompileSettings};
 
 use super::Attributes;
 use proc_macro2::TokenStream;
@@ -41,6 +41,7 @@ pub fn configure() -> Builder {
         disable_comments: HashSet::default(),
         use_arc_self: false,
         generate_default_stubs: false,
+        compile_settings: CompileSettings::default(),
     }
 }
 
@@ -61,61 +62,98 @@ pub fn compile_protos(proto: impl AsRef<Path>) -> io::Result<()> {
     Ok(())
 }
 
-const PROST_CODEC_PATH: &str = "tonic::codec::ProstCodec";
-
 /// Non-path Rust types allowed for request/response types.
 const NON_PATH_TYPE_ALLOWLIST: &[&str] = &["()"];
 
-impl crate::Service for Service {
-    type Method = Method;
+/// Newtype wrapper for prost to add tonic-specific extensions
+struct TonicBuildService {
+    prost_service: Service,
+    methods: Vec<TonicBuildMethod>,
+}
+
+impl TonicBuildService {
+    fn new(prost_service: Service, settings: CompileSettings) -> Self {
+        Self {
+            // CompileSettings are currently only consumed method-by-method but if you need them in the Service, here's your spot.
+            // The tonic_build::Service trait specifies that methods are borrowed, so they have to reified up front.
+            methods: prost_service
+                .methods
+                .iter()
+                .map(|prost_method| TonicBuildMethod {
+                    prost_method: prost_method.clone(),
+                    settings: settings.clone(),
+                })
+                .collect(),
+            prost_service,
+        }
+    }
+}
+
+/// Newtype wrapper for prost to add tonic-specific extensions
+struct TonicBuildMethod {
+    prost_method: Method,
+    settings: CompileSettings,
+}
+
+impl crate::Service for TonicBuildService {
+    type Method = TonicBuildMethod;
     type Comment = String;
 
     fn name(&self) -> &str {
-        &self.name
+        &self.prost_service.name
     }
 
     fn package(&self) -> &str {
-        &self.package
+        &self.prost_service.package
     }
 
     fn identifier(&self) -> &str {
-        &self.proto_name
+        &self.prost_service.proto_name
     }
 
     fn comment(&self) -> &[Self::Comment] {
-        &self.comments.leading[..]
+        &self.prost_service.comments.leading[..]
     }
 
     fn methods(&self) -> &[Self::Method] {
-        &self.methods[..]
+        &self.methods
     }
 }
 
-impl crate::Method for Method {
+impl crate::Method for TonicBuildMethod {
     type Comment = String;
 
     fn name(&self) -> &str {
-        &self.name
+        &self.prost_method.name
     }
 
     fn identifier(&self) -> &str {
-        &self.proto_name
+        &self.prost_method.proto_name
     }
 
+    /// For code generation, you can override the codec.
+    ///
+    /// You should set the codec path to an import path that has a free
+    /// function like `fn default()`. The default value is tonic::codec::ProstCodec,
+    /// which returns a default-configured ProstCodec. You may wish to configure
+    /// the codec, e.g., with a buffer configuration.
+    ///
+    /// Though ProstCodec implements Default, it is currently only required that
+    /// the function match the Default trait's function spec.
     fn codec_path(&self) -> &str {
-        PROST_CODEC_PATH
+        &self.settings.codec_path
     }
 
     fn client_streaming(&self) -> bool {
-        self.client_streaming
+        self.prost_method.client_streaming
     }
 
     fn server_streaming(&self) -> bool {
-        self.server_streaming
+        self.prost_method.server_streaming
     }
 
     fn comment(&self) -> &[Self::Comment] {
-        &self.comments.leading[..]
+        &self.prost_method.comments.leading[..]
     }
 
     fn request_response_name(
@@ -140,8 +178,14 @@ impl crate::Method for Method {
             }
         };
 
-        let request = convert_type(&self.input_proto_type, &self.input_type);
-        let response = convert_type(&self.output_proto_type, &self.output_type);
+        let request = convert_type(
+            &self.prost_method.input_proto_type,
+            &self.prost_method.input_type,
+        );
+        let response = convert_type(
+            &self.prost_method.output_proto_type,
+            &self.prost_method.output_type,
+        );
         (request, response)
     }
 }
@@ -176,7 +220,10 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
                 .disable_comments(self.builder.disable_comments.clone())
                 .use_arc_self(self.builder.use_arc_self)
                 .generate_default_stubs(self.builder.generate_default_stubs)
-                .generate_server(&service, &self.builder.proto_path);
+                .generate_server(
+                    &TonicBuildService::new(service.clone(), self.builder.compile_settings.clone()),
+                    &self.builder.proto_path,
+                );
 
             self.servers.extend(server);
         }
@@ -188,7 +235,10 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
                 .attributes(self.builder.client_attributes.clone())
                 .disable_comments(self.builder.disable_comments.clone())
                 .build_transport(self.builder.build_transport)
-                .generate_client(&service, &self.builder.proto_path);
+                .generate_client(
+                    &TonicBuildService::new(service, self.builder.compile_settings.clone()),
+                    &self.builder.proto_path,
+                );
 
             self.clients.extend(client);
         }
@@ -252,6 +302,7 @@ pub struct Builder {
     pub(crate) disable_comments: HashSet<String>,
     pub(crate) use_arc_self: bool,
     pub(crate) generate_default_stubs: bool,
+    pub(crate) compile_settings: CompileSettings,
 
     out_dir: Option<PathBuf>,
 }
@@ -524,6 +575,16 @@ impl Builder {
         self
     }
 
+    /// Override the default codec.
+    ///
+    /// If set, writes `{codec_path}::default()` in generated code wherever a codec is created.
+    ///
+    /// This defaults to `"tonic::codec::ProstCodec"`
+    pub fn codec_path(mut self, codec_path: impl Into<String>) -> Self {
+        self.compile_settings.codec_path = codec_path.into();
+        self
+    }
+
     /// Compile the .proto files and execute code generation.
     pub fn compile(
         self,