tonic icon indicating copy to clipboard operation
tonic copied to clipboard

feat(tonic): Custom codecs for generated code

Open kvcache opened this issue 1 year ago • 3 comments

Motivation

My company has a pubsub service powered in part by Tonic. It is designed to accept hundreds of thousands of subscriptions per EC2 instance. Subscriptions normally receive infrequent, very small messages - to the tune of 200 bytes. It is rare that a message is over 1kb, and the service has a cap at 4kb.

I observed that memory is the limiting factor for reaching the required memory density goals. I've tuned our structs and removed as many extra vtable references as possible. I even gave a concrete stream type up to tonic in the subscription rpc implementation to save 8 bytes per subscription. These optimizations, per Amdahl's Law, can only save on the order of 2% memory (in my system) with 16 kilobytes allocated to encoding and decoding per stream.

I'm not the only person at my company who writes services, and our ecosystem that touches protocol buffers is multi-language. Generated code helps to ensure consistency in cross-language service vs client definition. For these reasons and others, we strongly prefer to keep the tonic service traits as generated code. The Tonic code generator is awesome!

I experimented with changing decode.rs's const BUFFER_SIZE: usize = 8 * 1024; to smaller values, and it did the needful. This PR is an attempt to make it configurable.

Solution

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.

I've included a hello world example with a customized codec similar to that which I'm currently compiling in my service.

By exposing ProstEncoder/ProstDecoder, BufferSettings, and a tonic_build::prost::Builder::codec_path setting, I am able to meet the memory density requirements for our small item, low throughput streaming rpcs. I am also able to separate the buffer size settings between small message Services and large message Services.

kvcache avatar Jan 12 '24 18:01 kvcache

Force-pushed commit to rebase on the tip of master. Force push was used because this hasn't been reviewed yet.

kvcache avatar Feb 08 '24 23:02 kvcache

CI is failing due to

  error: unused import: `compile_settings::CompileSettings`
     --> tonic-build/src/lib.rs:101:16
      |
  101 | pub(crate) use compile_settings::CompileSettings;
      |                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      |
      = note: `-D unused-imports` implied by `-D warnings`
      = help: to override `-D warnings` add `#[allow(unused_imports)]`
  
  error: field `codec_path` is never read
   --> tonic-build/src/compile_settings.rs:3:16
    |
  2 | pub(crate) struct CompileSettings {
    |                   --------------- field in this struct
  3 |     pub(crate) codec_path: String,
    |                ^^^^^^^^^^
    |
    = note: `CompileSettings` has derived impls for the traits `Clone` and `Debug`, but these are intentionally ignored during dead code analysis
    = note: `-D dead-code` implied by `-D warnings`
    = help: to override `-D warnings` add `#[allow(dead_code)]`

LucioFranco avatar Feb 14 '24 13:02 LucioFranco

I installed nightly, cargo-hack, and cargo-udeps to reproduce the workflow error locally. I addressed it by gating the codec path on the prost feature. If it's needed outside of that feature, I think it can be consumed and the feature gates removed.

kvcache avatar Feb 15 '24 22:02 kvcache