Addressed comments

Author: hczhu-db

cmd/thanos/query.go

@@ -247,7 +247,7 @@ func registerQuery(app *extkingpin.App) {
 
 	rewriteAggregationLabelTo := cmd.Flag("query.aggregation-label-value-override", "The value override for __rollup__ label for aggregated metrics. If set to x, all queries on aggregated metrics will have a __rollup__=x matcher. Leave empty to disable this behavior. Default is empty.").Default("").String()
 
-	lazyRetrievalBufferSize := cmd.Flag("query.lazy-retrieval-buffer-size", "The fixed buffer size (in terms of how many store responses) for lazy retrieval strategy. This is to limit the memory usage.").
+	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
@@ -390,7 +390,7 @@ func registerQuery(app *extkingpin.App) {
 			*tenantLabel,
 			*enableGroupReplicaPartialStrategy,
 			*rewriteAggregationLabelTo,
-			*lazyRetrievalBufferSize,
+			*lazyRetrievalMaxBufferedResponses,
 		)
 	})
 }
@@ -476,7 +476,7 @@ func runQuery(
 	tenantLabel string,
 	groupReplicaPartialResponseStrategy bool,
 	rewriteAggregationLabelTo string,
-	lazyRetrievalBufferSize int,
+	lazyRetrievalMaxBufferedResponses int,
 ) error {
 	comp := component.Query
 	if alertQueryURL == "" {
@@ -564,7 +564,7 @@ func runQuery(
 		store.WithTSDBSelector(tsdbSelector),
 		store.WithProxyStoreDebugLogging(debugLogging),
 		store.WithQuorumChunkDedup(queryDeduplicationFunc == dedup.AlgorithmQuorum),
-		store.WithLazyRetrievalBufferSize(lazyRetrievalBufferSize),
+		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 {
@@ -612,6 +622,7 @@ func NewBucketStore(
 		enableChunkHashCalculation:      enableChunkHashCalculation,
 		seriesBatchSize:                 SeriesBatchSize,
 		sortingStrategy:                 sortingStrategyStore,
+		lazyRetrievalMaxBufferedResponses: 10,
 		indexHeaderLazyDownloadStrategy: indexheader.AlwaysEagerDownloadIndexHeader,
 		requestLoggerFunc:               NoopRequestLoggerFunc,
 	}
@@ -1623,7 +1634,7 @@ func (s *BucketStore) Series(req *storepb.SeriesRequest, seriesSrv storepb.Store
 						shardMatcher,
 						false,
 						s.metrics.emptyPostingCount.WithLabelValues(tenant),
-						100,
+						s.lazyRetrievalMaxBufferedResponses,
 					)
 				}
 

pkg/store/proxy.go

@@ -101,7 +101,7 @@ type ProxyStore struct {
 	quorumChunkDedup        bool
 	enableDedup             bool
 	matcherConverter        *storepb.MatcherConverter
-	lazyRetrievalBufferSize int
+	lazyRetrievalMaxBufferedResponses int
 }
 
 type proxyStoreMetrics struct {
@@ -138,9 +138,9 @@ func RegisterStoreServer(storeSrv storepb.StoreServer, logger log.Logger) func(*
 // BucketStoreOption are functions that configure BucketStore.
 type ProxyStoreOption func(s *ProxyStore)
 
-func WithLazyRetrievalBufferSize(buferSize int) ProxyStoreOption {
+func WithLazyRetrievalMaxBufferedResponsesForProxy(buferSize int) ProxyStoreOption {
 	return func(s *ProxyStore) {
-		s.lazyRetrievalBufferSize = buferSize
+		s.lazyRetrievalMaxBufferedResponses = buferSize
 	}
 }
 
@@ -392,16 +392,16 @@ func (s *ProxyStore) Series(originalRequest *storepb.SeriesRequest, srv storepb.
 		}
 	}
 	defer logGroupReplicaErrors()
-	lazyRetrievalBufferSize := s.lazyRetrievalBufferSize
-	if lazyRetrievalBufferSize <= 0 {
+	lazyRetrievalMaxBufferedResponses := s.lazyRetrievalMaxBufferedResponses
+	if lazyRetrievalMaxBufferedResponses <= 0 {
 		// Use 1 as default value for lazy retrieval buffer size.
 		// Unit tests hit this path so that corner cases can be tested with buffer size 1.
-		lazyRetrievalBufferSize = 1
+		lazyRetrievalMaxBufferedResponses = 1
 	}
 	for _, st := range stores {
 		st := st
 
-		respSet, err := newAsyncRespSet(ctx, st, r, s.responseTimeout, s.retrievalStrategy, &s.buffers, r.ShardInfo, reqLogger, s.metrics.emptyStreamResponses, lazyRetrievalBufferSize)
+		respSet, err := newAsyncRespSet(ctx, st, r, s.responseTimeout, s.retrievalStrategy, &s.buffers, r.ShardInfo, reqLogger, s.metrics.emptyStreamResponses, 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()