Implementing custom grpc codec that does not free the buffer when Unmarshalling

Signed-off-by: alanprot <alanprot@gmail.com>

Author: alanprot

integration/grpc_server_test.go

@@ -31,17 +31,20 @@ type mockGprcServer struct {
 	ingester_client.IngesterServer
 }
 
-func (m mockGprcServer) QueryStream(_ *ingester_client.QueryRequest, streamServer ingester_client.Ingester_QueryStreamServer) error {
+func (m mockGprcServer) QueryStream(req *ingester_client.QueryRequest, streamServer ingester_client.Ingester_QueryStreamServer) error {
 	md, _ := metadata.FromIncomingContext(streamServer.Context())
 	i, _ := strconv.Atoi(md["i"][0])
 	return streamServer.Send(createStreamResponse(i))
 }
 
 func (m mockGprcServer) Push(ctx context.Context, request *cortexpb.WriteRequest) (*cortexpb.WriteResponse, error) {
+	defer request.Free()
 	time.Sleep(time.Duration(rand.Int31n(100)) * time.Millisecond)
 	md, _ := metadata.FromIncomingContext(ctx)
 	i, _ := strconv.Atoi(md["i"][0])
 	expected := createRequest(i)
+	// Need to do this so the .String method return the same value for MessageWithBufRef
+	expected.MessageWithBufRef = request.MessageWithBufRef
 
 	if expected.String() != request.String() {
 		return nil, fmt.Errorf("expected %v, got %v", expected, request)

pkg/cortexpb/codec.go

@@ -0,0 +1,84 @@
+package cortexpb
+
+import (
+	"fmt"
+
+	"google.golang.org/grpc/encoding"
+	"google.golang.org/grpc/mem"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/protoadapt"
+)
+
+func init() {
+	c := encoding.GetCodecV2("proto")
+	encoding.RegisterCodecV2(&cortexCodec{c: c})
+}
+
+type ReleasableMessage interface {
+	RegisterBuffer(mem.Buffer)
+}
+
+type cortexCodec struct {
+	c encoding.CodecV2
+}
+
+func (c cortexCodec) Name() string {
+	return c.c.Name()
+}
+
+func (c cortexCodec) Marshal(v any) (mem.BufferSlice, error) {
+	return c.c.Marshal(v)
+}
+
+// Unmarshal Copied from https://github.com/grpc/grpc-go/blob/d2e836604b36400a54fbf04af495d12b38fa1e3a/encoding/proto/proto.go#L69-L81
+// but without releasing the buffer
+func (c *cortexCodec) Unmarshal(data mem.BufferSlice, v any) error {
+	vv := messageV2Of(v)
+	if vv == nil {
+		return fmt.Errorf("failed to unmarshal, message is %T, want proto.Message", v)
+	}
+
+	// To be in the safe side, we will never automatically release the buffer used to Unmarshal the message automatically.
+	// This should simulate the same behavior of grpc v1.65.0 and before.
+	buf := data.MaterializeToBuffer(mem.DefaultBufferPool())
+
+	err := proto.Unmarshal(buf.ReadOnlyData(), vv)
+
+	if err != nil {
+		buf.Free()
+		return err
+	}
+
+	// If v implements ReleasableMessage interface, we add the buff to be freed later when the request is no longer being used
+	if fm, ok := v.(ReleasableMessage); ok {
+		fm.RegisterBuffer(buf)
+	}
+
+	return err
+}
+
+func messageV2Of(v any) proto.Message {
+	switch v := v.(type) {
+	case protoadapt.MessageV1:
+		return protoadapt.MessageV2Of(v)
+	case protoadapt.MessageV2:
+		return v
+	}
+
+	return nil
+}
+
+var _ ReleasableMessage = &MessageWithBufRef{}
+
+type MessageWithBufRef struct {
+	bs mem.BufferSlice
+}
+
+func (m *MessageWithBufRef) RegisterBuffer(buffer mem.Buffer) {
+	m.bs = append(m.bs, buffer)
+}
+
+func (m *MessageWithBufRef) Free() {
+	m.bs.Free()
+	m.bs = m.bs[:0]
+}