Reusing grpc membuffer for QueryStreamResponse

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

Author: alanprot

integration/grpc_server_test.go

@@ -202,8 +202,12 @@ func TestConcurrentGrpcCalls(t *testing.T) {
 						require.NoError(t, err)
 						resp, err := s.Recv()
 						require.NoError(t, err)
+						require.NotEmpty(t, resp.Buffers())
 						expected := createStreamResponse(i)
+						expected.MessageWithBufRef = resp.MessageWithBufRef
 						require.Equal(t, expected.String(), resp.String())
+						resp.Free()
+
 					}(i)
 				}
 

pkg/cortexpb/codec.go

@@ -19,6 +19,8 @@ func init() {
 
 type ReleasableMessage interface {
 	RegisterBuffer(mem.Buffer)
+	buffers() []mem.Buffer
+	unref()
 }
 
 type GogoProtoMessage interface {
@@ -130,10 +132,25 @@ type MessageWithBufRef struct {
 	bs mem.BufferSlice
 }
 
+func (m *MessageWithBufRef) unref() {
+	m.bs = nil
+}
+
+func (m *MessageWithBufRef) buffers() []mem.Buffer {
+	return m.bs
+}
+
 func (m *MessageWithBufRef) RegisterBuffer(buffer mem.Buffer) {
 	m.bs = append(m.bs, buffer)
 }
 
+func (m *MessageWithBufRef) MergeBuffers(o any) {
+	if fm, ok := o.(ReleasableMessage); ok {
+		m.bs = append(m.bs, fm.buffers()...)
+		fm.unref()
+	}
+}
+
 func (m *MessageWithBufRef) Free() {
 	m.bs.Free()
 	m.bs = m.bs[:0]

pkg/cortexpb/codec_test.go

@@ -0,0 +1,39 @@
+package cortexpb
+
+import (
+	"testing"
+	"unsafe"
+
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/mem"
+)
+
+func TestMergingBuffers(t *testing.T) {
+	codec := &cortexCodec{}
+	wr := &WriteRequest{}
+
+	data, err := codec.Marshal(wr)
+	require.NoError(t, err)
+
+	wr1 := &WriteRequest{}
+	wr2 := &WriteRequest{}
+	require.NoError(t, codec.Unmarshal(data, wr1))
+	require.NoError(t, codec.Unmarshal(data, wr2))
+	// should not be the same reference but should have the same value
+	require.False(t, isSameReference(wr1.bs, wr2.bs))
+	require.NotEmpty(t, wr1.bs)
+	require.Equal(t, wr1.bs, wr2.bs)
+	finalWr := &WriteRequest{}
+	finalWr.MergeBuffers(wr1)
+	require.Empty(t, wr1.bs)
+	finalWr.MergeBuffers(wr2)
+	require.Empty(t, wr2.bs)
+	require.Len(t, finalWr.bs, 2)
+	finalWr.Free()
+	require.Empty(t, finalWr.bs)
+}
+
+func isSameReference(a, b mem.BufferSlice) bool {
+	return len(a) > 0 && len(b) > 0 &&
+		unsafe.Pointer(&a[0]) == unsafe.Pointer(&b[0])
+}

pkg/distributor/query.go

@@ -297,8 +297,11 @@ func (d *Distributor) queryIngesterStream(ctx context.Context, replicationSet ri
 	defer span.Finish()
 	hashToChunkseries := map[string]ingester_client.TimeSeriesChunk{}
 
+	resp := &ingester_client.QueryStreamResponse{}
+
 	for _, result := range results {
 		response := result.(*ingester_client.QueryStreamResponse)
+		resp.MergeBuffers(response)
 
 		// Parse any chunk series
 		for _, series := range response.Chunkseries {
@@ -310,9 +313,8 @@ func (d *Distributor) queryIngesterStream(ctx context.Context, replicationSet ri
 		}
 	}
 
-	resp := &ingester_client.QueryStreamResponse{
-		Chunkseries: make([]ingester_client.TimeSeriesChunk, 0, len(hashToChunkseries)),
-	}
+	resp.Chunkseries = make([]ingester_client.TimeSeriesChunk, 0, len(hashToChunkseries))
+	
 	for _, series := range hashToChunkseries {
 		resp.Chunkseries = append(resp.Chunkseries, series)
 	}

pkg/ingester/client/ingester.pb.go

@@ -57,6 +57,7 @@ message QueryStreamResponse {
   repeated TimeSeriesChunk chunkseries = 1 [(gogoproto.nullable) = false];
   // Not used anymore
   reserved  2;
+  cortexpb.MessageWithBufRef Ref = 1001 [(gogoproto.embed) = true, (gogoproto.customtype) = "github.com/cortexproject/cortex/pkg/cortexpb.MessageWithBufRef", (gogoproto.nullable) = false];
 }
 
 message ExemplarQueryResponse {

pkg/ingester/client/ingester.proto

@@ -92,6 +92,8 @@ type distributorQuerier struct {
 	queryIngestersWithin     time.Duration
 	isPartialDataEnabled     partialdata.IsCfgEnabledFunc
 	ingesterQueryMaxAttempts int
+
+	freeBuffers func()
 }
 
 // Select implements storage.Querier interface.
@@ -166,6 +168,8 @@ func (q *distributorQuerier) streamingSelect(ctx context.Context, sortSeries, pa
 		return storage.ErrSeriesSet(err)
 	}
 
+	q.freeBuffers = results.Free
+
 	serieses := make([]storage.Series, 0, len(results.Chunkseries))
 	for _, result := range results.Chunkseries {
 		// Sometimes the ingester can send series that have no data.
@@ -357,6 +361,10 @@ func (q *distributorQuerier) labelNamesWithMatchers(ctx context.Context, hints *
 }
 
 func (q *distributorQuerier) Close() error {
+	if q.freeBuffers != nil {
+		q.freeBuffers()
+		q.freeBuffers = nil
+	}
 	return nil
 }