Limit lazyRespSet memory buffer size

Addressed comments

Fix linter

goformat

Fix unit tests

Author: hczhu-db

cmd/thanos/query.go

@@ -248,6 +248,9 @@ func registerQuery(app *extkingpin.App) {
 	rewriteAggregationLabelStrategy := cmd.Flag("query.aggregation-label-strategy", "The strategy to use when rewriting aggregation labels. Used during aggregator migration only.").Default(string(query.NoopLabelRewriter)).Hidden().Enum(string(query.NoopLabelRewriter), string(query.UpsertLabelRewriter), string(query.InsertOnlyLabelRewriter))
 	rewriteAggregationLabelTo := cmd.Flag("query.aggregation-label-value-override", "The value override for aggregation label. If set to x, all queries on aggregated metrics will have a `__agg_rule_type__=x` matcher. If empty, this behavior is disabled. Default is empty.").Hidden().Default("").String()
 
+	lazyRetrievalMaxBufferedResponses := cmd.Flag("query.lazy-retrieval-max-buffered-responses", "The lazy retrieval strategy can buffer up to this number of responses. This is to limit the memory usage. This flag takes effect only when the lazy retrieval strategy is enabled.").
+		Default("20").Int()
+
 	var storeRateLimits store.SeriesSelectLimits
 	storeRateLimits.RegisterFlags(cmd)
 
@@ -389,6 +392,7 @@ func registerQuery(app *extkingpin.App) {
 			*enableGroupReplicaPartialStrategy,
 			*rewriteAggregationLabelStrategy,
 			*rewriteAggregationLabelTo,
+			*lazyRetrievalMaxBufferedResponses,
 		)
 	})
 }
@@ -475,6 +479,7 @@ func runQuery(
 	groupReplicaPartialResponseStrategy bool,
 	rewriteAggregationLabelStrategy string,
 	rewriteAggregationLabelTo string,
+	lazyRetrievalMaxBufferedResponses int,
 ) error {
 	comp := component.Query
 	if alertQueryURL == "" {
@@ -562,6 +567,7 @@ func runQuery(
 		store.WithTSDBSelector(tsdbSelector),
 		store.WithProxyStoreDebugLogging(debugLogging),
 		store.WithQuorumChunkDedup(queryDeduplicationFunc == dedup.AlgorithmQuorum),
+		store.WithLazyRetrievalMaxBufferedResponsesForProxy(lazyRetrievalMaxBufferedResponses),
 	}
 
 	// Parse and sanitize the provided replica labels flags.

pkg/store/bucket.go

@@ -422,6 +422,8 @@ type BucketStore struct {
 	enabledLazyExpandedPostings bool
 
 	sortingStrategy sortingStrategy
+	// This flag limits memory usage when lazy retrieval strategy, newLazyRespSet(), is used.
+	lazyRetrievalMaxBufferedResponses int
 
 	blockEstimatedMaxSeriesFunc BlockEstimator
 	blockEstimatedMaxChunkFunc  BlockEstimator
@@ -561,6 +563,14 @@ func WithDontResort(true bool) BucketStoreOption {
 	}
 }
 
+func WithLazyRetrievalMaxBufferedResponsesForBucket(n int) BucketStoreOption {
+	return func(s *BucketStore) {
+		if true {
+			s.lazyRetrievalMaxBufferedResponses = n
+		}
+	}
+}
+
 // WithIndexHeaderLazyDownloadStrategy specifies what block to lazy download its index header.
 // Only used when lazy mmap is enabled at the same time.
 func WithIndexHeaderLazyDownloadStrategy(strategy indexheader.LazyDownloadIndexHeaderFunc) BucketStoreOption {
@@ -597,23 +607,24 @@ func NewBucketStore(
 			b := make([]byte, 0, initialBufSize)
 			return &b
 		}},
-		chunkPool:                       pool.NoopPool[byte]{},
-		blocks:                          map[ulid.ULID]*bucketBlock{},
-		blockSets:                       map[uint64]*bucketBlockSet{},
-		blockSyncConcurrency:            blockSyncConcurrency,
-		queryGate:                       gate.NewNoop(),
-		chunksLimiterFactory:            chunksLimiterFactory,
-		seriesLimiterFactory:            seriesLimiterFactory,
-		bytesLimiterFactory:             bytesLimiterFactory,
-		partitioner:                     partitioner,
-		enableCompatibilityLabel:        enableCompatibilityLabel,
-		postingOffsetsInMemSampling:     postingOffsetsInMemSampling,
-		enableSeriesResponseHints:       enableSeriesResponseHints,
-		enableChunkHashCalculation:      enableChunkHashCalculation,
-		seriesBatchSize:                 SeriesBatchSize,
-		sortingStrategy:                 sortingStrategyStore,
-		indexHeaderLazyDownloadStrategy: indexheader.AlwaysEagerDownloadIndexHeader,
-		requestLoggerFunc:               NoopRequestLoggerFunc,
+		chunkPool:                         pool.NoopPool[byte]{},
+		blocks:                            map[ulid.ULID]*bucketBlock{},
+		blockSets:                         map[uint64]*bucketBlockSet{},
+		blockSyncConcurrency:              blockSyncConcurrency,
+		queryGate:                         gate.NewNoop(),
+		chunksLimiterFactory:              chunksLimiterFactory,
+		seriesLimiterFactory:              seriesLimiterFactory,
+		bytesLimiterFactory:               bytesLimiterFactory,
+		partitioner:                       partitioner,
+		enableCompatibilityLabel:          enableCompatibilityLabel,
+		postingOffsetsInMemSampling:       postingOffsetsInMemSampling,
+		enableSeriesResponseHints:         enableSeriesResponseHints,
+		enableChunkHashCalculation:        enableChunkHashCalculation,
+		seriesBatchSize:                   SeriesBatchSize,
+		sortingStrategy:                   sortingStrategyStore,
+		lazyRetrievalMaxBufferedResponses: 1,
+		indexHeaderLazyDownloadStrategy:   indexheader.AlwaysEagerDownloadIndexHeader,
+		requestLoggerFunc:                 NoopRequestLoggerFunc,
 	}
 
 	for _, option := range options {
@@ -1613,6 +1624,11 @@ func (s *BucketStore) Series(req *storepb.SeriesRequest, seriesSrv storepb.Store
 						nil,
 					)
 				} else {
+					lazyRetrievalMaxBufferedResponses := s.lazyRetrievalMaxBufferedResponses
+					if lazyRetrievalMaxBufferedResponses < 1 {
+						// Some unit and e2e tests hit this path.
+						lazyRetrievalMaxBufferedResponses = 1
+					}
 					resp = newLazyRespSet(
 						span,
 						10*time.Minute,
@@ -1623,6 +1639,7 @@ func (s *BucketStore) Series(req *storepb.SeriesRequest, seriesSrv storepb.Store
 						shardMatcher,
 						false,
 						s.metrics.emptyPostingCount.WithLabelValues(tenant),
+						lazyRetrievalMaxBufferedResponses,
 					)
 				}
 

pkg/store/bucket_test.go

@@ -1780,12 +1780,13 @@ func TestBucketSeries_OneBlock_InMemIndexCacheSegfault(t *testing.T) {
 			b1.meta.ULID: b1,
 			b2.meta.ULID: b2,
 		},
-		queryGate:            gate.NewNoop(),
-		chunksLimiterFactory: NewChunksLimiterFactory(0),
-		seriesLimiterFactory: NewSeriesLimiterFactory(0),
-		bytesLimiterFactory:  NewBytesLimiterFactory(0),
-		seriesBatchSize:      SeriesBatchSize,
-		requestLoggerFunc:    NoopRequestLoggerFunc,
+		queryGate:                         gate.NewNoop(),
+		chunksLimiterFactory:              NewChunksLimiterFactory(0),
+		seriesLimiterFactory:              NewSeriesLimiterFactory(0),
+		bytesLimiterFactory:               NewBytesLimiterFactory(0),
+		seriesBatchSize:                   SeriesBatchSize,
+		requestLoggerFunc:                 NoopRequestLoggerFunc,
+		lazyRetrievalMaxBufferedResponses: 1,
 	}
 
 	t.Run("invoke series for one block. Fill the cache on the way.", func(t *testing.T) {

pkg/store/proxy.go

@@ -93,14 +93,15 @@ type ProxyStore struct {
 	selectorLabels labels.Labels
 	buffers        sync.Pool
 
-	responseTimeout   time.Duration
-	metrics           *proxyStoreMetrics
-	retrievalStrategy RetrievalStrategy
-	debugLogging      bool
-	tsdbSelector      *TSDBSelector
-	quorumChunkDedup  bool
-	enableDedup       bool
-	matcherConverter  *storepb.MatcherConverter
+	responseTimeout                   time.Duration
+	metrics                           *proxyStoreMetrics
+	retrievalStrategy                 RetrievalStrategy
+	debugLogging                      bool
+	tsdbSelector                      *TSDBSelector
+	quorumChunkDedup                  bool
+	enableDedup                       bool
+	matcherConverter                  *storepb.MatcherConverter
+	lazyRetrievalMaxBufferedResponses int
 }
 
 type proxyStoreMetrics struct {
@@ -137,6 +138,12 @@ func RegisterStoreServer(storeSrv storepb.StoreServer, logger log.Logger) func(*
 // BucketStoreOption are functions that configure BucketStore.
 type ProxyStoreOption func(s *ProxyStore)
 
+func WithLazyRetrievalMaxBufferedResponsesForProxy(buferSize int) ProxyStoreOption {
+	return func(s *ProxyStore) {
+		s.lazyRetrievalMaxBufferedResponses = buferSize
+	}
+}
+
 // WithProxyStoreDebugLogging toggles debug logging.
 func WithProxyStoreDebugLogging(enable bool) ProxyStoreOption {
 	return func(s *ProxyStore) {
@@ -390,11 +397,10 @@ func (s *ProxyStore) Series(originalRequest *storepb.SeriesRequest, srv storepb.
 		}
 	}
 	defer logGroupReplicaErrors()
-
 	for _, st := range stores {
 		st := st
 
-		respSet, err := newAsyncRespSet(ctx, st, r, s.responseTimeout, s.retrievalStrategy, &s.buffers, r.ShardInfo, reqLogger, s.metrics.emptyStreamResponses)
+		respSet, err := newAsyncRespSet(ctx, st, r, s.responseTimeout, s.retrievalStrategy, &s.buffers, r.ShardInfo, reqLogger, s.metrics.emptyStreamResponses, s.lazyRetrievalMaxBufferedResponses)
 		if err != nil {
 			level.Warn(s.logger).Log("msg", "Store failure", "group", st.GroupKey(), "replica", st.ReplicaKey(), "err", err)
 			s.metrics.storeFailureCount.WithLabelValues(st.GroupKey(), st.ReplicaKey()).Inc()

pkg/store/proxy_merge.go

@@ -255,8 +255,20 @@ type lazyRespSet struct {
 	frameTimeout   time.Duration
 
 	// Internal bookkeeping.
-	dataOrFinishEvent    *sync.Cond
-	bufferedResponses    []*storepb.SeriesResponse
+	dataOrFinishEvent *sync.Cond
+	// This event firing means the buffer has a slot for more data.
+	bufferSlotEvent *sync.Cond
+	fixedBufferSize int
+	// This a ring buffer of size fixedBufferSize.
+	// A ring buffer of size N can hold N - 1 elements at most in order to distinguish being empty from being full.
+	bufferedResponses []*storepb.SeriesResponse
+	// ringHead points to the first element in the ring buffer.
+	// ringTail points to the slot after the last element in the ring buffer.
+	// if ringHead == ringTail then the buffer is empty.
+	// if ringHead == (ringTail + 1) % fixedBufferSize then the buffer is full.
+	ringHead             int
+	ringTail             int
+	closed               bool
 	bufferedResponsesMtx *sync.Mutex
 	lastResp             *storepb.SeriesResponse
 
@@ -266,24 +278,32 @@ type lazyRespSet struct {
 	shardMatcher *storepb.ShardMatcher
 }
 
+func (l *lazyRespSet) isEmpty() bool {
+	return l.ringHead == l.ringTail
+}
+
+func (l *lazyRespSet) isFull() bool {
+	return (l.ringTail+1)%l.fixedBufferSize == l.ringHead
+}
+
 func (l *lazyRespSet) Empty() bool {
 	l.bufferedResponsesMtx.Lock()
 	defer l.bufferedResponsesMtx.Unlock()
 
 	// NOTE(GiedriusS): need to wait here for at least one
 	// response so that we could build the heap properly.
-	if l.noMoreData && len(l.bufferedResponses) == 0 {
+	if l.noMoreData && l.isEmpty() {
 		return true
 	}
 
-	for len(l.bufferedResponses) == 0 {
+	for l.isEmpty() {
 		l.dataOrFinishEvent.Wait()
-		if l.noMoreData && len(l.bufferedResponses) == 0 {
+		if l.noMoreData && l.isEmpty() {
 			break
 		}
 	}
 
-	return len(l.bufferedResponses) == 0 && l.noMoreData
+	return l.isEmpty() && l.noMoreData
 }
 
 // Next either blocks until more data is available or reads
@@ -295,23 +315,24 @@ func (l *lazyRespSet) Next() bool {
 
 	l.initialized = true
 
-	if l.noMoreData && len(l.bufferedResponses) == 0 {
+	if l.noMoreData && l.isEmpty() {
 		l.lastResp = nil
 
 		return false
 	}
 
-	for len(l.bufferedResponses) == 0 {
+	for l.isEmpty() {
 		l.dataOrFinishEvent.Wait()
-		if l.noMoreData && len(l.bufferedResponses) == 0 {
+		if l.noMoreData && l.isEmpty() {
 			break
 		}
 	}
 
-	if len(l.bufferedResponses) > 0 {
-		l.lastResp = l.bufferedResponses[0]
+	if !l.isEmpty() {
+		l.lastResp = l.bufferedResponses[l.ringHead]
 		if l.initialized {
-			l.bufferedResponses = l.bufferedResponses[1:]
+			l.ringHead = (l.ringHead + 1) % l.fixedBufferSize
+			l.bufferSlotEvent.Signal()
 		}
 		return true
 	}
@@ -338,8 +359,12 @@ func newLazyRespSet(
 	shardMatcher *storepb.ShardMatcher,
 	applySharding bool,
 	emptyStreamResponses prometheus.Counter,
+	fixedBufferSize int,
 ) respSet {
-	bufferedResponses := []*storepb.SeriesResponse{}
+	// A ring buffer of size N can hold N - 1 elements at most in order to distinguish being empty from being full.
+	// That's why the size is increased by 1 internally.
+	fixedBufferSize++
+	bufferedResponses := make([]*storepb.SeriesResponse, fixedBufferSize)
 	bufferedResponsesMtx := &sync.Mutex{}
 	dataAvailable := sync.NewCond(bufferedResponsesMtx)
 
@@ -351,9 +376,14 @@ func newLazyRespSet(
 		closeSeries:          closeSeries,
 		span:                 span,
 		dataOrFinishEvent:    dataAvailable,
+		bufferSlotEvent:      sync.NewCond(bufferedResponsesMtx),
 		bufferedResponsesMtx: bufferedResponsesMtx,
 		bufferedResponses:    bufferedResponses,
 		shardMatcher:         shardMatcher,
+		fixedBufferSize:      fixedBufferSize,
+		ringHead:             0,
+		ringTail:             0,
+		closed:               false,
 	}
 	respSet.storeLabels = make(map[string]struct{})
 	for _, ls := range storeLabelSets {
@@ -406,11 +436,16 @@ func newLazyRespSet(
 				} else {
 					rerr = errors.Wrapf(err, "receive series from %s", st)
 				}
-
 				l.span.SetTag("err", rerr.Error())
 
 				l.bufferedResponsesMtx.Lock()
-				l.bufferedResponses = append(l.bufferedResponses, storepb.NewWarnSeriesResponse(rerr))
+				for l.isFull() && !l.closed {
+					l.bufferSlotEvent.Wait()
+				}
+				if !l.closed {
+					l.bufferedResponses[l.ringTail] = storepb.NewWarnSeriesResponse(rerr)
+					l.ringTail = (l.ringTail + 1) % l.fixedBufferSize
+				}
 				l.noMoreData = true
 				l.dataOrFinishEvent.Signal()
 				l.bufferedResponsesMtx.Unlock()
@@ -429,8 +464,14 @@ func newLazyRespSet(
 			}
 
 			l.bufferedResponsesMtx.Lock()
-			l.bufferedResponses = append(l.bufferedResponses, resp)
-			l.dataOrFinishEvent.Signal()
+			for l.isFull() && !l.closed {
+				l.bufferSlotEvent.Wait()
+			}
+			if !l.closed {
+				l.bufferedResponses[l.ringTail] = resp
+				l.ringTail = (l.ringTail + 1) % l.fixedBufferSize
+				l.dataOrFinishEvent.Signal()
+			}
 			l.bufferedResponsesMtx.Unlock()
 			return true
 		}
@@ -474,6 +515,7 @@ func newAsyncRespSet(
 	shardInfo *storepb.ShardInfo,
 	logger log.Logger,
 	emptyStreamResponses prometheus.Counter,
+	lazyRetrievalMaxBufferedResponses int,
 ) (respSet, error) {
 
 	var (
@@ -525,6 +567,11 @@ func newAsyncRespSet(
 	switch retrievalStrategy {
 	case LazyRetrieval:
 		span.SetTag("retrival_strategy", LazyRetrieval)
+		if lazyRetrievalMaxBufferedResponses < 1 {
+			// Some unit and e2e tests hit this path.
+			lazyRetrievalMaxBufferedResponses = 1
+		}
+
 		return newLazyRespSet(
 			span,
 			frameTimeout,
@@ -535,6 +582,7 @@ func newAsyncRespSet(
 			shardMatcher,
 			applySharding,
 			emptyStreamResponses,
+			lazyRetrievalMaxBufferedResponses,
 		), nil
 	case EagerRetrieval:
 		span.SetTag("retrival_strategy", EagerRetrieval)
@@ -560,6 +608,8 @@ func (l *lazyRespSet) Close() {
 	defer l.bufferedResponsesMtx.Unlock()
 
 	l.closeSeries()
+	l.closed = true
+	l.bufferSlotEvent.Signal()
 	l.noMoreData = true
 	l.dataOrFinishEvent.Signal()