Cherry-pick upstream goroutine leak fixes from OSS Thanos (#245)

Author: yuchen-db

CHANGELOG.md

@@ -12,6 +12,8 @@ We use *breaking :warning:* to mark changes that are not backward compatible (re
 
 ### Fixed
 
+- [#8378](https://github.com/thanos-io/thanos/pull/8378): Store: fix the reuse of dirty posting slices
+
 ### Added
 
 ### Changed

pkg/store/bucket.go

@@ -34,6 +34,7 @@ import (
 	"github.com/prometheus/prometheus/tsdb/chunks"
 	"github.com/prometheus/prometheus/tsdb/encoding"
 	"github.com/prometheus/prometheus/tsdb/index"
+	"github.com/prometheus/prometheus/util/zeropool"
 	"github.com/weaveworks/common/httpgrpc"
 	"golang.org/x/exp/slices"
 	"golang.org/x/sync/errgroup"
@@ -125,8 +126,22 @@ const (
 var (
 	errBlockSyncConcurrencyNotValid = errors.New("the block sync concurrency must be equal or greater than 1.")
 	hashPool                        = sync.Pool{New: func() interface{} { return xxhash.New() }}
+	postingsPool                    zeropool.Pool[[]storage.SeriesRef]
 )
 
+func getPostingsSlice() []storage.SeriesRef {
+	if p := postingsPool.Get(); p != nil {
+		return p
+	}
+
+	// Pre-allocate slice with initial capacity.
+	return make([]storage.SeriesRef, 0, 1024)
+}
+
+func putPostingsSlice(p []storage.SeriesRef) {
+	postingsPool.Put(p[:0])
+}
+
 type bucketStoreMetrics struct {
 	blocksLoaded          prometheus.Gauge
 	blockLoads            prometheus.Counter
@@ -1693,6 +1708,9 @@ func (s *BucketStore) Series(req *storepb.SeriesRequest, seriesSrv storepb.Store
 			err = g.Wait()
 		})
 		if err != nil {
+			for _, resp := range respSets {
+				resp.Close()
+			}
 			code := codes.Aborted
 			if s, ok := status.FromError(errors.Cause(err)); ok {
 				code = s.Code()
@@ -2539,6 +2557,10 @@ type bucketIndexReader struct {
 
 	indexVersion int
 	logger       log.Logger
+
+	// Posting slice to return to the postings pool on close.
+	// A single bucketIndexReader should have at most 1 postings slice to return.
+	postings []storage.SeriesRef
 }
 
 func newBucketIndexReader(block *bucketBlock, logger log.Logger) *bucketIndexReader {
@@ -2659,13 +2681,13 @@ func (r *bucketIndexReader) ExpandedPostings(ctx context.Context, ms sortedMatch
 
 // ExpandPostingsWithContext returns the postings expanded as a slice and considers context.
 func ExpandPostingsWithContext(ctx context.Context, p index.Postings) ([]storage.SeriesRef, error) {
-	res := make([]storage.SeriesRef, 0, 1024) // Pre-allocate slice with initial capacity
+	res := getPostingsSlice()
 	i := 0
 	for p.Next() {
 		i++
 		if i%checkContextEveryNIterations == 0 {
 			if err := ctx.Err(); err != nil {
-				return nil, err
+				return res, err
 			}
 		}
 		res = append(res, p.At())
@@ -2958,6 +2980,7 @@ func (r *bucketIndexReader) fetchExpandedPostingsFromCache(ctx context.Context,
 	}
 
 	ps, err := ExpandPostingsWithContext(ctx, p)
+	r.postings = ps
 	if err != nil {
 		level.Error(r.logger).Log("msg", "failed to expand cached expanded postings, refetch postings", "id", r.block.meta.ULID.String(), "err", err)
 		return false, nil, nil
@@ -3394,6 +3417,10 @@ func (r *bucketIndexReader) LoadSeriesForTime(ref storage.SeriesRef, lset *[]sym
 // Close released the underlying resources of the reader.
 func (r *bucketIndexReader) Close() error {
 	r.block.pendingReaders.Done()
+
+	if r.postings != nil {
+		putPostingsSlice(r.postings)
+	}
 	return nil
 }
 

pkg/store/bucket_e2e_test.go

@@ -131,7 +131,7 @@ func prepareTestBlocks(t testing.TB, now time.Time, count int, dir string, bkt o
 	return
 }
 
-func prepareStoreWithTestBlocks(t testing.TB, dir string, bkt objstore.Bucket, manyParts bool, chunksLimiterFactory ChunksLimiterFactory, seriesLimiterFactory SeriesLimiterFactory, bytesLimiterFactory BytesLimiterFactory, relabelConfig []*relabel.Config, filterConf *FilterConfig) *storeSuite {
+func prepareStoreWithTestBlocks(t testing.TB, dir string, bkt objstore.Bucket, manyParts bool, chunksLimiterFactory ChunksLimiterFactory, seriesLimiterFactory SeriesLimiterFactory, bytesLimiterFactory BytesLimiterFactory, relabelConfig []*relabel.Config, filterConf *FilterConfig, opts ...BucketStoreOption) *storeSuite {
 	series := []labels.Labels{
 		labels.FromStrings("a", "1", "b", "1"),
 		labels.FromStrings("a", "1", "b", "2"),
@@ -176,10 +176,10 @@ func prepareStoreWithTestBlocks(t testing.TB, dir string, bkt objstore.Bucket, m
 		true,
 		true,
 		time.Minute,
-		WithLogger(s.logger),
-		WithIndexCache(s.cache),
-		WithFilterConfig(filterConf),
-		WithRegistry(reg),
+		append(opts, WithLogger(s.logger),
+			WithIndexCache(s.cache),
+			WithFilterConfig(filterConf),
+			WithRegistry(reg))...,
 	)
 	testutil.Ok(t, err)
 	defer func() { testutil.Ok(t, store.Close()) }()
@@ -619,6 +619,7 @@ func TestBucketStore_Series_ChunksLimiter_e2e(t *testing.T) {
 		maxChunksLimit uint64
 		maxSeriesLimit uint64
 		maxBytesLimit  int64
+		storeOpts      []BucketStoreOption
 		expectedErr    string
 		code           codes.Code
 	}{
@@ -630,19 +631,40 @@ func TestBucketStore_Series_ChunksLimiter_e2e(t *testing.T) {
 			expectedErr:    "exceeded chunks limit",
 			code:           codes.ResourceExhausted,
 		},
+		"should fail if the max chunks limit is exceeded - ResourceExhausted (sortingStrategyNone)": {
+			maxChunksLimit: expectedChunks - 1,
+			expectedErr:    "exceeded chunks limit",
+			storeOpts:      []BucketStoreOption{WithDontResort(true)},
+			code:           codes.ResourceExhausted,
+		},
 		"should fail if the max series limit is exceeded - ResourceExhausted": {
 			maxChunksLimit: expectedChunks,
 			expectedErr:    "exceeded series limit",
 			maxSeriesLimit: 1,
 			code:           codes.ResourceExhausted,
 		},
+		"should fail if the max series limit is exceeded - ResourceExhausted (sortingStrategyNone)": {
+			maxChunksLimit: expectedChunks,
+			expectedErr:    "exceeded series limit",
+			maxSeriesLimit: 1,
+			storeOpts:      []BucketStoreOption{WithDontResort(true)},
+			code:           codes.ResourceExhausted,
+		},
 		"should fail if the max bytes limit is exceeded - ResourceExhausted": {
 			maxChunksLimit: expectedChunks,
 			expectedErr:    "exceeded bytes limit",
 			maxSeriesLimit: 2,
 			maxBytesLimit:  1,
 			code:           codes.ResourceExhausted,
 		},
+		"should fail if the max bytes limit is exceeded - ResourceExhausted (sortingStrategyNone)": {
+			maxChunksLimit: expectedChunks,
+			expectedErr:    "exceeded bytes limit",
+			maxSeriesLimit: 2,
+			maxBytesLimit:  1,
+			storeOpts:      []BucketStoreOption{WithDontResort(true)},
+			code:           codes.ResourceExhausted,
+		},
 	}
 
 	for testName, testData := range cases {
@@ -653,7 +675,7 @@ func TestBucketStore_Series_ChunksLimiter_e2e(t *testing.T) {
 
 			dir := t.TempDir()
 
-			s := prepareStoreWithTestBlocks(t, dir, bkt, false, NewChunksLimiterFactory(testData.maxChunksLimit), NewSeriesLimiterFactory(testData.maxSeriesLimit), NewBytesLimiterFactory(units.Base2Bytes(testData.maxBytesLimit)), emptyRelabelConfig, allowAllFilterConf)
+			s := prepareStoreWithTestBlocks(t, dir, bkt, false, NewChunksLimiterFactory(testData.maxChunksLimit), NewSeriesLimiterFactory(testData.maxSeriesLimit), NewBytesLimiterFactory(units.Base2Bytes(testData.maxBytesLimit)), emptyRelabelConfig, allowAllFilterConf, testData.storeOpts...)
 			testutil.Ok(t, s.store.SyncBlocks(ctx))
 
 			req := &storepb.SeriesRequest{

pkg/store/bucket_test.go

@@ -2971,7 +2971,7 @@ func TestExpandPostingsWithContextCancel(t *testing.T) {
 	res, err := ExpandPostingsWithContext(ctx, p)
 	testutil.NotOk(t, err)
 	testutil.Equals(t, context.Canceled, err)
-	testutil.Equals(t, []storage.SeriesRef(nil), res)
+	testutil.Equals(t, true, cap(res) == 1024)
 }
 
 func samePostingGroup(a, b *postingGroup) bool {

pkg/store/lazy_postings.go

@@ -279,6 +279,19 @@ func fetchAndExpandPostingGroups(ctx context.Context, r *bucketIndexReader, post
 		return nil, nil, errors.Wrap(err, "get postings")
 	}
 
+	result := mergeFetchedPostings(ctx, fetchedPostings, postingGroups)
+	if err := ctx.Err(); err != nil {
+		return nil, nil, err
+	}
+	ps, err := ExpandPostingsWithContext(ctx, result)
+	r.postings = ps
+	if err != nil {
+		return nil, nil, errors.Wrap(err, "expand")
+	}
+	return ps, lazyMatchers, nil
+}
+
+func mergeFetchedPostings(ctx context.Context, fetchedPostings []index.Postings, postingGroups []*postingGroup) index.Postings {
 	// Get "add" and "remove" postings from groups. We iterate over postingGroups and their keys
 	// again, and this is exactly the same order as before (when building the groups), so we can simply
 	// use one incrementing index to fetch postings from returned slice.
@@ -306,14 +319,5 @@ func fetchAndExpandPostingGroups(ctx context.Context, r *bucketIndexReader, post
 		}
 	}
 
-	result := index.Without(index.Intersect(groupAdds...), index.Merge(ctx, groupRemovals...))
-
-	if err := ctx.Err(); err != nil {
-		return nil, nil, err
-	}
-	ps, err := ExpandPostingsWithContext(ctx, result)
-	if err != nil {
-		return nil, nil, errors.Wrap(err, "expand")
-	}
-	return ps, lazyMatchers, nil
+	return index.Without(index.Intersect(groupAdds...), index.Merge(ctx, groupRemovals...))
 }

pkg/store/proxy_merge.go

@@ -19,6 +19,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/prometheus/model/labels"
+
 	grpc_opentracing "github.com/thanos-io/thanos/pkg/tracing/tracing_middleware"
 
 	"github.com/thanos-io/thanos/pkg/losertree"
@@ -277,6 +278,8 @@ type lazyRespSet struct {
 	initialized bool
 
 	shardMatcher *storepb.ShardMatcher
+
+	donec chan struct{}
 }
 
 func (l *lazyRespSet) isEmpty() bool {
@@ -385,6 +388,7 @@ func newLazyRespSet(
 		ringHead:             0,
 		ringTail:             0,
 		closed:               false,
+		donec:                make(chan struct{}),
 	}
 	respSet.storeLabels = make(map[string]struct{})
 	for _, ls := range storeLabelSets {
@@ -403,6 +407,8 @@ func newLazyRespSet(
 			l.span.SetTag("processed.samples", seriesStats.Samples)
 			l.span.SetTag("processed.bytes", bytesProcessed)
 			l.span.Finish()
+
+			close(l.donec)
 		}()
 
 		numResponses := 0
@@ -611,13 +617,14 @@ func newAsyncRespSet(
 
 func (l *lazyRespSet) Close() {
 	l.bufferedResponsesMtx.Lock()
-	defer l.bufferedResponsesMtx.Unlock()
-
 	l.closeSeries()
 	l.closed = true
 	l.bufferSlotEvent.Signal()
 	l.noMoreData = true
 	l.dataOrFinishEvent.Signal()
+	l.bufferedResponsesMtx.Unlock()
+
+	<-l.donec // Wait for the internal goroutine to complete its work.
 
 	l.shardMatcher.Close()
 	_ = l.cl.CloseSend()
@@ -806,11 +813,15 @@ func sortWithoutLabels(set []*storepb.SeriesResponse, labelsToRemove map[string]
 }
 
 func (l *eagerRespSet) Close() {
+	l.wg.Wait()
+
 	if l.closeSeries != nil {
 		l.closeSeries()
 	}
+	l.wg.Wait()
 	l.shardMatcher.Close()
 	_ = l.cl.CloseSend()
+
 }
 
 func (l *eagerRespSet) At() *storepb.SeriesResponse {

pkg/store/proxy_test.go

@@ -2055,10 +2055,6 @@ func TestProxyStore_SeriesSlowStores(t *testing.T) {
 			return
 		}
 	}
-
-	// Wait until the last goroutine exits which is stuck on time.Sleep().
-	// Otherwise, goleak complains.
-	time.Sleep(5 * time.Second)
 }
 
 func TestProxyStore_Series_RequestParamsProxied(t *testing.T) {