GRPC custom codec for measuring compression/decompression times

Question

I want to measure the cost/time of compressing and decompressing the payloads handled by the GRPC library. I've created some functions that I believe will give me reasonably accurate timings (see MeasureTimeAndCpu) for this to occur and from what I have read I need to create my own codec, wrapping the default proto codec so that I can 'override' the Marshal/Unmarshal functions:

package codec

import (
    "bytes"
    "compress/gzip"
    "fmt"
    "io/ioutil"

    "google.golang.org/grpc/encoding"
    "log"
    "my/json-over-grpc/pkg/metrics"
)

func init() {
    fmt.Println("registering custom codec")
    encoding.RegisterCodec(&TimerCodec{
        encoding.GetCodec("proto"),
    })
}

type TimerCodec struct {
    encoding.Codec
}

func (g *TimerCodec) Marshal(v interface{}) ([]byte, error) {
        fmt.Println("g.codec ", g.Codec)
        return g.Codec.Marshal(v)
}

func (g *TimerCodec) Unmarshal(data []byte, v interface{}) error {
    fmt.Println("unmarshalling")
    // Check if the data is compressed with gzip
    if len(data) >= 2 && data[0] == 0x1f && data[1] == 0x8b {
        var err error
        startTime, startCpuTime := metrics.MeasureTimeAndCpu(func() {
            var reader *gzip.Reader
            reader, err = gzip.NewReader(bytes.NewReader(data))
            if err != nil {
                return
            }
            defer reader.Close()
            uncompressed, _ := ioutil.ReadAll(reader)
            data = uncompressed
        })
        if err != nil {
            return err
        }

        log.Printf("Decompression wall time: %d, CPU time: %d", startTime, startCpuTime)
    }

    return g.Codec.Unmarshal(data, v)
}

func (g *TimerCodec) Name() string {
    return "mytimercodec" // use a unique name for your codec
}

In my server then I am doing

    cdc := &codec.TimerCodec{
        Codec: encoding.GetCodec("proto"),
    }
    encoding.RegisterCodec(cdc)

and on my client I'm doing

    conn, err := grpc.Dial(":10000",
        grpc.WithTransportCredentials(insecure.NewCredentials()),
        grpc.WithDefaultCallOptions(grpc.CallContentSubtype("mytimercodec")),
    )

and I've also tried in when making the call in the client to do

    opts := []grpc.CallOption{
        grpc.CallContentSubtype((&codec.TimerCodec{
            Codec: encoding.GetCodec("proto"),
        }).Name()), // Use TimerCodec for message transmission
    }

but what I keep facing on the client side is panic: runtime error: invalid memory address or nil pointer dereference coming from the line return g.Codec.Marshal(v) in the marshal function in my custom codec, i.e g.Codec is nil.

I'm struggling to find the docs surrounding the way to do this so I suspect I am slightly off somewhere or whether there may even be a simpler way to do this... thanks!

Answer 1

The "proto" codec never gets registered, so encoding.GetCodec("proto") always returns nil. encoding/proto/proto.go registers this codec in its own init() function, so you need to make sure this subpackage gets loaded.

Based on the way the package is organized, I suggest:

import (
    "fmt"

    "google.golang.org/grpc/encoding"
    "google.golang.org/grpc/encoding/proto"
)

func init() {
    encoding.RegisterCodec(&TimerCodec{
        encoding.GetCodec(proto.Name),
    })
}

type TimerCodec struct {
    encoding.Codec
}

However, you should use probably use the built-in Go Benchmarking library instead of writing your own timing logic: https://pkg.go.dev/testing#hdr-Benchmarks