From 735db72a4ba638c87bfbd626b495ad9f6194df5e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Giedrius=20Statkevi=C4=8Dius?=
 <giedrius.statkevicius@vinted.com>
Date: Fri, 20 Sep 2024 13:00:02 +0300
Subject: [PATCH] query/store: memoize PromLabels() call (#7767)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

We use the stringlabels call so some allocations are inevitable but we
can be much smarter about it:

```
func (s *storeSeriesSet) At() (labels.Labels, []*storepb.AggrChunk) {
	return s.series[s.i].PromLabels(), s.series[s.i].Chunks <--- not memoized, new alloc on every At() call; need to memoize because of stringlabel. One alloc is inevitable.
}
```

```
lset, chks := s.SeriesSet.At()
if s.peek == nil {
	s.peek = &Series{Labels: labelpb.PromLabelsToLabelpbLabels(lset), Chunks: chks} <-- converting back to labelpb ?
	continue
}
```

```
if labels.Compare(lset, s.peek.PromLabels()) != 0 { <--- PromLabels() called; we can avoid this call
	s.lset, s.chunks = s.peek.PromLabels(), s.peek.Chunks <- PromLabels() called; we can avoid this
	s.peek = &Series{Labels: labelpb.PromLabelsToLabelpbLabels(lset), Chunks: chks} <--- converting back to labelpb; we can avoid this
	return true
}
```

Signed-off-by: Giedrius Statkevičius <giedrius.statkevicius@vinted.com>
---
 pkg/query/iter.go           | 10 ++++++++--
 pkg/store/storepb/custom.go | 22 ++++++++++++++++------
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/pkg/query/iter.go b/pkg/query/iter.go
index 253b587a95..de3ec28705 100644
--- a/pkg/query/iter.go
+++ b/pkg/query/iter.go
@@ -64,10 +64,12 @@ type storeSeriesSet struct {
 	// TODO(bwplotka): Don't buffer all, we have to buffer single series (to sort and dedup chunks), but nothing more.
 	series []storepb.Series
 	i      int
+
+	promLabels []labels.Labels
 }
 
 func newStoreSeriesSet(s []storepb.Series) *storeSeriesSet {
-	return &storeSeriesSet{series: s, i: -1}
+	return &storeSeriesSet{series: s, i: -1, promLabels: make([]labels.Labels, len(s))}
 }
 
 func (s *storeSeriesSet) Next() bool {
@@ -83,7 +85,11 @@ func (*storeSeriesSet) Err() error {
 }
 
 func (s *storeSeriesSet) At() (labels.Labels, []*storepb.AggrChunk) {
-	return s.series[s.i].PromLabels(), s.series[s.i].Chunks
+	// stringlabels are immutable, so we can cache them.
+	if s.promLabels[s.i].IsEmpty() {
+		s.promLabels[s.i] = s.series[s.i].PromLabels()
+	}
+	return s.promLabels[s.i], s.series[s.i].Chunks
 }
 
 // chunkSeries implements storage.Series for a series on storepb types.
diff --git a/pkg/store/storepb/custom.go b/pkg/store/storepb/custom.go
index 81adbdeabd..ac010a452f 100644
--- a/pkg/store/storepb/custom.go
+++ b/pkg/store/storepb/custom.go
@@ -214,12 +214,18 @@ Outer:
 	return true
 }
 
+type seriesWithLabels struct {
+	*Series
+
+	labels.Labels
+}
+
 // uniqueSeriesSet takes one series set and ensures each iteration contains single, full series.
 type uniqueSeriesSet struct {
 	SeriesSet
 	done bool
 
-	peek *Series
+	peek *seriesWithLabels
 
 	lset   labels.Labels
 	chunks []*AggrChunk
@@ -244,13 +250,17 @@ func (s *uniqueSeriesSet) Next() bool {
 		}
 		lset, chks := s.SeriesSet.At()
 		if s.peek == nil {
-			s.peek = &Series{Labels: labelpb.PromLabelsToLabelpbLabels(lset), Chunks: chks}
+			s.peek = &seriesWithLabels{Series: &Series{
+				Chunks: chks,
+			}, Labels: lset}
 			continue
 		}
 
-		if labels.Compare(lset, s.peek.PromLabels()) != 0 {
-			s.lset, s.chunks = s.peek.PromLabels(), s.peek.Chunks
-			s.peek = &Series{Labels: labelpb.PromLabelsToLabelpbLabels(lset), Chunks: chks}
+		if labels.Compare(lset, s.peek.Labels) != 0 {
+			s.lset, s.chunks = s.peek.Labels, s.peek.Chunks
+			s.peek = &seriesWithLabels{Series: &Series{
+				Chunks: chks,
+			}, Labels: lset}
 			return true
 		}
 
@@ -263,7 +273,7 @@ func (s *uniqueSeriesSet) Next() bool {
 		return false
 	}
 
-	s.lset, s.chunks = s.peek.PromLabels(), s.peek.Chunks
+	s.lset, s.chunks = s.peek.Labels, s.peek.Chunks
 	s.peek = nil
 	return true
 }