vtprotocodec: user codecV2 for grpc client/server

Switch to use codecV2 backed by the default buffer pool to handle all
Marshaling/Unmarshaling by grpc clients/servers. The default buffer pool
is already setup here.

```
server/util/grpc_client/grpc_client.go
254:		experimental.WithBufferPool(mem.DefaultBufferPool()),

server/util/grpc_server/grpc_server.go
240:		experimental.BufferPool(mem.DefaultBufferPool()),
```

This does not affect the manual Marshaling/Unmarshaling that we does via
`server/util/proto.{Marshal,Unmarshal}()`. These should still use V1 for
the time being.

Author: sluongng

server/util/proto/proto.go

@@ -20,22 +20,26 @@ var Bool = gproto.Bool
 type Message = gproto.Message
 type MarshalOptions = gproto.MarshalOptions
 
-type vtprotoMessage interface {
+type VtprotoMessage interface {
 	MarshalVT() ([]byte, error)
 	UnmarshalVT([]byte) error
 	CloneMessageVT() Message
+
+	// For vtprotoCodecV2
+	MarshalToSizedBufferVT(data []byte) (int, error)
+	SizeVT() int
 }
 
 func Marshal(v Message) ([]byte, error) {
-	vt, ok := v.(vtprotoMessage)
+	vt, ok := v.(VtprotoMessage)
 	if ok {
 		return vt.MarshalVT()
 	}
 	return MarshalOld(v)
 }
 
 func Unmarshal(b []byte, v Message) error {
-	vt, ok := v.(vtprotoMessage)
+	vt, ok := v.(VtprotoMessage)
 	if ok {
 		return vt.UnmarshalVT(b)
 	}
@@ -44,7 +48,7 @@ func Unmarshal(b []byte, v Message) error {
 }
 
 func Clone(v Message) Message {
-	vt, ok := v.(vtprotoMessage)
+	vt, ok := v.(VtprotoMessage)
 	if ok {
 		return vt.CloneMessageVT()
 	}

server/util/vtprotocodec/BUILD

@@ -9,5 +9,6 @@ go_library(
         "//server/util/proto",
         "@org_golang_google_grpc//encoding",
         "@org_golang_google_grpc//encoding/proto",
+        "@org_golang_google_grpc//mem",
     ],
 )

server/util/vtprotocodec/vtprotocodec.go

@@ -1,42 +1,67 @@
 package vtprotocodec
 
 import (
-	"fmt"
-
 	"github.com/buildbuddy-io/buildbuddy/server/util/proto"
 
 	"google.golang.org/grpc/encoding"
+	"google.golang.org/grpc/mem"
+
 	_ "google.golang.org/grpc/encoding/proto" // for default proto registration purposes
 )
 
 const Name = "proto"
 
-// vtprotoCodec represents a codec able to encode and decode vt enabled
-// proto messages.
-type vtprotoCodec struct{}
+// vtprotoCodecV2 implements encoding.CodecV2 and uses vtproto and default buffer pool
+// to encode/decode proto messages. The implementation is heavily inspired by
+// https://github.com/planetscale/vtprotobuf/pull/138
+// and https://github.com/vitessio/vitess/pull/16790.
+type vtprotoCodecV2 struct {
+	fallback encoding.CodecV2
+}
+
+func (vtprotoCodecV2) Name() string {
+	return Name
+}
 
-func (vtprotoCodec) Marshal(v any) ([]byte, error) {
-	vv, ok := v.(proto.Message)
+func (c *vtprotoCodecV2) Marshal(v any) (mem.BufferSlice, error) {
+	m, ok := v.(proto.VtprotoMessage)
 	if !ok {
-		return nil, fmt.Errorf("failed to marshal, message is %T, want proto.Message", v)
+		return c.fallback.Marshal(v)
 	}
-	return proto.Marshal(vv)
+
+	size := m.SizeVT()
+	if mem.IsBelowBufferPoolingThreshold(size) {
+		buf := make([]byte, 0, size)
+		if _, err := m.MarshalToSizedBufferVT(buf[:size]); err != nil {
+			return nil, err
+		}
+		return mem.BufferSlice{mem.SliceBuffer(buf)}, nil
+	}
+
+	pool := mem.DefaultBufferPool()
+	buf := pool.Get(size)
+	if _, err := m.MarshalToSizedBufferVT((*buf)[:size]); err != nil {
+		pool.Put(buf)
+		return nil, err
+	}
+	return mem.BufferSlice{mem.NewBuffer(buf, pool)}, nil
 }
 
-func (vtprotoCodec) Unmarshal(data []byte, v any) error {
-	vv, ok := v.(proto.Message)
+func (c *vtprotoCodecV2) Unmarshal(data mem.BufferSlice, v any) error {
+	m, ok := v.(proto.VtprotoMessage)
 	if !ok {
-		return fmt.Errorf("failed to unmarshal, message is %T, want proto.Message", v)
+		return c.fallback.Unmarshal(data, v)
 	}
-	return proto.Unmarshal(data, vv)
-}
 
-func (vtprotoCodec) Name() string {
-	return Name
+	buf := data.MaterializeToBuffer(mem.DefaultBufferPool())
+	defer buf.Free()
+	return m.UnmarshalVT(buf.ReadOnlyData())
 }
 
 // RegisterCodec registers the vtprotoCodec to encode/decode proto messages with
 // all gRPC clients and servers.
 func Register() {
-	encoding.RegisterCodec(vtprotoCodec{})
+	encoding.RegisterCodecV2(&vtprotoCodecV2{
+		fallback: encoding.GetCodecV2("proto"),
+	})
 }