feat: add QPR mem size limit

api/storeapi/store_api.proto

@@ -144,6 +144,7 @@ enum SearchErrorCode {
   TOO_MANY_GROUP_TOKENS = 5;
   TOO_MANY_FRACTION_TOKENS = 6;
   TOO_MANY_FIELD_VALUES = 7;
+  MEMORY_LIMIT_EXCEEDED = 8;
 }
 
 message StartAsyncSearchRequest {

cmd/seq-db/seq-db.go

@@ -297,6 +297,7 @@ func startStore(
 				WorkersCount:          cfg.Resources.SearchWorkers,
 				MaxFractionHits:       cfg.Limits.FractionHits,
 				FractionsPerIteration: config.NumCPU,
+				MaxQprMemory:          uint64(cfg.Limits.QprMemoryUsage),
 				RequestsLimit:         uint64(cfg.Limits.SearchRequests),
 				LogThreshold:          cfg.SlowLogs.SearchThreshold,
 				Async: asyncsearcher.AsyncSearcherConfig{

config/config.go

@@ -131,6 +131,9 @@ type Config struct {
 		// DocSize specifies maximum possible size for single document.
 		// Document larger than this threshold will be skipped.
 		DocSize Bytes `config:"doc_size" default:"128KiB"`
+		// QprMemoryUsage specifies maximum heap memory which a single QPR (query partial result)
+		// can use in either store or proxy.
+		QprMemoryUsage Bytes `config:"qpr_memory_usage" default:"0B"`
 
 		Aggregation struct {
 			// FieldTokens specifies maximum amount of unique field tokens

consts/consts.go

@@ -88,4 +88,5 @@ var (
 	ErrTooManyGroupTokens        = errors.New("aggregation has too many group tokens")
 	ErrTooManyFractionTokens     = errors.New("aggregation has too many fraction tokens")
 	ErrTooManyFractionsHit       = errors.New("too many fractions hit")
+	ErrMemoryLimitExceeded       = errors.New("memory limit exceeded")
 )

fracmanager/searcher.go

@@ -20,6 +20,7 @@ import (
 type SearcherCfg struct {
 	MaxFractionHits       int // the maximum number of fractions used in the search
 	FractionsPerIteration int
+	MaxQprMemory          uint64 // max heap memory a single QPR can use. 0 if no limit set
 }
 
 type Searcher struct {
@@ -59,6 +60,7 @@ func (s *Searcher) SearchDocs(ctx context.Context, fracs []frac.Fraction, params
 
 	var totalSearchTimeNanos int64
 	var totalWaitTimeNanos int64
+	var totalMemUsage uint64
 
 	for len(remainingFracs) > 0 && (scanAll || params.Limit > 0) {
 		subQPRs, searchTimeNanos, waitTimeNanos, err := s.searchDocsAsync(ctx, remainingFracs.Shift(fracsChunkSize), params)
@@ -71,6 +73,12 @@ func (s *Searcher) SearchDocs(ctx context.Context, fracs []frac.Fraction, params
 
 		seq.MergeQPRs(total, subQPRs, origLimit, seq.MillisToMID(params.HistInterval), params.Order)
 
+		totalMemUsage = total.MemUsage()
+
+		if s.cfg.MaxQprMemory > 0 && totalMemUsage > s.cfg.MaxQprMemory {
+			return nil, fmt.Errorf("%w: used %d bytes, limit %d", consts.ErrMemoryLimitExceeded, total.MemUsage(), s.cfg.MaxQprMemory)
+		}
+
 		// reduce the limit on the number of ensured docs in response
 		params.Limit = origLimit - calcEnsuredIDsCount(total.IDs, remainingFracs, params.Order)
 

metric/store.go

@@ -60,6 +60,13 @@ var (
 		Help:      "Search request duration time (only successful searches)",
 		Buckets:   SecondsBuckets,
 	})
+	SearchQprMemSize = promauto.NewHistogram(prometheus.HistogramOpts{
+		Namespace: "seq_db_store",
+		Subsystem: "search",
+		Name:      "qpr_mem_size_bytes",
+		Help:      "QPR heap memory size in bytes",
+		Buckets:   prometheus.ExponentialBuckets(1024, 3, 21),
+	})
 
 	SearchRangesSeconds = promauto.NewHistogram(prometheus.HistogramOpts{
 		Namespace: "seq_db_store",

proxy/search/ingestor.go

@@ -601,6 +601,9 @@ func (si *Ingestor) searchShard(
 			if errMessage == consts.ErrTooManyFractionTokens.Error() {
 				return nil, source, fmt.Errorf("store forbids aggregation request: %w", consts.ErrTooManyFractionTokens)
 			}
+			if errMessage == consts.ErrMemoryLimitExceeded.Error() {
+				return nil, source, fmt.Errorf("store forbids search request: %w", consts.ErrMemoryLimitExceeded)
+			}
 			errs = append(errs, err)
 			continue
 		}
@@ -620,6 +623,8 @@ func (si *Ingestor) searchShard(
 			return nil, source, fmt.Errorf("store forbids aggregation request: %w", consts.ErrTooManyGroupTokens)
 		case storeapi.SearchErrorCode_TOO_MANY_FRACTION_TOKENS:
 			return nil, source, fmt.Errorf("store forbids aggregation request: %w", consts.ErrTooManyFractionTokens)
+		case storeapi.SearchErrorCode_MEMORY_LIMIT_EXCEEDED:
+			return nil, source, fmt.Errorf("store forbids search request: %w", consts.ErrMemoryLimitExceeded)
 		case storeapi.SearchErrorCode_TOO_MANY_FRACTIONS_HIT:
 			return nil, source, fmt.Errorf("store forbids request: %w", consts.ErrTooManyFractionsHit)
 		}