Loki: Fix bug where items are returned to a sync.Pool incorrectly

pkg/iter/sample_iterator.go

@@ -4,6 +4,7 @@ import (
 	"container/heap"
 	"context"
 	"io"
+	"sync"
 
 	"github.com/grafana/loki/pkg/logproto"
 	"github.com/grafana/loki/pkg/logqlmodel/stats"
@@ -344,26 +345,31 @@ type seriesIterator struct {
 }
 
 type withCloseSampleIterator struct {
-	closeFn func() error
+	closeOnce sync.Once
+	closeFn   func() error
+	errs      []error
 	SampleIterator
 }
 
 func (w *withCloseSampleIterator) Close() error {
-	var errs []error
-	if err := w.SampleIterator.Close(); err != nil {
-		errs = append(errs, err)
-	}
-	if err := w.closeFn(); err != nil {
-		errs = append(errs, err)
-	}
-	if len(errs) == 0 {
+	w.closeOnce.Do(func() {
+		if err := w.SampleIterator.Close(); err != nil {
+			w.errs = append(w.errs, err)
+		}
+		if err := w.closeFn(); err != nil {
+			w.errs = append(w.errs, err)
+		}
+
+	})
+	if len(w.errs) == 0 {
 		return nil
 	}
-	return util.MultiError(errs)
+	return util.MultiError(w.errs)
 }
 
 func SampleIteratorWithClose(it SampleIterator, closeFn func() error) SampleIterator {
 	return &withCloseSampleIterator{
+		closeOnce:      sync.Once{},
 		closeFn:        closeFn,
 		SampleIterator: it,
 	}

pkg/iter/sample_iterator_test.go

@@ -6,10 +6,13 @@ import (
 	"testing"
 	"time"
 
+	"github.com/pkg/errors"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/atomic"
 
 	"github.com/grafana/loki/pkg/logproto"
+	"github.com/grafana/loki/pkg/util"
 )
 
 func TestNewPeekingSampleIterator(t *testing.T) {
@@ -223,3 +226,51 @@ func TestNonOverlappingSampleClose(t *testing.T) {
 	require.Equal(t, true, a.closed.Load())
 	require.Equal(t, true, b.closed.Load())
 }
+
+func TestSampleIteratorWithClose_CloseIdempotent(t *testing.T) {
+	c := 0
+	closeFn := func() error {
+		c++
+		return nil
+	}
+	ni := noOpIterator{}
+	it := SampleIteratorWithClose(ni, closeFn)
+	// Multiple calls to close should result in c only ever having been incremented one time from 0 to 1
+	err := it.Close()
+	assert.NoError(t, err)
+	assert.EqualValues(t, 1, c)
+	err = it.Close()
+	assert.NoError(t, err)
+	assert.EqualValues(t, 1, c)
+	err = it.Close()
+	assert.NoError(t, err)
+	assert.EqualValues(t, 1, c)
+}
+
+type alwaysErrorIterator struct {
+	noOpIterator
+}
+
+func (alwaysErrorIterator) Close() error {
+	return errors.New("i always error")
+}
+
+func TestSampleIteratorWithClose_ReturnsError(t *testing.T) {
+	closeFn := func() error {
+		return errors.New("i broke")
+	}
+	ei := alwaysErrorIterator{}
+	it := SampleIteratorWithClose(ei, closeFn)
+	err := it.Close()
+	// Verify that a proper multi error is returned when both the iterator and the close function return errors
+	if me, ok := err.(util.MultiError); ok {
+		assert.True(t, len(me) == 2, "Expected 2 errors, one from the iterator and one from the close function")
+		assert.EqualError(t, me[0], "i always error")
+		assert.EqualError(t, me[1], "i broke")
+	} else {
+		t.Error("Expected returned error to be of type util.MultiError")
+	}
+	// A second call to Close should return the same error
+	err2 := it.Close()
+	assert.Equal(t, err, err2)
+}