ingester/client.QueryStreamResponse: Use memory pooling

Signed-off-by: Arve Knudsen <[email protected]>
grafana · Jan 6, 2025 · a79ef47 · a79ef47
1 parent b1546d6
commit a79ef47
Show file tree

Hide file tree

Showing 30 changed files with 2,349 additions and 671 deletions.
diff --git a/pkg/distributor/distributor.go b/pkg/distributor/distributor.go
@@ -2712,7 +2712,7 @@ func (d *Distributor) MetricsForLabelMatchers(ctx context.Context, from, through
 	}
 	defer func() {
 		for _, resp := range resps {
-			resp.FreeBuffer()
+			resp.Release()
 		}
 	}()
 

diff --git a/pkg/distributor/distributor_test.go b/pkg/distributor/distributor_test.go
@@ -6283,28 +6283,34 @@ func (i *mockIngester) QueryStream(ctx context.Context, req *client.QueryRequest
 
 		if i.disableStreamingResponse || req.StreamingChunksBatchSize == 0 {
 			nonStreamingResponses = append(nonStreamingResponses, &client.QueryStreamResponse{
-				Chunkseries: []client.TimeSeriesChunk{
+				Chunkseries: []client.CustomTimeSeriesChunk{
 					{
-						Labels: ts.Labels,
-						Chunks: wireChunks,
+						TimeSeriesChunk: &client.TimeSeriesChunk{
+							Labels: ts.Labels,
+							Chunks: wireChunks,
+						},
 					},
 				},
 			})
 		} else {
 			streamingLabelResponses = append(streamingLabelResponses, &client.QueryStreamResponse{
-				StreamingSeries: []client.QueryStreamSeries{
+				StreamingSeries: []client.CustomQueryStreamSeries{
 					{
-						Labels:     ts.Labels,
-						ChunkCount: int64(len(wireChunks)),
+						QueryStreamSeries: &client.QueryStreamSeries{
+							Labels:     ts.Labels,
+							ChunkCount: int64(len(wireChunks)),
+						},
 					},
 				},
 			})
 
 			streamingChunkResponses = append(streamingChunkResponses, &client.QueryStreamResponse{
-				StreamingSeriesChunks: []client.QueryStreamSeriesChunks{
+				StreamingSeriesChunks: []client.CustomQueryStreamSeriesChunks{
 					{
-						SeriesIndex: uint64(seriesIndex),
-						Chunks:      wireChunks,
+						QueryStreamSeriesChunks: &client.QueryStreamSeriesChunks{
+							SeriesIndex: uint64(seriesIndex),
+							Chunks:      wireChunks,
+						},
 					},
 				},
 			})
@@ -6384,15 +6390,17 @@ func (i *mockIngester) QueryExemplars(ctx context.Context, req *client.ExemplarQ
 		}
 
 		if len(exemplars) > 0 {
-			res.Timeseries = append(res.Timeseries, mimirpb.TimeSeries{
-				Labels:    series.Labels,
-				Exemplars: exemplars,
+			res.Timeseries = append(res.Timeseries, mimirpb.CustomTimeSeries{
+				TimeSeries: &mimirpb.TimeSeries{
+					Labels:    series.Labels,
+					Exemplars: exemplars,
+				},
 			})
 		}
 	}
 
 	// Sort series by labels because the real ingester returns sorted ones.
-	slices.SortFunc(res.Timeseries, func(a, b mimirpb.TimeSeries) int {
+	slices.SortFunc(res.Timeseries, func(a, b mimirpb.CustomTimeSeries) int {
 		aKey := mimirpb.FromLabelAdaptersToKeyString(a.Labels)
 		bKey := mimirpb.FromLabelAdaptersToKeyString(b.Labels)
 		return strings.Compare(aKey, bKey)
@@ -6424,7 +6432,11 @@ func (i *mockIngester) MetricsForLabelMatchers(ctx context.Context, req *client.
 	for _, matchers := range multiMatchers {
 		for _, ts := range i.timeseries {
 			if match(ts.Labels, matchers) {
-				response.Metric = append(response.Metric, &mimirpb.Metric{Labels: ts.Labels})
+				response.Metric = append(response.Metric, mimirpb.CustomMetric{
+					Metric: &mimirpb.Metric{
+						Labels: ts.Labels,
+					},
+				})
 			}
 		}
 	}

diff --git a/pkg/distributor/query.go b/pkg/distributor/query.go
@@ -66,11 +66,13 @@ func (d *Distributor) QueryExemplars(ctx context.Context, from, to model.Time, m
 		if err != nil {
 			return err
 		}
+		/* TODO: Fix me.
 		defer func() {
 			for _, r := range results {
-				r.FreeBuffer()
+				r.Release()
 			}
 		}()
+		*/
 
 		result = mergeExemplarQueryResponses(results)
 
@@ -180,7 +182,7 @@ func mergeExemplarSets(a, b []mimirpb.Exemplar) []mimirpb.Exemplar {
 
 func mergeExemplarQueryResponses(results []*ingester_client.ExemplarQueryResponse) *ingester_client.ExemplarQueryResponse {
 	var keys []string
-	exemplarResults := make(map[string]mimirpb.TimeSeries)
+	exemplarResults := make(map[string]mimirpb.CustomTimeSeries)
 	for _, r := range results {
 		for _, ts := range r.Timeseries {
 			lbls := mimirpb.FromLabelAdaptersToKeyString(ts.Labels)
@@ -199,7 +201,7 @@ func mergeExemplarQueryResponses(results []*ingester_client.ExemplarQueryRespons
 	// Query results from each ingester were sorted, but are not necessarily still sorted after merging.
 	slices.Sort(keys)
 
-	result := make([]mimirpb.TimeSeries, len(exemplarResults))
+	result := make([]mimirpb.CustomTimeSeries, len(exemplarResults))
 	for i, k := range keys {
 		ts := exemplarResults[k]
 		for i, l := range ts.Labels {
@@ -220,8 +222,8 @@ func mergeExemplarQueryResponses(results []*ingester_client.ExemplarQueryRespons
 
 type ingesterQueryResult struct {
 	// Why retain the batches rather than build a single slice? We don't need a single slice for each ingester, so building a single slice for each ingester is a waste of time.
-	chunkseriesBatches [][]ingester_client.TimeSeriesChunk
-	timeseriesBatches  [][]mimirpb.TimeSeries
+	chunkseriesBatches [][]ingester_client.CustomTimeSeriesChunk
+	timeseriesBatches  [][]mimirpb.CustomTimeSeries
 	streamingSeries    seriesChunksStream
 }
 
@@ -327,8 +329,8 @@ func (d *Distributor) queryIngesterStream(ctx context.Context, replicationSets [
 		queryMetrics.IngesterChunksTotal.Add(float64(totalChunks))
 	}()
 
-	hashToChunkseries := map[string]ingester_client.TimeSeriesChunk{}
-	hashToTimeSeries := map[string]mimirpb.TimeSeries{}
+	hashToChunkseries := map[string]ingester_client.CustomTimeSeriesChunk{}
+	hashToTimeSeries := map[string]mimirpb.CustomTimeSeries{}
 
 	for _, res := range results {
 		// Accumulate any chunk series
@@ -370,8 +372,8 @@ func (d *Distributor) queryIngesterStream(ctx context.Context, replicationSets [
 
 	// Now turn the accumulated maps into slices.
 	resp := ingester_client.CombinedQueryStreamResponse{
-		Chunkseries:     make([]ingester_client.TimeSeriesChunk, 0, len(hashToChunkseries)),
-		Timeseries:      make([]mimirpb.TimeSeries, 0, len(hashToTimeSeries)),
+		Chunkseries:     make([]ingester_client.CustomTimeSeriesChunk, 0, len(hashToChunkseries)),
+		Timeseries:      make([]mimirpb.CustomTimeSeries, 0, len(hashToTimeSeries)),
 		StreamingSeries: mergeSeriesChunkStreams(results, d.estimatedIngestersPerSeries(replicationSets)),
 	}
 	for _, series := range hashToChunkseries {
@@ -395,7 +397,6 @@ func receiveResponse(stream ingester_client.Ingester_QueryStreamClient, streamin
 	if err != nil {
 		return 0, nil, false, err
 	}
-	defer resp.FreeBuffer()
 
 	if len(resp.Timeseries) > 0 {
 		for _, series := range resp.Timeseries {