Initial query limit overriding

src/dbnode/generated/thrift/rpc/rpc.go

@@ -1,4 +1,4 @@
-// Copyright (c) 2020 Uber Technologies, Inc.
+// Copyright (c) 2021 Uber Technologies, Inc.
 //
 // Permission is hereby granted, free of charge, to any person obtaining a copy
 // of this software and associated documentation files (the "Software"), to deal

src/dbnode/integration/query_limit_test.go

@@ -87,8 +87,9 @@ func newTestOptionsWithIndexedNamespace(t *testing.T) (TestOptions, namespace.Me
 
 func newTestSetupWithQueryLimits(t *testing.T, opts TestOptions) TestSetup {
 	storageLimitsFn := func(storageOpts storage.Options) storage.Options {
+		limit := int64(1)
 		queryLookback := limits.DefaultLookbackLimitOptions()
-		queryLookback.Limit = 1
+		queryLookback.Limit = &limit
 		queryLookback.Lookback = time.Hour
 
 		limitOpts := limits.NewOptions().

src/dbnode/kvconfig/keys.go

@@ -49,4 +49,16 @@ const (
 	// ClientWriteConsistencyLevel is the KV config key for the runtime
 	// configuration specifying the client write consistency level
 	ClientWriteConsistencyLevel = "m3db.client.write-consistency-level"
+
+	// DocsLimit is the KV config key for the docs matched query limit.
+	// Settings in string form "{limit},{lookback}", e.g. "1000,15s".
+	DocsLimit = "m3db.limits.docs"
+
+	// DiskBytesReadLimit is the KV config key for the disk bytes read query limit.
+	// Settings in string form "{limit},{lookback}", e.g. "1000,15s".
+	DiskBytesReadLimit = "m3db.limits.disk-bytes-read"
+
+	// DiskSeriesReadLimit is the KV config key for the disk series read query limit.
+	// Settings in string form "{limit},{lookback}", e.g. "1000,15s".
+	DiskSeriesReadLimit = "m3db.limits.disk-series-read"
 )

src/dbnode/persist/fs/retriever_test.go

@@ -805,12 +805,13 @@ func TestBlockRetrieverHandlesSeekByIndexEntryErrors(t *testing.T) {
 
 func TestLimitSeriesReadFromDisk(t *testing.T) {
 	scope := tally.NewTestScope("test", nil)
+	limit := int64(2)
 	limitOpts := limits.NewOptions().
 		SetInstrumentOptions(instrument.NewOptions().SetMetricsScope(scope)).
 		SetBytesReadLimitOpts(limits.DefaultLookbackLimitOptions()).
 		SetDocsLimitOpts(limits.DefaultLookbackLimitOptions()).
 		SetDiskSeriesReadLimitOpts(limits.LookbackLimitOptions{
-			Limit:    1,
+			Limit:    &limit,
 			Lookback: time.Second * 1,
 		})
 	queryLimits, err := limits.NewQueryLimits(limitOpts)

src/dbnode/server/server.go

@@ -32,6 +32,7 @@ import (
 	"path"
 	"runtime"
 	"runtime/debug"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -452,15 +453,21 @@ func Run(runOpts RunOptions) {
 	bytesReadLimit := limits.DefaultLookbackLimitOptions()
 	diskSeriesReadLimit := limits.DefaultLookbackLimitOptions()
 	if limitConfig := runOpts.Config.Limits.MaxRecentlyQueriedSeriesBlocks; limitConfig != nil {
-		docsLimit.Limit = limitConfig.Value
+		if limitConfig.Value != 0 {
+			docsLimit.Limit = &limitConfig.Value
+		}
 		docsLimit.Lookback = limitConfig.Lookback
 	}
 	if limitConfig := runOpts.Config.Limits.MaxRecentlyQueriedSeriesDiskBytesRead; limitConfig != nil {
-		bytesReadLimit.Limit = limitConfig.Value
+		if limitConfig.Value != 0 {
+			bytesReadLimit.Limit = &limitConfig.Value
+		}
 		bytesReadLimit.Lookback = limitConfig.Lookback
 	}
 	if limitConfig := runOpts.Config.Limits.MaxRecentlyQueriedSeriesDiskRead; limitConfig != nil {
-		diskSeriesReadLimit.Limit = limitConfig.Value
+		if limitConfig.Value != 0 {
+			diskSeriesReadLimit.Limit = &limitConfig.Value
+		}
 		diskSeriesReadLimit.Lookback = limitConfig.Lookback
 	}
 	limitOpts := limits.NewOptions().
@@ -993,6 +1000,9 @@ func Run(runOpts RunOptions) {
 			runtimeOptsMgr, cfg.Limits.WriteNewSeriesPerSecond)
 		kvWatchEncodersPerBlockLimit(syncCfg.KVStore, logger,
 			runtimeOptsMgr, cfg.Limits.MaxEncodersPerBlock)
+		kvWatchQueryLimit(syncCfg.KVStore, logger, queryLimits.DocsLimit(), kvconfig.DocsLimit)
+		kvWatchQueryLimit(syncCfg.KVStore, logger, queryLimits.DiskSeriesReadLimit(), kvconfig.DiskBytesReadLimit)
+		kvWatchQueryLimit(syncCfg.KVStore, logger, queryLimits.BytesReadLimit(), kvconfig.DiskSeriesReadLimit)
 	}()
 
 	// Wait for process interrupt.
@@ -1165,6 +1175,79 @@ func kvWatchEncodersPerBlockLimit(
 	}()
 }
 
+func kvWatchQueryLimit(
+	store kv.Store,
+	logger *zap.Logger,
+	limit limits.LookbackLimit,
+	kvName string,
+) {
+	options := limit.Options()
+
+	value, err := store.Get(kvName)
+	if err == nil {
+		protoValue := &commonpb.StringProto{}
+		err = value.Unmarshal(protoValue)
+		if err == nil {
+			options = parseLookbackLimitOptions(logger, kvName, protoValue.Value, options)
+		}
+	} else if errors.Is(err, kv.ErrNotFound) {
+		logger.Warn("error resolving query limit", zap.Error(err), zap.String("name", kvName))
+	}
+
+	if err := limit.Update(options); err != nil {
+		logger.Warn("unable to set query limit", zap.Error(err), zap.String("name", kvName))
+	}
+
+	watch, err := store.Watch(kvName)
+	if err != nil {
+		logger.Error("could not watch query limit", zap.Error(err), zap.String("name", kvName))
+		return
+	}
+
+	go func() {
+		protoValue := &commonpb.StringProto{}
+		for range watch.C() {
+			value := options
+			if newValue := watch.Get(); newValue != nil {
+				if err := newValue.Unmarshal(protoValue); err != nil {
+					logger.Warn("unable to parse new query limit", zap.Error(err), zap.String("name", kvName))
+					continue
+				}
+				value = parseLookbackLimitOptions(logger, kvName, protoValue.Value, value)
+			}
+
+			if err := limit.Update(value); err != nil {
+				logger.Warn("unable to set query limit", zap.Error(err), zap.String("name", kvName))
+			}
+		}
+	}()
+}
+
+func parseLookbackLimitOptions(logger *zap.Logger,
+	kvName string,
+	val string,
+	defaultOpts limits.LookbackLimitOptions,
+) limits.LookbackLimitOptions {
+	parts := strings.Split(val, ",")
+	if val == "" {
+		defaultOpts.Limit = nil
+	} else if len(parts) == 2 {
+		parsedLimit, err := strconv.ParseInt(parts[0], 10, 64)
+		if err != nil {
+			logger.Warn("error parsing query limit value", zap.Error(err), zap.String("name", kvName))
+		} else {
+			defaultOpts.Limit = &parsedLimit
+		}
+		parsedLookback, err := time.ParseDuration(parts[1])
+		if err != nil {
+			logger.Warn("error parsing query limit lookback", zap.Error(err), zap.String("name", kvName))
+		} else {
+			defaultOpts.Lookback = parsedLookback
+		}
+	}
+	return defaultOpts
+}
+
 func kvWatchClientConsistencyLevels(
 	store kv.Store,
 	logger *zap.Logger,

src/dbnode/storage/limits/noop_query_limits.go

@@ -58,6 +58,14 @@ func (q *noOpQueryLimits) Stop() {
 func (q *noOpQueryLimits) Start() {
 }
 
+func (q *noOpLookbackLimit) Options() LookbackLimitOptions {
+	return LookbackLimitOptions{}
+}
+
+func (q *noOpLookbackLimit) Update(LookbackLimitOptions) error {
+	return nil
+}
+
 func (q *noOpLookbackLimit) Inc(int, []byte) error {
 	return nil
 }

src/dbnode/storage/limits/query_limits.go

@@ -22,13 +22,15 @@ package limits
 
 import (
 	"fmt"
+	"sync"
 	"time"
 
 	xerrors "github.com/m3db/m3/src/x/errors"
 	"github.com/m3db/m3/src/x/instrument"
 
 	"github.com/uber-go/tally"
 	"go.uber.org/atomic"
+	"go.uber.org/zap"
 )
 
 const (
@@ -45,15 +47,19 @@ type lookbackLimit struct {
 	name    string
 	options LookbackLimitOptions
 	metrics lookbackLimitMetrics
+	logger  *zap.Logger
 	recent  *atomic.Int64
 	stopCh  chan struct{}
+	lock    sync.RWMutex
 }
 
 type lookbackLimitMetrics struct {
-	recentCount tally.Gauge
-	recentMax   tally.Gauge
-	total       tally.Counter
-	exceeded    tally.Counter
+	optionsMax      tally.Gauge
+	optionsLookback tally.Gauge
+	recentCount     tally.Gauge
+	recentMax       tally.Gauge
+	total           tally.Counter
+	exceeded        tally.Counter
 
 	sourceLogger SourceLogger
 }
@@ -67,7 +73,7 @@ var (
 func DefaultLookbackLimitOptions() LookbackLimitOptions {
 	return LookbackLimitOptions{
 		// Default to no limit.
-		Limit:    0,
+		Limit:    nil,
 		Lookback: defaultLookback,
 	}
 }
@@ -110,6 +116,7 @@ func newLookbackLimit(
 		name:    name,
 		options: opts,
 		metrics: newLookbackLimitMetrics(instrumentOpts, name, sourceLoggerBuilder),
+		logger:  instrumentOpts.Logger(),
 		recent:  atomic.NewInt64(0),
 		stopCh:  make(chan struct{}),
 	}
@@ -171,6 +178,36 @@ func (q *queryLimits) AnyExceeded() error {
 	return q.bytesReadLimit.exceeded()
 }
 
+func (q *lookbackLimit) Options() LookbackLimitOptions {
+	return q.options
+}
+
+// Override overrides the limit set on construction.
+func (q *lookbackLimit) Update(opts LookbackLimitOptions) error {
+	if err := opts.validate(); err != nil {
+		return err
+	}
+
+	q.lock.Lock()
+	defer q.lock.Unlock()
+
+	old := q.options
+	q.options = opts
+
+	// If the lookback changed, replace the background goroutine that manages the periodic resetting.
+	if q.options.Lookback != old.Lookback {
+		q.stop()
+		q.start()
+	}
+
+	q.logger.Info("query limit options updated",
+		zap.String("name", q.name),
+		zap.Any("new", opts),
+		zap.Any("old", old))
+
+	return nil
+}
+
 // Inc increments the current value and returns an error if above the limit.
 func (q *lookbackLimit) Inc(val int, source []byte) error {
 	if val < 0 {
@@ -199,7 +236,15 @@ func (q *lookbackLimit) exceeded() error {
 }
 
 func (q *lookbackLimit) checkLimit(recent int64) error {
-	if q.options.Limit > 0 && recent > q.options.Limit {
+	q.lock.RLock()
+	limit := q.options.Limit
+	q.lock.RUnlock()
+
+	if limit == nil {
+		return nil
+	}
+
+	if recent >= *limit {
 		q.metrics.exceeded.Inc(1)
 		return xerrors.NewInvalidParamsError(NewQueryLimitExceededError(fmt.Sprintf(
 			"query aborted due to limit: name=%s, limit=%d, current=%d, within=%s",
@@ -210,7 +255,9 @@ func (q *lookbackLimit) checkLimit(recent int64) error {
 
 func (q *lookbackLimit) start() {
 	ticker := time.NewTicker(q.options.Lookback)
+	ticker.Reset(q.options.Lookback)
 	go func() {
+		q.logger.Info("query limit interval started", zap.String("name", q.name))
 		for {
 			select {
 			case <-ticker.C:
@@ -221,10 +268,16 @@ func (q *lookbackLimit) start() {
 			}
 		}
 	}()
+
+	q.metrics.optionsMax.Update(float64(*q.options.Limit))
+	q.metrics.optionsLookback.Update(q.options.Lookback.Seconds())
 }
 
 func (q *lookbackLimit) stop() {
 	close(q.stopCh)
+	q.stopCh = make(chan struct{})
+
+	q.logger.Info("query limit interval stopped", zap.String("name", q.name))
 }
 
 func (q *lookbackLimit) current() int64 {
@@ -244,8 +297,8 @@ func (q *lookbackLimit) reset() {
 }
 
 func (opts LookbackLimitOptions) validate() error {
-	if opts.Limit < 0 {
-		return fmt.Errorf("query limit requires limit >= 0 (%d)", opts.Limit)
+	if opts.Limit != nil && *opts.Limit < 0 {
+		return fmt.Errorf("query limit requires limit >= 0 or nil (%d)", *opts.Limit)
 	}
 	if opts.Lookback <= 0 {
 		return fmt.Errorf("query limit requires lookback > 0 (%d)", opts.Lookback)