feat(build): Custom codecs for generated code

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

examples/src/json-codec/common.rs

@@ -30,6 +30,10 @@ impl<T: serde::Serialize> Encoder for JsonEncoder<T> {
     fn encode(&mut self, item: Self::Item, buf: &mut EncodeBuf<'_>) -> Result<(), Self::Error> {
         serde_json::to_writer(buf.writer(), &item).map_err(|e| Status::internal(e.to_string()))
     }
+
+    fn buffer_settings(&self) -> tonic::codec::BufferSettings {
+        Default::default()
+    }
 }
 
 #[derive(Debug)]
@@ -48,6 +52,10 @@ impl<U: serde::de::DeserializeOwned> Decoder for JsonDecoder<U> {
             serde_json::from_reader(buf.reader()).map_err(|e| Status::internal(e.to_string()))?;
         Ok(Some(item))
     }
+
+    fn buffer_settings(&self) -> tonic::codec::BufferSettings {
+        Default::default()
+    }
 }
 
 /// A [`Codec`] that implements `application/grpc+json` via the serde library.

tonic-build/src/client.rs

@@ -221,7 +221,7 @@ fn generate_unary<T: Service>(
     proto_path: &str,
     compile_well_known_types: bool,
 ) -> TokenStream {
-    let codec_name = syn::parse_str::<syn::Path>(method.codec_path()).unwrap();
+    let codec_name = syn::parse_str::<syn::Path>(&method.codec_path()).unwrap();
     let ident = format_ident!("{}", method.name());
     let (request, response) = method.request_response_name(proto_path, compile_well_known_types);
     let service_name = format_service_name(service, emit_package);
@@ -252,7 +252,7 @@ fn generate_server_streaming<T: Service>(
     proto_path: &str,
     compile_well_known_types: bool,
 ) -> TokenStream {
-    let codec_name = syn::parse_str::<syn::Path>(method.codec_path()).unwrap();
+    let codec_name = syn::parse_str::<syn::Path>(&method.codec_path()).unwrap();
     let ident = format_ident!("{}", method.name());
     let (request, response) = method.request_response_name(proto_path, compile_well_known_types);
     let service_name = format_service_name(service, emit_package);
@@ -283,7 +283,7 @@ fn generate_client_streaming<T: Service>(
     proto_path: &str,
     compile_well_known_types: bool,
 ) -> TokenStream {
-    let codec_name = syn::parse_str::<syn::Path>(method.codec_path()).unwrap();
+    let codec_name = syn::parse_str::<syn::Path>(&method.codec_path()).unwrap();
     let ident = format_ident!("{}", method.name());
     let (request, response) = method.request_response_name(proto_path, compile_well_known_types);
     let service_name = format_service_name(service, emit_package);
@@ -314,7 +314,7 @@ fn generate_streaming<T: Service>(
     proto_path: &str,
     compile_well_known_types: bool,
 ) -> TokenStream {
-    let codec_name = syn::parse_str::<syn::Path>(method.codec_path()).unwrap();
+    let codec_name = syn::parse_str::<syn::Path>(&method.codec_path()).unwrap();
     let ident = format_ident!("{}", method.name());
     let (request, response) = method.request_response_name(proto_path, compile_well_known_types);
     let service_name = format_service_name(service, emit_package);

tonic-build/src/compile_settings.rs

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

tonic-build/src/lib.rs

@@ -97,6 +97,9 @@
 mod code_gen;
 pub use code_gen::CodeGenBuilder;
 
+mod compile_settings;
+pub(crate) use compile_settings::CompileSettings;
+
 /// Service generation trait.
 ///
 /// This trait can be implemented and consumed