Repository: grpc/grpc-go
File: encoding/proto/proto.go (Go)

Code before:
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/protoadapt"

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

// Name is the name registered for the proto compressor.
const Name = "proto"

func init() {
	encoding.RegisterCodecV2(&codecV2{})
}

// codec is a CodecV2 implementation with protobuf. It is the default codec for
// gRPC.
type codecV2 struct{}

func (c *codecV2) Marshal(v any) (data mem.BufferSlice, err error) {
	vv := messageV2Of(v)
	if vv == nil {
		return nil, fmt.Errorf("proto: failed to marshal, message is %T, want proto.Message", v)
	}

	size := proto.Size(vv)
	if mem.IsBelowBufferPoolingThreshold(size) {
		buf, err := proto.MarshalOptions{UseCachedSize: true}.Marshal(vv)
		if err != nil {
			return nil, err
		}
		data = append(data, mem.SliceBuffer(buf))
	} else {
		pool := mem.DefaultBufferPool()
		buf := pool.Get(size)
		if _, err := (proto.MarshalOptions{UseCachedSize: true}).MarshalAppend((*buf)[:0], vv); err != nil {
			pool.Put(buf)
			return nil, err
		}
		data = append(data, mem.NewBuffer(buf, pool))
	}

	return data, nil
}

func (c *codecV2) Unmarshal(data mem.BufferSlice, v any) (err error) {
	vv := messageV2Of(v)
	if vv == nil {
		return fmt.Errorf("failed to unmarshal, message is %T, want proto.Message", v)
	}

	buf := data.MaterializeToBuffer(mem.DefaultBufferPool())

Review comment (performance):
@eshitachandwani `proto.Size()` is being called here, which should cause the size to be calculated and cached. `proto.Marshal()` is called just below, without performing any mutations on the proto message. I believe it should always be safe to use the cache here, right?

Code after:
	"google.golang.org/grpc/encoding"
	"google.golang.org/grpc/mem"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/protoadapt"
)

// Name is the name registered for the proto compressor.
const Name = "proto"

func init() {
	encoding.RegisterCodecV2(&codecV2{})
}

// codec is a CodecV2 implementation with protobuf. It is the default codec for
// gRPC.
type codecV2 struct{}

func (c *codecV2) Marshal(v any) (data mem.BufferSlice, err error) {
	vv := messageV2Of(v)
	if vv == nil {
		return nil, fmt.Errorf("proto: failed to marshal, message is %T, want proto.Message", v)
	}

	// Important: if we remove this Size call then we cannot use
	// UseCachedSize in MarshalOptions below.
	size := proto.Size(vv)

	// MarshalOptions with UseCachedSize allows reusing the result from the
	// previous Size call. This is safe here because:
	//
	// 1. We just computed the size.
	// 2. We assume the message is not being mutated concurrently.
	//
	// Important: If the proto.Size call above is removed, using UseCachedSize
	// becomes unsafe and may lead to incorrect marshaling.
	//
	// For more details, see the doc of UseCachedSize:
	// https://pkg.go.dev/google.golang.org/protobuf/proto#MarshalOptions
	marshalOptions := proto.MarshalOptions{UseCachedSize: true}

	if mem.IsBelowBufferPoolingThreshold(size) {
		buf, err := marshalOptions.Marshal(vv)
		if err != nil {
			return nil, err
		}
		data = append(data, mem.SliceBuffer(buf))
	} else {
		pool := mem.DefaultBufferPool()
		buf := pool.Get(size)
		if _, err := marshalOptions.MarshalAppend((*buf)[:0], vv); err != nil {
			pool.Put(buf)