[dbnode] Query stats tracker for metrics and limits

src/cmd/services/m3dbnode/config/limits.go

@@ -20,6 +20,8 @@
 
 package config
 
+import "time"
+
 // Limits contains configuration for configurable limits that can be applied to M3DB.
 type Limits struct {
 	// MaxOutstandingWriteRequests controls the maximum number of outstanding write requests
@@ -39,4 +41,15 @@ type Limits struct {
 	// process would pause until some of the repaired bytes had been persisted to disk (and subsequently
 	// evicted from memory) at which point it would resume.
 	MaxOutstandingRepairedBytes int64 `yaml:"maxOutstandingRepairedBytes" validate:"min=0"`
+
+	// MaxRecentlyQueriedBlocks controls the max blocks being queried within a given
+	// lookback period. Queries which are issued while this max is surpassed are abandonded.
+	MaxRecentlyQueriedBlocks *MaxRecentlyQueriedBlocks `yaml:"maxRecentlyQueriedBlocks"`
+}
+
+// MaxRecentlyQueriedBlocks controls the max blocks being queried within a given
+// lookback period. Queries which are issued while this max is surpassed are abandonded.
+type MaxRecentlyQueriedBlocks struct {
+	Max      int64         `yaml:"max" validate:"min=0"`
+	Lookback time.Duration `yaml:"lookback"`
 }

src/dbnode/server/server.go

@@ -138,7 +138,7 @@ type RunOptions struct {
 	InterruptCh <-chan error
 
 	// QueryStatsTrackerFn returns a tracker for tracking query stats.
-	QueryStatsTrackerFn func(instrument.Options) stats.QueryStatsTracker
+	QueryStatsTrackerFn func(opts instrument.Options, config stats.QueryStatsConfig) stats.QueryStatsTracker
 
 	// CustomOptions are custom options to apply to the session.
 	CustomOptions []client.CustomAdminOption
@@ -409,11 +409,24 @@ func Run(runOpts RunOptions) {
 	defer stopReporting()
 
 	// Setup query stats tracking.
+	var statsConfig stats.QueryStatsConfig
+	if runOpts.Config.Limits.MaxRecentlyQueriedBlocks == nil {
+		statsConfig = stats.DefaultQueryStatsConfigForMetrics()
+	} else {
+		maxQueriedBlocks := runOpts.Config.Limits.MaxRecentlyQueriedBlocks
+		statsConfig = stats.DefaultQueryStatsConfigForMetricsAndLimits(
+			maxQueriedBlocks.Max,
+			maxQueriedBlocks.Lookback,
+		)
+	}
 	var tracker stats.QueryStatsTracker
 	if runOpts.QueryStatsTrackerFn == nil {
-		tracker = stats.DefaultQueryStatsTrackerForMetrics(iopts)
+		tracker, err = stats.DefaultQueryStatsTracker(iopts, statsConfig)
+		if err != nil {
+			logger.Fatal("could not construct query stats tracker from config", zap.Error(err))
+		}
 	} else {
-		tracker = runOpts.QueryStatsTrackerFn(iopts)
+		tracker = runOpts.QueryStatsTrackerFn(iopts, statsConfig)
 	}
 	queryStats := stats.NewQueryStats(tracker)
 	queryStats.Start()

src/dbnode/storage/stats/query_stats.go

@@ -49,6 +49,15 @@ type QueryStats interface {
 	Stop()
 }
 
+// QueryStatsConfig holds configuration for how a tracker should handle query stats.
+type QueryStatsConfig struct {
+	// MaxDocs limits how many recently queried max
+	// documents are allowed before queries are abandoned.
+	MaxDocs *int64
+	// Lookback specifies the lookback period over which stats are aggregated.
+	Lookback time.Duration
+}
+
 // QueryStatsValues stores values of query stats.
 type QueryStatsValues struct {
 	RecentDocs int64

src/dbnode/storage/stats/query_stats_metrics.go → src/dbnode/storage/stats/query_stats_default_tracker.go

@@ -21,6 +21,7 @@
 package stats
 
 import (
+	"fmt"
 	"time"
 
 	"github.com/m3db/m3/src/x/instrument"
@@ -31,31 +32,71 @@ import (
 const defaultLookback = time.Second * 5
 
 // Tracker implementation that emits query stats as metrics.
-type queryStatsMetricsTracker struct {
+type queryStatsTracker struct {
 	recentDocs tally.Gauge
 	totalDocs  tally.Counter
+
+	config QueryStatsConfig
+}
+
+var _ QueryStatsTracker = (*queryStatsTracker)(nil)
+
+// DefaultQueryStatsConfigForMetrics returns a default
+// config for tracking query stats as metrics.
+func DefaultQueryStatsConfigForMetrics() QueryStatsConfig {
+	return QueryStatsConfig{
+		MaxDocs:  nil,
+		Lookback: defaultLookback,
+	}
 }
 
-var _ QueryStatsTracker = (*queryStatsMetricsTracker)(nil)
+// DefaultQueryStatsConfigForMetricsAndLimits returns a
+// default config for tracking query stats as metrics.
+func DefaultQueryStatsConfigForMetricsAndLimits(
+	maxDocs int64,
+	lookback time.Duration,
+) QueryStatsConfig {
+	return QueryStatsConfig{
+		MaxDocs:  &maxDocs,
+		Lookback: lookback,
+	}
+}
 
-// DefaultQueryStatsTrackerForMetrics provides a tracker
-// implementation that emits query stats as metrics.
-func DefaultQueryStatsTrackerForMetrics(opts instrument.Options) QueryStatsTracker {
+// DefaultQueryStatsTracker provides a tracker
+// implementation that emits query stats as metrics
+// and enforces limits.
+func DefaultQueryStatsTracker(
+	opts instrument.Options,
+	config QueryStatsConfig,
+) (QueryStatsTracker, error) {
+	if config.MaxDocs != nil && *config.MaxDocs <= 0 {
+		return nil, fmt.Errorf("query stats tracker requires max docs > 0 if not nil (%d)", *config.MaxDocs)
+	}
+	if config.Lookback <= 0 {
+		return nil, fmt.Errorf("query stats tracker requires lookback > 0 (%d)", config.Lookback)
+	}
 	scope := opts.
 		MetricsScope().
 		SubScope("query-stats")
-	return &queryStatsMetricsTracker{
+	return &queryStatsTracker{
+		config:     config,
 		recentDocs: scope.Gauge("recent-docs-per-block"),
 		totalDocs:  scope.Counter("total-docs-per-block"),
-	}
+	}, nil
 }
 
-func (t *queryStatsMetricsTracker) TrackStats(values QueryStatsValues) error {
+func (t *queryStatsTracker) TrackStats(values QueryStatsValues) error {
+	// Track stats as metrics.
 	t.recentDocs.Update(float64(values.RecentDocs))
 	t.totalDocs.Inc(values.NewDocs)
+
+	// Enforce max queried docs (if specified).
+	if t.config.MaxDocs != nil && values.RecentDocs > *t.config.MaxDocs {
+		return fmt.Errorf("query was aborted due to too many recent docs queried (%d)", values.RecentDocs)
+	}
 	return nil
 }
 
-func (t *queryStatsMetricsTracker) Lookback() time.Duration {
-	return defaultLookback
+func (t *queryStatsTracker) Lookback() time.Duration {
+	return t.config.Lookback
 }

src/dbnode/storage/stats/query_stats_default_tracker_test.go

@@ -0,0 +1,155 @@
+// Copyright (c) 2020 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
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package stats
+
+import (
+	"testing"
+	"time"
+
+	"github.com/m3db/m3/src/x/instrument"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/uber-go/tally"
+)
+
+func TestValidateTrackerInputs(t *testing.T) {
+	scope := tally.NewTestScope("", nil)
+	opts := instrument.NewOptions().SetMetricsScope(scope)
+
+	validMax := int64(1)
+	negativeMax := int64(-1)
+	zeroMax := int64(0)
+
+	for _, test := range []struct {
+		name          string
+		maxDocs       *int64
+		lookback      time.Duration
+		expectedError string
+	}{
+		{"valid lookback without limit", nil, time.Millisecond, ""},
+		{"valid lookback with valid limit", &validMax, time.Millisecond, ""},
+		{"negative lookback", nil, -time.Millisecond, "query stats tracker requires lookback > 0 (-1000000)"},
+		{"zero lookback", nil, time.Duration(0), "query stats tracker requires lookback > 0 (0)"},
+		{"negative max", &negativeMax, time.Millisecond, "query stats tracker requires max docs > 0 if not nil (-1)"},
+		{"zero max", &zeroMax, time.Millisecond, "query stats tracker requires max docs > 0 if not nil (0)"},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			_, err := DefaultQueryStatsTracker(opts, QueryStatsConfig{
+				MaxDocs:  test.maxDocs,
+				Lookback: test.lookback,
+			})
+			if test.expectedError != "" {
+				require.Error(t, err)
+				require.Equal(t, test.expectedError, err.Error())
+			} else {
+				require.NoError(t, err)
+			}
+		})
+	}
+}
+
+func TestEmitQueryStatsBasedMetrics(t *testing.T) {
+	for _, test := range []struct {
+		name   string
+		config QueryStatsConfig
+	}{
+		{"metrics only", DefaultQueryStatsConfigForMetrics()},
+		{"metrics and limits", DefaultQueryStatsConfigForMetricsAndLimits(1000, time.Second)},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			scope := tally.NewTestScope("", nil)
+			opts := instrument.NewOptions().SetMetricsScope(scope)
+
+			tracker, err := DefaultQueryStatsTracker(opts, test.config)
+			require.NoError(t, err)
+
+			err = tracker.TrackStats(QueryStatsValues{RecentDocs: 100, NewDocs: 5})
+			require.NoError(t, err)
+			verifyMetrics(t, scope, 100, 5)
+
+			err = tracker.TrackStats(QueryStatsValues{RecentDocs: 140, NewDocs: 10})
+			require.NoError(t, err)
+			verifyMetrics(t, scope, 140, 15)
+		})
+	}
+}
+
+func TestLimitMaxDocs(t *testing.T) {
+	scope := tally.NewTestScope("", nil)
+	opts := instrument.NewOptions().SetMetricsScope(scope)
+
+	maxDocs := int64(100)
+
+	for _, test := range []struct {
+		name             string
+		config           QueryStatsConfig
+		expectLimitError bool
+	}{
+		{"metrics only", DefaultQueryStatsConfigForMetrics(), false},
+		{"metrics and limits", DefaultQueryStatsConfigForMetricsAndLimits(maxDocs, time.Second), true},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			tracker, err := DefaultQueryStatsTracker(opts, test.config)
+			require.NoError(t, err)
+
+			err = tracker.TrackStats(QueryStatsValues{RecentDocs: maxDocs + 1})
+			if test.expectLimitError {
+				require.Error(t, err)
+				require.Equal(t, "query was aborted due to too many recent docs queried (101)", err.Error())
+			} else {
+				require.NoError(t, err)
+			}
+
+			err = tracker.TrackStats(QueryStatsValues{RecentDocs: maxDocs - 1})
+			require.NoError(t, err)
+
+			err = tracker.TrackStats(QueryStatsValues{RecentDocs: 0})
+			require.NoError(t, err)
+
+			err = tracker.TrackStats(QueryStatsValues{RecentDocs: maxDocs + 1})
+			if test.expectLimitError {
+				require.Error(t, err)
+				require.Equal(t, "query was aborted due to too many recent docs queried (101)", err.Error())
+			} else {
+				require.NoError(t, err)
+			}
+
+			err = tracker.TrackStats(QueryStatsValues{RecentDocs: maxDocs - 1})
+			require.NoError(t, err)
+
+			err = tracker.TrackStats(QueryStatsValues{RecentDocs: 0})
+			require.NoError(t, err)
+		})
+	}
+}
+
+func verifyMetrics(t *testing.T, scope tally.TestScope, expectedRecent float64, expectedTotal int64) {
+	snapshot := scope.Snapshot()
+
+	recent, exists := snapshot.Gauges()["query-stats.recent-docs-per-block+"]
+	assert.True(t, exists)
+	assert.Equal(t, expectedRecent, recent.Value())
+
+	total, exists := snapshot.Counters()["query-stats.total-docs-per-block+"]
+	assert.True(t, exists)
+	assert.Equal(t, expectedTotal, total.Value())
+}