Limit lazyRespSet memory buffer size

Author: hczhu-db

cmd/thanos/query.go

@@ -247,6 +247,9 @@ 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.").
+		Default("20").Int()
+
 	var storeRateLimits store.SeriesSelectLimits
 	storeRateLimits.RegisterFlags(cmd)
 
@@ -387,6 +390,7 @@ func registerQuery(app *extkingpin.App) {
 			*tenantLabel,
 			*enableGroupReplicaPartialStrategy,
 			*rewriteAggregationLabelTo,
+			*lazyRetrievalBufferSize,
 		)
 	})
 }
@@ -472,6 +476,7 @@ func runQuery(
 	tenantLabel string,
 	groupReplicaPartialResponseStrategy bool,
 	rewriteAggregationLabelTo string,
+	lazyRetrievalBufferSize int,
 ) error {
 	comp := component.Query
 	if alertQueryURL == "" {
@@ -559,6 +564,7 @@ func runQuery(
 		store.WithTSDBSelector(tsdbSelector),
 		store.WithProxyStoreDebugLogging(debugLogging),
 		store.WithQuorumChunkDedup(queryDeduplicationFunc == dedup.AlgorithmQuorum),
+		store.WithLazyRetrievalBufferSize(lazyRetrievalBufferSize),
 	}
 
 	// Parse and sanitize the provided replica labels flags.

pkg/store/bucket.go

@@ -1623,6 +1623,7 @@ func (s *BucketStore) Series(req *storepb.SeriesRequest, seriesSrv storepb.Store
 						shardMatcher,
 						false,
 						s.metrics.emptyPostingCount.WithLabelValues(tenant),
+						100,
 					)
 				}
 

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
+	lazyRetrievalBufferSize 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 WithLazyRetrievalBufferSize(buferSize int) ProxyStoreOption {
+	return func(s *ProxyStore) {
+		s.lazyRetrievalBufferSize = buferSize
+	}
+}
+
 // WithProxyStoreDebugLogging toggles debug logging.
 func WithProxyStoreDebugLogging(enable bool) ProxyStoreOption {
 	return func(s *ProxyStore) {
@@ -385,11 +392,16 @@ func (s *ProxyStore) Series(originalRequest *storepb.SeriesRequest, srv storepb.
 		}
 	}
 	defer logGroupReplicaErrors()
-
+	lazyRetrievalBufferSize := s.lazyRetrievalBufferSize
+	if lazyRetrievalBufferSize <= 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
+	}
 	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, lazyRetrievalBufferSize)
 		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,
+	fixedBufferSizeForLazyRetrieval int,
 ) (respSet, error) {
 
 	var (
@@ -535,6 +577,7 @@ func newAsyncRespSet(
 			shardMatcher,
 			applySharding,
 			emptyStreamResponses,
+			fixedBufferSizeForLazyRetrieval,
 		), nil
 	case EagerRetrieval:
 		span.SetTag("retrival_strategy", EagerRetrieval)
@@ -560,6 +603,8 @@ func (l *lazyRespSet) Close() {
 	defer l.bufferedResponsesMtx.Unlock()
 
 	l.closeSeries()
+	l.closed = true
+	l.bufferSlotEvent.Signal()
 	l.noMoreData = true
 	l.dataOrFinishEvent.Signal()