store: add chunk pooling (#7771)

Pool byte slices inside of aggrchunks.

Signed-off-by: Giedrius Statkevičius <giedrius.statkevicius@vinted.com>

go.mod

@@ -64,7 +64,7 @@ require (
 	github.com/sony/gobreaker v0.5.0
 	github.com/stretchr/testify v1.9.0
 	github.com/thanos-io/objstore v0.0.0-20240913074259-63feed0da069
-	github.com/thanos-io/promql-engine v0.0.0-20240718195911-cdbd6dfed36b
+	github.com/thanos-io/promql-engine v0.0.0-20240921092401-37747eddbd31
 	github.com/uber/jaeger-client-go v2.30.0+incompatible
 	github.com/uber/jaeger-lib v2.4.1+incompatible // indirect
 	github.com/vimeo/galaxycache v0.0.0-20210323154928-b7e5d71c067a

go.sum

@@ -1369,8 +1369,8 @@ github.com/thanos-community/galaxycache v0.0.0-20211122094458-3a32041a1f1e h1:f1
 github.com/thanos-community/galaxycache v0.0.0-20211122094458-3a32041a1f1e/go.mod h1:jXcofnrSln/cLI6/dhlBxPQZEEQHVPCcFaH75M+nSzM=
 github.com/thanos-io/objstore v0.0.0-20240913074259-63feed0da069 h1:TUPZ6euAh8I62KrpDnBIg7k2C5HjgXQnVHoUUMacGwM=
 github.com/thanos-io/objstore v0.0.0-20240913074259-63feed0da069/go.mod h1:Cba80S8NbVBBdyZKzra7San/jXvpAxArbpFymWzIZhg=
-github.com/thanos-io/promql-engine v0.0.0-20240718195911-cdbd6dfed36b h1:V06gjM1OFiJydoClwiGOMCpBWLSpxa5FZBvBc3coQg4=
-github.com/thanos-io/promql-engine v0.0.0-20240718195911-cdbd6dfed36b/go.mod h1:Gtv7CJIxGyiGsT+bNDg4nOAsL/bVKLlpfOZUSLSyYfY=
+github.com/thanos-io/promql-engine v0.0.0-20240921092401-37747eddbd31 h1:xPaP58g+3EPohdw4cv+6jv5+LcX6LynhHvQcYwTAMxQ=
+github.com/thanos-io/promql-engine v0.0.0-20240921092401-37747eddbd31/go.mod h1:wx0JlRZtsB2S10JYUgeg5GqLfMxw31SzArP+28yyE00=
 github.com/themihai/gomemcache v0.0.0-20180902122335-24332e2d58ab h1:7ZR3hmisBWw77ZpO1/o86g+JV3VKlk3d48jopJxzTjU=
 github.com/themihai/gomemcache v0.0.0-20180902122335-24332e2d58ab/go.mod h1:eheTFp954zcWZXCU8d0AT76ftsQOTo4DTqkN/h3k1MY=
 github.com/tklauser/go-sysconf v0.3.4/go.mod h1:Cl2c8ZRWfHD5IrfHo9VN+FX9kCFjIOyVklgXycLB6ek=

pkg/query/querier.go

@@ -133,6 +133,16 @@ type querier struct {
 	selectTimeout           time.Duration
 	shardInfo               *storepb.ShardInfo
 	seriesStatsReporter     seriesStatsReporter
+
+	returnChunksMtx sync.Mutex
+	returnChunks    []*storepb.AggrChunk
+}
+
+var returnChunksSlicePool = sync.Pool{
+	New: func() interface{} {
+		r := make([]*storepb.AggrChunk, 0)
+		return &r
+	},
 }
 
 // newQuerier creates implementation of storage.Querier that fetches data from the proxy
@@ -165,6 +175,8 @@ func newQuerier(
 	if partialResponse {
 		partialResponseStrategy = storepb.PartialResponseStrategy_WARN
 	}
+
+	returnChunks := returnChunksSlicePool.Get().(*[]*storepb.AggrChunk)
 	return &querier{
 		logger:        logger,
 		selectGate:    selectGate,
@@ -181,6 +193,7 @@ func newQuerier(
 		skipChunks:              skipChunks,
 		shardInfo:               shardInfo,
 		seriesStatsReporter:     seriesStatsReporter,
+		returnChunks:            *returnChunks,
 	}
 }
 
@@ -347,6 +360,12 @@ func (q *querier) selectFn(ctx context.Context, hints *storage.SelectHints, ms .
 	if err := q.proxy.Series(&req, resp); err != nil {
 		return nil, storepb.SeriesStatsCounter{}, errors.Wrap(err, "proxy Series()")
 	}
+	q.returnChunksMtx.Lock()
+	for i := range resp.seriesSet {
+		q.returnChunks = append(q.returnChunks, resp.seriesSet[i].Chunks...)
+	}
+	q.returnChunksMtx.Unlock()
+
 	warns := annotations.New().Merge(resp.warnings)
 
 	if !q.isDedupEnabled() {
@@ -459,4 +478,15 @@ func (q *querier) LabelNames(ctx context.Context, hints *storage.LabelHints, mat
 	return resp.Names, warns, nil
 }
 
-func (q *querier) Close() error { return nil }
+func (q *querier) Close() error {
+	q.returnChunksMtx.Lock()
+	defer q.returnChunksMtx.Unlock()
+
+	for _, ch := range q.returnChunks {
+		ch.ReturnToVTPool()
+	}
+	q.returnChunks = q.returnChunks[:0]
+	returnChunksSlicePool.Put(&q.returnChunks)
+
+	return nil
+}

pkg/store/local.go

@@ -18,6 +18,7 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/encoding/protojson"
+	"google.golang.org/protobuf/proto"
 
 	"github.com/thanos-io/thanos/pkg/component"
 	"github.com/thanos-io/thanos/pkg/runutil"
@@ -180,7 +181,7 @@ func (s *LocalStore) Series(r *storepb.SeriesRequest, srv storepb.Store_SeriesSe
 
 		sort.Ints(chosen)
 		for _, ci := range chosen {
-			resp.Chunks = append(resp.Chunks, series.Chunks[ci])
+			resp.Chunks = append(resp.Chunks, proto.Clone(series.Chunks[ci]).(*storepb.AggrChunk))
 		}
 
 		if err := srv.Send(storepb.NewSeriesResponse(resp)); err != nil {

pkg/store/storepb/types.proto

@@ -7,9 +7,10 @@ package thanos;
 option go_package = "github.com/thanos-io/thanos/pkg/store/storepb";
 
 import "store/labelpb/types.proto";
-
+import "github.com/planetscale/vtprotobuf/vtproto/ext.proto";
 
 message Chunk {
+  option (vtproto.mempool) = true;
   enum Encoding {
     XOR = 0;
     HISTOGRAM = 1;
@@ -26,6 +27,7 @@ message Series {
 }
 
 message AggrChunk {
+  option (vtproto.mempool) = true;
   int64 min_time = 1;
   int64 max_time = 2;
 

scripts/genproto.sh

@@ -39,7 +39,7 @@ for dir in ${DIRS}; do
     --go-grpc_out=. --go-grpc_opt=paths=source_relative -I=. -I=${INCLUDE_PATH} -I=${VTPROTOBUF_INCLUDE_PATH} \
     --go_opt=Mstore/storepb/types.proto=github.com/thanos-io/thanos/pkg/store/storepb \
     --go_opt=Mrules/rulespb/rpc.proto=github.com/thanos-io/thanos/pkg/rules/rulespb \
-    --go-vtproto_out=. --go-vtproto_opt=features=marshal+unmarshal+size+equal,paths=source_relative \
+    --go-vtproto_out=. --go-vtproto_opt=features=marshal+unmarshal+size+equal+pool,paths=source_relative \
     --go_opt=Mstore/storepb/prompb/types.proto=github.com/thanos-io/thanos/pkg/store/storepb/prompb \
     --go_opt=Mmetadata/metadatapb/rpc.proto=github.com/thanos-io/thanos/pkg/metadata/metadatapb \
     ${LIST}