Fixes race conditions in the batch iterator.

pkg/storage/cache.go → pkg/iter/cache.go

@@ -1,66 +1,56 @@
-package storage
+package iter
 
 import (
-	"github.com/grafana/loki/pkg/iter"
 	"github.com/grafana/loki/pkg/logproto"
 )
 
+type CacheEntryIterator interface {
+	EntryIterator
+	Reset()
+}
+
 // cachedIterator is an iterator that caches iteration to be replayed later on.
 type cachedIterator struct {
-	cache []*logproto.Entry
-	base  iter.EntryIterator
+	cache []entryWithLabels
+	base  EntryIterator // once set to nil it means we have to use the cache.
 
-	labels string
-	curr   int
+	curr int
 
 	closeErr error
 	iterErr  error
 }
 
-// newCachedIterator creates an iterator that cache iteration result and can be iterated again
+// NewCachedIterator creates an iterator that cache iteration result and can be iterated again
 // after closing it without re-using the underlaying iterator `it`.
-// The cache iterator should be used for entries that belongs to the same stream only.
-func newCachedIterator(it iter.EntryIterator, cap int) *cachedIterator {
+func NewCachedIterator(it EntryIterator, cap int) CacheEntryIterator {
 	c := &cachedIterator{
 		base:  it,
-		cache: make([]*logproto.Entry, 0, cap),
+		cache: make([]entryWithLabels, 0, cap),
 		curr:  -1,
 	}
-	c.load()
 	return c
 }
 
-func (it *cachedIterator) reset() {
+func (it *cachedIterator) Reset() {
 	it.curr = -1
 }
 
-func (it *cachedIterator) load() {
+func (it *cachedIterator) Next() bool {
 	if it.base != nil {
-		defer func() {
+		ok := it.base.Next()
+		// we're done with the base iterator.
+		if !ok {
 			it.closeErr = it.base.Close()
 			it.iterErr = it.base.Error()
 			it.base = nil
-			it.reset()
-		}()
-		// set labels using the first entry
-		if !it.base.Next() {
-			return
+			return false
 		}
-		it.labels = it.base.Labels()
-
-		// add all entries until the base iterator is exhausted
-		for {
-			e := it.base.Entry()
-			it.cache = append(it.cache, &e)
-			if !it.base.Next() {
-				break
-			}
-		}
-
+		// we're caching entries
+		it.cache = append(it.cache, entryWithLabels{entry: it.base.Entry(), labels: it.base.Labels()})
+		it.curr++
+		return true
 	}
-}
-
-func (it *cachedIterator) Next() bool {
+	// second pass
 	if len(it.cache) == 0 {
 		it.cache = nil
 		return false
@@ -73,81 +63,74 @@ func (it *cachedIterator) Next() bool {
 }
 
 func (it *cachedIterator) Entry() logproto.Entry {
-	if len(it.cache) == 0 {
+	if len(it.cache) == 0 || it.curr < 0 {
 		return logproto.Entry{}
 	}
-	if it.curr < 0 {
-		return *it.cache[0]
-	}
-	return *it.cache[it.curr]
+
+	return it.cache[it.curr].entry
 }
 
 func (it *cachedIterator) Labels() string {
-	return it.labels
+	if len(it.cache) == 0 || it.curr < 0 {
+		return ""
+	}
+	return it.cache[it.curr].labels
 }
 
 func (it *cachedIterator) Error() error { return it.iterErr }
 
 func (it *cachedIterator) Close() error {
-	it.reset()
+	it.Reset()
 	return it.closeErr
 }
 
+type CacheSampleIterator interface {
+	SampleIterator
+	Reset()
+}
+
 // cachedIterator is an iterator that caches iteration to be replayed later on.
 type cachedSampleIterator struct {
-	cache []logproto.Sample
-	base  iter.SampleIterator
+	cache []sampleWithLabels
+	base  SampleIterator
 
-	labels string
-	curr   int
+	curr int
 
 	closeErr error
 	iterErr  error
 }
 
 // newSampleCachedIterator creates an iterator that cache iteration result and can be iterated again
 // after closing it without re-using the underlaying iterator `it`.
-// The cache iterator should be used for entries that belongs to the same stream only.
-func newCachedSampleIterator(it iter.SampleIterator, cap int) *cachedSampleIterator {
+func NewCachedSampleIterator(it SampleIterator, cap int) CacheSampleIterator {
 	c := &cachedSampleIterator{
 		base:  it,
-		cache: make([]logproto.Sample, 0, cap),
+		cache: make([]sampleWithLabels, 0, cap),
 		curr:  -1,
 	}
-	c.load()
 	return c
 }
 
-func (it *cachedSampleIterator) reset() {
+func (it *cachedSampleIterator) Reset() {
 	it.curr = -1
 }
 
-func (it *cachedSampleIterator) load() {
+func (it *cachedSampleIterator) Next() bool {
 	if it.base != nil {
-		defer func() {
+		ok := it.base.Next()
+		// we're done with the base iterator.
+		if !ok {
 			it.closeErr = it.base.Close()
 			it.iterErr = it.base.Error()
 			it.base = nil
-			it.reset()
-		}()
-		// set labels using the first entry
-		if !it.base.Next() {
-			return
+			return false
 		}
-		it.labels = it.base.Labels()
-
-		// add all entries until the base iterator is exhausted
-		for {
-			it.cache = append(it.cache, it.base.Sample())
-			if !it.base.Next() {
-				break
-			}
-		}
-
+		// we're caching entries
+		it.cache = append(it.cache, sampleWithLabels{Sample: it.base.Sample(), labels: it.base.Labels()})
+		it.curr++
+		return true
 	}
-}
-
-func (it *cachedSampleIterator) Next() bool {
+	// second pass
 	if len(it.cache) == 0 {
 		it.cache = nil
 		return false
@@ -160,22 +143,22 @@ func (it *cachedSampleIterator) Next() bool {
 }
 
 func (it *cachedSampleIterator) Sample() logproto.Sample {
-	if len(it.cache) == 0 {
+	if len(it.cache) == 0 || it.curr < 0 {
 		return logproto.Sample{}
 	}
-	if it.curr < 0 {
-		return it.cache[0]
-	}
-	return it.cache[it.curr]
+	return it.cache[it.curr].Sample
 }
 
 func (it *cachedSampleIterator) Labels() string {
-	return it.labels
+	if len(it.cache) == 0 || it.curr < 0 {
+		return ""
+	}
+	return it.cache[it.curr].labels
 }
 
 func (it *cachedSampleIterator) Error() error { return it.iterErr }
 
 func (it *cachedSampleIterator) Close() error {
-	it.reset()
+	it.Reset()
 	return it.closeErr
 }

pkg/storage/cache_test.go → pkg/iter/cache_test.go

@@ -1,4 +1,4 @@
-package storage
+package iter
 
 import (
 	"errors"
@@ -7,7 +7,6 @@ import (
 
 	"github.com/stretchr/testify/require"
 
-	"github.com/grafana/loki/pkg/iter"
 	"github.com/grafana/loki/pkg/logproto"
 )
 
@@ -20,12 +19,11 @@ func Test_CachedIterator(t *testing.T) {
 			{Timestamp: time.Unix(0, 3), Line: "3"},
 		},
 	}
-	c := newCachedIterator(iter.NewStreamIterator(stream), 3)
+	c := NewCachedIterator(NewStreamIterator(stream), 3)
 
 	assert := func() {
-		// we should crash for call of entry without next although that's not expected.
-		require.Equal(t, stream.Labels, c.Labels())
-		require.Equal(t, stream.Entries[0], c.Entry())
+		require.Equal(t, "", c.Labels())
+		require.Equal(t, logproto.Entry{}, c.Entry())
 		require.Equal(t, true, c.Next())
 		require.Equal(t, stream.Entries[0], c.Entry())
 		require.Equal(t, true, c.Next())
@@ -48,7 +46,7 @@ func Test_CachedIterator(t *testing.T) {
 
 func Test_EmptyCachedIterator(t *testing.T) {
 
-	c := newCachedIterator(iter.NoopIterator, 0)
+	c := NewCachedIterator(NoopIterator, 0)
 
 	require.Equal(t, "", c.Labels())
 	require.Equal(t, logproto.Entry{}, c.Entry())
@@ -68,7 +66,7 @@ func Test_EmptyCachedIterator(t *testing.T) {
 
 func Test_ErrorCachedIterator(t *testing.T) {
 
-	c := newCachedIterator(&errorIter{}, 0)
+	c := NewCachedIterator(&errorIter{}, 0)
 
 	require.Equal(t, false, c.Next())
 	require.Equal(t, "", c.Labels())
@@ -86,12 +84,11 @@ func Test_CachedSampleIterator(t *testing.T) {
 			{Timestamp: time.Unix(0, 3).UnixNano(), Hash: 3, Value: 3.},
 		},
 	}
-	c := newCachedSampleIterator(iter.NewSeriesIterator(series), 3)
+	c := NewCachedSampleIterator(NewSeriesIterator(series), 3)
 
 	assert := func() {
-		// we should crash for call of entry without next although that's not expected.
-		require.Equal(t, series.Labels, c.Labels())
-		require.Equal(t, series.Samples[0], c.Sample())
+		require.Equal(t, "", c.Labels())
+		require.Equal(t, logproto.Sample{}, c.Sample())
 		require.Equal(t, true, c.Next())
 		require.Equal(t, series.Samples[0], c.Sample())
 		require.Equal(t, true, c.Next())
@@ -114,7 +111,7 @@ func Test_CachedSampleIterator(t *testing.T) {
 
 func Test_EmptyCachedSampleIterator(t *testing.T) {
 
-	c := newCachedSampleIterator(iter.NoopIterator, 0)
+	c := NewCachedSampleIterator(NoopIterator, 0)
 
 	require.Equal(t, "", c.Labels())
 	require.Equal(t, logproto.Sample{}, c.Sample())
@@ -134,7 +131,7 @@ func Test_EmptyCachedSampleIterator(t *testing.T) {
 
 func Test_ErrorCachedSampleIterator(t *testing.T) {
 
-	c := newCachedSampleIterator(&errorIter{}, 0)
+	c := NewCachedSampleIterator(&errorIter{}, 0)
 
 	require.Equal(t, false, c.Next())
 	require.Equal(t, "", c.Labels())