From ea085accd6782f23961da11fbd04ff7f6fc96bd9 Mon Sep 17 00:00:00 2001
From: Matthias Loibl <mail@matthiasloibl.com>
Date: Tue, 21 Dec 2021 11:36:54 +0100
Subject: [PATCH] pkg/storage: Fix bug so samples seek to correct index

---
 pkg/storage/series_iterator.go            |  2 +-
 pkg/storage/series_iterator_range.go      | 16 +++++++++---
 pkg/storage/series_iterator_range_test.go | 32 +++++++++++------------
 3 files changed, 29 insertions(+), 21 deletions(-)

diff --git a/pkg/storage/series_iterator.go b/pkg/storage/series_iterator.go
index 4f1e49267bd..d8e581a0ad5 100644
--- a/pkg/storage/series_iterator.go
+++ b/pkg/storage/series_iterator.go
@@ -376,7 +376,7 @@ type MemSeriesInstantFlatProfile struct {
 	durationsIterator  MemSeriesValuesIterator
 	periodsIterator    MemSeriesValuesIterator
 
-	sampleIterators map[string]MemSeriesValuesIterator
+	sampleIterators map[string]*MultiChunksIterator
 	locations       map[string][]*metastore.Location
 }
 
diff --git a/pkg/storage/series_iterator_range.go b/pkg/storage/series_iterator_range.go
index abb5d645718..50ae8ba0d20 100644
--- a/pkg/storage/series_iterator_range.go
+++ b/pkg/storage/series_iterator_range.go
@@ -60,7 +60,7 @@ func (rs *MemRangeSeries) Iterator() ProfileSeriesIterator {
 		rootIt.Seek(start)
 	}
 
-	var sampleIterators map[string]MemSeriesValuesIterator
+	var sampleIterators map[string]*MultiChunksIterator
 
 	root := &MemSeriesIteratorTreeNode{}
 	if rs.trees {
@@ -109,7 +109,7 @@ func (rs *MemRangeSeries) Iterator() ProfileSeriesIterator {
 			memItStack.Pop()
 		}
 	} else {
-		sampleIterators = make(map[string]MemSeriesValuesIterator, len(rs.s.samples))
+		sampleIterators = make(map[string]*MultiChunksIterator, len(rs.s.samples))
 		for key, chunks := range rs.s.samples {
 			sampleIterators[key] = NewMultiChunkIterator(chunks)
 		}
@@ -123,6 +123,12 @@ func (rs *MemRangeSeries) Iterator() ProfileSeriesIterator {
 		timestampIterator.Seek(start)
 		durationsIterator.Seek(start)
 		periodsIterator.Seek(start)
+
+		if !rs.trees {
+			for _, sampleIterator := range sampleIterators {
+				sampleIterator.Seek(start)
+			}
+		}
 	}
 
 	if end-start < numSamples {
@@ -159,7 +165,7 @@ type MemRangeSeriesIterator struct {
 	durationsIterator  MemSeriesValuesIterator
 	periodsIterator    MemSeriesValuesIterator
 
-	sampleIterators map[string]MemSeriesValuesIterator
+	sampleIterators map[string]*MultiChunksIterator
 
 	numSamples uint64 // uint16 might not be enough for many chunks (~500+)
 	err        error
@@ -251,6 +257,10 @@ func (it *MemRangeSeriesIterator) Next() bool {
 				it.err = errors.New("unexpected end of samples iterator")
 				return false
 			}
+			if sread := sit.Read(); sread != read {
+				it.err = fmt.Errorf("sample iterator in wrong iteration, expected %d got %d", read, sread)
+				return false
+			}
 		}
 	}
 
diff --git a/pkg/storage/series_iterator_range_test.go b/pkg/storage/series_iterator_range_test.go
index c6e4126fb3e..2352b3da4c5 100644
--- a/pkg/storage/series_iterator_range_test.go
+++ b/pkg/storage/series_iterator_range_test.go
@@ -18,6 +18,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/prometheus/prometheus/pkg/labels"
 	"github.com/stretchr/testify/require"
 )
@@ -28,39 +29,36 @@ func TestMemRangeSeries_Iterator(t *testing.T) {
 	app, err := s.Appender()
 	require.NoError(t, err)
 
+	s1 := makeSample(2, []uuid.UUID{
+		uuid.MustParse("00000000-0000-0000-0000-000000000002"),
+		uuid.MustParse("00000000-0000-0000-0000-000000000001"),
+	})
+	k1 := makeStacktraceKey(s1)
+
 	for i := 1; i <= 500; i++ {
-		p := Profile{
+		s1.Value = int64(i)
+		p := &FlatProfile{
 			Meta: InstantProfileMeta{
 				Timestamp: int64(i),
 				Duration:  time.Second.Nanoseconds(),
 				Period:    time.Second.Nanoseconds(),
 			},
-			Tree: &ProfileTree{
-				Roots: &ProfileTreeRootNode{
-					ProfileTreeNode: &ProfileTreeNode{
-						flatValues: []*ProfileTreeValueNode{{Value: int64(i)}},
-					},
-				},
+			samples: map[string]*Sample{
+				string(k1): s1,
 			},
 		}
-		err = app.Append(ctx, &p)
+		err = app.AppendFlat(ctx, p)
 		require.NoError(t, err)
 	}
 
-	it := (&MemRangeSeries{s: s, mint: 74, maxt: 420, trees: true}).Iterator()
+	it := (&MemRangeSeries{s: s, mint: 74, maxt: 420}).Iterator()
 
 	seen := int64(75)
 	for it.Next() {
 		p := it.At()
 		require.Equal(t, seen, p.ProfileMeta().Timestamp)
-
-		itt := p.ProfileTree().Iterator()
-		for itt.HasMore() {
-			if itt.NextChild() {
-				require.Equal(t, seen, itt.At().FlatValues()[0].Value)
-				itt.StepInto()
-			}
-			itt.StepUp()
+		for _, sample := range p.Samples() {
+			require.Equal(t, seen, sample.Value)
 		}
 		seen++
 	}