Fix a bug where sample in a single sample chunk might get skipped by batch iterator (#4218)

* Fix a bug where sample in a single sample chunk might get skipped by batch iterator

Signed-off-by: Alvin Lin <[email protected]>

* Add reslease notes

Signed-off-by: Alvin Lin <[email protected]>

* Address PR comment

Signed-off-by: Alvin Lin <[email protected]>

* Validate value from iterator in test too

Signed-off-by: Alvin Lin <[email protected]>

* Revert accidental change in ingester.go

Signed-off-by: Alvin Lin <[email protected]>

* Address CR comment

Signed-off-by: Alvin Lin <[email protected]>

* Address PR comments

Signed-off-by: Alvin Lin <[email protected]>

* Refactor optimization a bit such that the batch iterator  don't go 'backwards'

Signed-off-by: Alvin Lin <[email protected]>

Author: alvinlin123

CHANGELOG.md

@@ -52,6 +52,7 @@
 * [BUGFIX] Querier: fix queries failing with "at least 1 healthy replica required, could only find 0" error right after scaling up store-gateways until they're ACTIVE in the ring. #4263
 * [BUGFIX] Store-gateway: when blocks sharding is enabled, do not load all blocks in each store-gateway in case of a cold startup, but load only blocks owned by the store-gateway replica. #4271
 * [BUGFIX] Memberlist: fix to setting the default configuration value for `-memberlist.retransmit-factor` when not provided. This should improve propagation delay of the ring state (including, but not limited to, tombstones). Note that if the configuration is already explicitly given, this fix has no effect. #4269
+* [BUGFIX] Querier: Fix issue where samples in a chunk might get skipped by batch iterator. #4218
 
 ## Blocksconvert
 

pkg/querier/batch/batch.go

@@ -83,13 +83,21 @@ func newIteratorAdapter(underlying iterator) chunkenc.Iterator {
 
 // Seek implements storage.SeriesIterator.
 func (a *iteratorAdapter) Seek(t int64) bool {
-	// Optimisation: see if the seek is within the current batch.
-	if a.curr.Length > 0 && t >= a.curr.Timestamps[0] && t <= a.curr.Timestamps[a.curr.Length-1] {
-		a.curr.Index = 0
-		for a.curr.Index < a.curr.Length && t > a.curr.Timestamps[a.curr.Index] {
-			a.curr.Index++
+
+	// Optimisation: fulfill the seek using current batch if possible.
+	if a.curr.Length > 0 && a.curr.Index < a.curr.Length {
+		if t <= a.curr.Timestamps[a.curr.Index] {
+			//In this case, the interface's requirement is met, so state of this
+			//iterator does not need any change.
+			return true
+		} else if t <= a.curr.Timestamps[a.curr.Length-1] {
+			//In this case, some timestamp between current sample and end of batch can fulfill
+			//the seek. Let's find it.
+			for a.curr.Index < a.curr.Length && t > a.curr.Timestamps[a.curr.Index] {
+				a.curr.Index++
+			}
+			return true
 		}
-		return true
 	}
 
 	a.curr.Length = -1

pkg/querier/batch/batch_test.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/prometheus/common/model"
+	"github.com/stretchr/testify/require"
 
 	"github.com/cortexproject/cortex/pkg/chunk"
 	promchunk "github.com/cortexproject/cortex/pkg/chunk/encoding"
@@ -55,6 +56,22 @@ func BenchmarkNewChunkMergeIterator_CreateAndIterate(b *testing.B) {
 	}
 }
 
+func TestSeekCorrectlyDealWithSinglePointChunks(t *testing.T) {
+	chunkOne := mkChunk(t, model.Time(1*step/time.Millisecond), 1, promchunk.PrometheusXorChunk)
+	chunkTwo := mkChunk(t, model.Time(10*step/time.Millisecond), 1, promchunk.PrometheusXorChunk)
+	chunks := []chunk.Chunk{chunkOne, chunkTwo}
+
+	sut := NewChunkMergeIterator(chunks, 0, 0)
+
+	// Following calls mimics Prometheus's query engine behaviour for VectorSelector.
+	require.True(t, sut.Next())
+	require.True(t, sut.Seek(0))
+
+	actual, val := sut.At()
+	require.Equal(t, float64(1*time.Second/time.Millisecond), val) // since mkChunk use ts as value.
+	require.Equal(t, int64(1*time.Second/time.Millisecond), actual)
+}
+
 func createChunks(b *testing.B, numChunks, numSamplesPerChunk, duplicationFactor int, enc promchunk.Encoding) []chunk.Chunk {
 	result := make([]chunk.Chunk, 0, numChunks)
 

pkg/querier/batch/chunk_test.go

@@ -24,8 +24,11 @@ func TestChunkIter(t *testing.T) {
 	forEncodings(t, func(t *testing.T, enc promchunk.Encoding) {
 		chunk := mkGenericChunk(t, 0, 100, enc)
 		iter := &chunkIterator{}
+
 		iter.reset(chunk)
 		testIter(t, 100, newIteratorAdapter(iter))
+
+		iter.reset(chunk)
 		testSeek(t, 100, newIteratorAdapter(iter))
 	})
 }
@@ -56,7 +59,7 @@ func mkChunk(t require.TestingT, from model.Time, points int, enc promchunk.Enco
 		require.Nil(t, npc)
 		ts = ts.Add(step)
 	}
-	return chunk.NewChunk(userID, fp, metric, pc, model.Time(0), ts)
+	return chunk.NewChunk(userID, fp, metric, pc, from, ts)
 }
 
 func mkGenericChunk(t require.TestingT, from model.Time, points int, enc promchunk.Encoding) GenericChunk {

pkg/querier/batch/merge_test.go

@@ -16,8 +16,11 @@ func TestMergeIter(t *testing.T) {
 		chunk3 := mkGenericChunk(t, model.TimeFromUnix(50), 100, enc)
 		chunk4 := mkGenericChunk(t, model.TimeFromUnix(75), 100, enc)
 		chunk5 := mkGenericChunk(t, model.TimeFromUnix(100), 100, enc)
+
 		iter := newMergeIterator([]GenericChunk{chunk1, chunk2, chunk3, chunk4, chunk5})
 		testIter(t, 200, newIteratorAdapter(iter))
+
+		iter = newMergeIterator([]GenericChunk{chunk1, chunk2, chunk3, chunk4, chunk5})
 		testSeek(t, 200, newIteratorAdapter(iter))
 	})
 }
@@ -37,6 +40,8 @@ func TestMergeHarder(t *testing.T) {
 		}
 		iter := newMergeIterator(chunks)
 		testIter(t, offset*numChunks+samples-offset, newIteratorAdapter(iter))
+
+		iter = newMergeIterator(chunks)
 		testSeek(t, offset*numChunks+samples-offset, newIteratorAdapter(iter))
 	})
 }