Add max-chunks-bytes-per-query limiter (#4216)

* Add max-chunks-bytes-per-query limiter

Signed-off-by: Tyler Reid <tyler.reid@grafana.com>

* Fix for distributor test

Signed-off-by: Tyler Reid <tyler.reid@grafana.com>

* Fix spacing and add pr number to changelog

Signed-off-by: Tyler Reid <tyler.reid@grafana.com>

* Add unit test for ingester and address code review comments

Signed-off-by: Tyler Reid <tyler.reid@grafana.com>

* Fix chunk bytes unit test and code review comments

Signed-off-by: Tyler Reid <tyler.reid@grafana.com>

* Fix linter error

Signed-off-by: Tyler Reid <tyler.reid@grafana.com>

* Move context to a global to per test for distributor tests. Address other code review comments.

Signed-off-by: Tyler Reid <tyler.reid@grafana.com>

* Update config docs

Signed-off-by: Tyler Reid <tyler.reid@grafana.com>

* Change username to userId

Signed-off-by: Tyler Reid <tyler.reid@grafana.com>

* Fix linter error

Signed-off-by: Tyler Reid <tyler.reid@grafana.com>

Co-authored-by: Tyler Reid <tyler.reid@grafana.com>

CHANGELOG.md

@@ -8,6 +8,7 @@
   - `-alertmanager.receivers-firewall.block.private-addresses` renamed to `-alertmanager.receivers-firewall-block-private-addresses`
 * [CHANGE] Change default value of `-server.grpc.keepalive.min-time-between-pings` to `10s` and `-server.grpc.keepalive.ping-without-stream-allowed` to `true`. #4168
 * [FEATURE] Querier: Added new `-querier.max-fetched-series-per-query` flag. When Cortex is running with blocks storage, the max series per query limit is enforced in the querier and applies to unique series received from ingesters and store-gateway (long-term storage). #4179
+* [FEATURE] Querier/Ruler: Added new `-querier.max-fetched-chunk-bytes-per-query` flag. When Cortex is running with blocks storage, the max chunk bytes limit is enforced in the querier and ruler and limits the size of all aggregated chunks returned from ingesters and storage as bytes for a query. #4216
 * [FEATURE] Alertmanager: Added rate-limits to notifiers. Rate limits used by all integrations can be configured using `-alertmanager.notification-rate-limit`, while per-integration rate limits can be specified via `-alertmanager.notification-rate-limit-per-integration` parameter. Both shared and per-integration limits can be overwritten using overrides mechanism. These limits are applied on individual (per-tenant) alertmanagers. Rate-limited notifications are failed notifications. It is possible to monitor rate-limited notifications via new `cortex_alertmanager_notification_rate_limited_total` metric. #4135 #4163
 * [FEATURE] Alertmanager: Added `-alertmanager.max-config-size-bytes` limit to control size of configuration files that Cortex users can upload to Alertmanager via API. This limit is configurable per-tenant. #4201
 * [FEATURE] Alertmanager: Added `-alertmanager.max-templates-count` and `-alertmanager.max-template-size-bytes` options to control number and size of templates uploaded to Alertmanager via API. These limits are configurable per-tenant. #4223

docs/configuration/config-file-reference.md

@@ -4041,6 +4041,12 @@ The `limits_config` configures default and per-tenant limits imposed by Cortex s
 # CLI flag: -querier.max-fetched-series-per-query
 [max_fetched_series_per_query: <int> | default = 0]
 
+# The maximum size of all chunks in bytes that a query can fetch from each
+# ingester and storage. This limit is enforced in the querier and ruler only
+# when running Cortex with blocks storage. 0 to disable.
+# CLI flag: -querier.max-fetched-chunk-bytes-per-query
+[max_fetched_chunk_bytes_per_query: <int> | default = 0]
+
 # Limit how long back data (series and metadata) can be queried, up until
 # <lookback> duration ago. This limit is enforced in the query-frontend, querier
 # and ruler. If the requested time range is outside the allowed range, the

pkg/distributor/distributor_test.go

@@ -50,7 +50,6 @@ import (
 var (
 	errFail       = fmt.Errorf("Fail")
 	emptyResponse = &cortexpb.WriteResponse{}
-	ctx           = user.InjectOrgID(context.Background(), "user")
 )
 
 func TestConfig_Validate(t *testing.T) {
@@ -110,6 +109,7 @@ func TestDistributor_Push(t *testing.T) {
 	lastSeenTimestamp := "cortex_distributor_latest_seen_sample_timestamp_seconds"
 	distributorAppend := "cortex_distributor_ingester_appends_total"
 	distributorAppendFailure := "cortex_distributor_ingester_append_failures_total"
+	ctx := user.InjectOrgID(context.Background(), "user")
 
 	type samplesIn struct {
 		num              int
@@ -380,6 +380,7 @@ func TestDistributor_PushIngestionRateLimiter(t *testing.T) {
 		expectedError error
 	}
 
+	ctx := user.InjectOrgID(context.Background(), "user")
 	tests := map[string]struct {
 		distributors          int
 		ingestionRateStrategy string
@@ -469,12 +470,14 @@ func TestDistributor_PushIngestionRateLimiter(t *testing.T) {
 }
 
 func TestDistributor_PushInstanceLimits(t *testing.T) {
+
 	type testPush struct {
 		samples       int
 		metadata      int
 		expectedError error
 	}
 
+	ctx := user.InjectOrgID(context.Background(), "user")
 	tests := map[string]struct {
 		preInflight    int
 		preRateSamples int        // initial rate before first push
@@ -618,7 +621,7 @@ func TestDistributor_PushInstanceLimits(t *testing.T) {
 }
 
 func TestDistributor_PushHAInstances(t *testing.T) {
-	ctx = user.InjectOrgID(context.Background(), "user")
+	ctx := user.InjectOrgID(context.Background(), "user")
 
 	for i, tc := range []struct {
 		enableTracker    bool
@@ -720,6 +723,7 @@ func TestDistributor_PushHAInstances(t *testing.T) {
 func TestDistributor_PushQuery(t *testing.T) {
 	const shuffleShardSize = 5
 
+	ctx := user.InjectOrgID(context.Background(), "user")
 	nameMatcher := mustEqualMatcher(model.MetricNameLabel, "foo")
 	barMatcher := mustEqualMatcher("bar", "baz")
 
@@ -895,6 +899,7 @@ func TestDistributor_PushQuery(t *testing.T) {
 func TestDistributor_QueryStream_ShouldReturnErrorIfMaxChunksPerQueryLimitIsReached(t *testing.T) {
 	const maxChunksLimit = 30 // Chunks are duplicated due to replication factor.
 
+	ctx := user.InjectOrgID(context.Background(), "user")
 	limits := &validation.Limits{}
 	flagext.DefaultValues(limits)
 	limits.MaxChunksPerQuery = maxChunksLimit
@@ -949,9 +954,11 @@ func TestDistributor_QueryStream_ShouldReturnErrorIfMaxChunksPerQueryLimitIsReac
 func TestDistributor_QueryStream_ShouldReturnErrorIfMaxSeriesPerQueryLimitIsReached(t *testing.T) {
 	const maxSeriesLimit = 10
 
+	ctx := user.InjectOrgID(context.Background(), "user")
 	limits := &validation.Limits{}
 	flagext.DefaultValues(limits)
-	ctx = limiter.AddQueryLimiterToContext(ctx, limiter.NewQueryLimiter(maxSeriesLimit))
+	ctx = limiter.AddQueryLimiterToContext(ctx, limiter.NewQueryLimiter(maxSeriesLimit, 0))
+
 	// Prepare distributors.
 	ds, _, r, _ := prepare(t, prepConfig{
 		numIngesters:     3,
@@ -994,10 +1001,78 @@ func TestDistributor_QueryStream_ShouldReturnErrorIfMaxSeriesPerQueryLimitIsReac
 	_, err = ds[0].QueryStream(ctx, math.MinInt32, math.MaxInt32, allSeriesMatchers...)
 	require.Error(t, err)
 	assert.Contains(t, err.Error(), "max number of series limit")
+
+}
+
+func TestDistributor_QueryStream_ShouldReturnErrorIfMaxChunkBytesPerQueryLimitIsReached(t *testing.T) {
+	const seriesToAdd = 10
+
+	ctx := user.InjectOrgID(context.Background(), "user")
+	limits := &validation.Limits{}
+	flagext.DefaultValues(limits)
+
+	// Prepare distributors.
+	ds, _, r, _ := prepare(t, prepConfig{
+		numIngesters:     3,
+		happyIngesters:   3,
+		numDistributors:  1,
+		shardByAllLabels: true,
+		limits:           limits,
+	})
+	defer stopAll(ds, r)
+
+	allSeriesMatchers := []*labels.Matcher{
+		labels.MustNewMatcher(labels.MatchRegexp, model.MetricNameLabel, ".+"),
+	}
+	// Push a single series to allow us to calculate the chunk size to calculate the limit for the test.
+	writeReq := &cortexpb.WriteRequest{}
+	writeReq.Timeseries = append(writeReq.Timeseries,
+		makeWriteRequestTimeseries([]cortexpb.LabelAdapter{{Name: model.MetricNameLabel, Value: "another_series"}}, 0, 0),
+	)
+	writeRes, err := ds[0].Push(ctx, writeReq)
+	assert.Equal(t, &cortexpb.WriteResponse{}, writeRes)
+	assert.Nil(t, err)
+	chunkSizeResponse, err := ds[0].QueryStream(ctx, math.MinInt32, math.MaxInt32, allSeriesMatchers...)
+	require.NoError(t, err)
+
+	// Use the resulting chunks size to calculate the limit as (series to add + our test series) * the response chunk size.
+	var responseChunkSize = chunkSizeResponse.ChunksSize()
+	var maxBytesLimit = (seriesToAdd) * responseChunkSize
+
+	// Update the limiter with the calculated limits.
+	ctx = limiter.AddQueryLimiterToContext(ctx, limiter.NewQueryLimiter(0, maxBytesLimit))
+
+	// Push a number of series below the max chunk bytes limit. Subtract one for the series added above.
+	writeReq = makeWriteRequest(0, seriesToAdd-1, 0)
+	writeRes, err = ds[0].Push(ctx, writeReq)
+	assert.Equal(t, &cortexpb.WriteResponse{}, writeRes)
+	assert.Nil(t, err)
+
+	// Since the number of chunk bytes is equal to the limit (but doesn't
+	// exceed it), we expect a query running on all series to succeed.
+	queryRes, err := ds[0].QueryStream(ctx, math.MinInt32, math.MaxInt32, allSeriesMatchers...)
+	require.NoError(t, err)
+	assert.Len(t, queryRes.Chunkseries, seriesToAdd)
+
+	// Push another series to exceed the chunk bytes limit once we'll query back all series.
+	writeReq = &cortexpb.WriteRequest{}
+	writeReq.Timeseries = append(writeReq.Timeseries,
+		makeWriteRequestTimeseries([]cortexpb.LabelAdapter{{Name: model.MetricNameLabel, Value: "another_series_1"}}, 0, 0),
+	)
+
+	writeRes, err = ds[0].Push(ctx, writeReq)
+	assert.Equal(t, &cortexpb.WriteResponse{}, writeRes)
+	assert.Nil(t, err)
+
+	// Since the aggregated chunk size is exceeding the limit, we expect
+	// a query running on all series to fail.
+	_, err = ds[0].QueryStream(ctx, math.MinInt32, math.MaxInt32, allSeriesMatchers...)
+	require.Error(t, err)
+	assert.Equal(t, err, validation.LimitError(fmt.Sprintf(limiter.ErrMaxChunkBytesHit, maxBytesLimit)))
 }
 
 func TestDistributor_Push_LabelRemoval(t *testing.T) {
-	ctx = user.InjectOrgID(context.Background(), "user")
+	ctx := user.InjectOrgID(context.Background(), "user")
 
 	type testcase struct {
 		inputSeries    labels.Labels
@@ -1086,6 +1161,7 @@ func TestDistributor_Push_LabelRemoval(t *testing.T) {
 }
 
 func TestDistributor_Push_ShouldGuaranteeShardingTokenConsistencyOverTheTime(t *testing.T) {
+	ctx := user.InjectOrgID(context.Background(), "user")
 	tests := map[string]struct {
 		inputSeries    labels.Labels
 		expectedSeries labels.Labels
@@ -1161,8 +1237,6 @@ func TestDistributor_Push_ShouldGuaranteeShardingTokenConsistencyOverTheTime(t *
 	limits.DropLabels = []string{"dropped"}
 	limits.AcceptHASamples = true
 
-	ctx = user.InjectOrgID(context.Background(), "user")
-
 	for testName, testData := range tests {
 		t.Run(testName, func(t *testing.T) {
 			ds, ingesters, r, _ := prepare(t, prepConfig{
@@ -1198,6 +1272,8 @@ func TestDistributor_Push_LabelNameValidation(t *testing.T) {
 		{Name: model.MetricNameLabel, Value: "foo"},
 		{Name: "999.illegal", Value: "baz"},
 	}
+	ctx := user.InjectOrgID(context.Background(), "user")
+
 	tests := map[string]struct {
 		inputLabels                labels.Labels
 		skipLabelNameValidationCfg bool
@@ -1245,7 +1321,7 @@ func TestDistributor_Push_LabelNameValidation(t *testing.T) {
 }
 
 func TestDistributor_Push_ExemplarValidation(t *testing.T) {
-
+	ctx := user.InjectOrgID(context.Background(), "user")
 	manyLabels := []string{model.MetricNameLabel, "test"}
 	for i := 1; i < 31; i++ {
 		manyLabels = append(manyLabels, fmt.Sprintf("name_%d", i), fmt.Sprintf("value_%d", i))
@@ -1292,7 +1368,6 @@ func TestDistributor_Push_ExemplarValidation(t *testing.T) {
 				numDistributors:  1,
 				shuffleShardSize: 1,
 			})
-
 			_, err := ds[0].Push(ctx, tc.req)
 			if tc.errMsg != "" {
 				fromError, _ := status.FromError(err)
@@ -1308,6 +1383,7 @@ func BenchmarkDistributor_Push(b *testing.B) {
 	const (
 		numSeriesPerRequest = 1000
 	)
+	ctx := user.InjectOrgID(context.Background(), "user")
 
 	tests := map[string]struct {
 		prepareConfig func(limits *validation.Limits)
@@ -1499,6 +1575,7 @@ func BenchmarkDistributor_Push(b *testing.B) {
 
 	for testName, testData := range tests {
 		b.Run(testName, func(b *testing.B) {
+
 			// Create an in-memory KV store for the ring with 1 ingester registered.
 			kvStore := consul.NewInMemoryClient(ring.GetCodec())
 			err := kvStore.CAS(context.Background(), ring.IngesterRingKey,
@@ -1573,6 +1650,7 @@ func BenchmarkDistributor_Push(b *testing.B) {
 }
 
 func TestSlowQueries(t *testing.T) {
+	ctx := user.InjectOrgID(context.Background(), "user")
 	nameMatcher := mustEqualMatcher(model.MetricNameLabel, "foo")
 	nIngesters := 3
 	for _, shardByAllLabels := range []bool{true, false} {
@@ -2526,7 +2604,7 @@ func TestSortLabels(t *testing.T) {
 }
 
 func TestDistributor_Push_Relabel(t *testing.T) {
-	ctx = user.InjectOrgID(context.Background(), "user")
+	ctx := user.InjectOrgID(context.Background(), "user")
 
 	type testcase struct {
 		inputSeries          labels.Labels

pkg/distributor/ha_tracker_test.go

@@ -15,6 +15,7 @@ import (
 	"github.com/prometheus/prometheus/pkg/timestamp"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"github.com/weaveworks/common/user"
 
 	"github.com/cortexproject/cortex/pkg/cortexpb"
 	"github.com/cortexproject/cortex/pkg/ring"
@@ -660,7 +661,8 @@ func TestHATracker_MetricsCleanup(t *testing.T) {
 func TestCheckReplicaCleanup(t *testing.T) {
 	replica := "r1"
 	cluster := "c1"
-	user := "user"
+	userID := "user"
+	ctx := user.InjectOrgID(context.Background(), userID)
 
 	reg := prometheus.NewPedanticRegistry()
 
@@ -678,32 +680,32 @@ func TestCheckReplicaCleanup(t *testing.T) {
 
 	now := time.Now()
 
-	err = c.checkReplica(context.Background(), user, cluster, replica, now)
+	err = c.checkReplica(context.Background(), userID, cluster, replica, now)
 	assert.NoError(t, err)
-	checkReplicaTimestamp(t, time.Second, c, user, cluster, replica, now)
+	checkReplicaTimestamp(t, time.Second, c, userID, cluster, replica, now)
 
 	// Replica is not marked for deletion yet.
-	checkReplicaDeletionState(t, time.Second, c, user, cluster, true, true, false)
+	checkReplicaDeletionState(t, time.Second, c, userID, cluster, true, true, false)
 
 	// This will mark replica for deletion (with time.Now())
 	c.cleanupOldReplicas(ctx, now.Add(1*time.Second))
 
 	// Verify marking for deletion.
-	checkReplicaDeletionState(t, time.Second, c, user, cluster, false, true, true)
+	checkReplicaDeletionState(t, time.Second, c, userID, cluster, false, true, true)
 
 	// This will "revive" the replica.
 	now = time.Now()
-	err = c.checkReplica(context.Background(), user, cluster, replica, now)
+	err = c.checkReplica(context.Background(), userID, cluster, replica, now)
 	assert.NoError(t, err)
-	checkReplicaTimestamp(t, time.Second, c, user, cluster, replica, now) // This also checks that entry is not marked for deletion.
+	checkReplicaTimestamp(t, time.Second, c, userID, cluster, replica, now) // This also checks that entry is not marked for deletion.
 
 	// This will mark replica for deletion again (with new time.Now())
 	c.cleanupOldReplicas(ctx, now.Add(1*time.Second))
-	checkReplicaDeletionState(t, time.Second, c, user, cluster, false, true, true)
+	checkReplicaDeletionState(t, time.Second, c, userID, cluster, false, true, true)
 
 	// Delete entry marked for deletion completely.
 	c.cleanupOldReplicas(ctx, time.Now().Add(5*time.Second))
-	checkReplicaDeletionState(t, time.Second, c, user, cluster, false, false, false)
+	checkReplicaDeletionState(t, time.Second, c, userID, cluster, false, false, false)
 
 	require.NoError(t, testutil.GatherAndCompare(reg, strings.NewReader(`
 		# HELP cortex_ha_tracker_replicas_cleanup_marked_for_deletion_total Number of elected replicas marked for deletion.

pkg/distributor/query.go

@@ -232,11 +232,17 @@ func (d *Distributor) queryIngesterStream(ctx context.Context, userID string, re
 					return nil, validation.LimitError(fmt.Sprintf(errMaxChunksPerQueryLimit, util.LabelMatchersToString(matchers), chunksLimit))
 				}
 			}
+
 			for _, series := range resp.Chunkseries {
 				if limitErr := queryLimiter.AddSeries(series.Labels); limitErr != nil {
 					return nil, limitErr
 				}
 			}
+
+			if chunkBytesLimitErr := queryLimiter.AddChunkBytes(resp.ChunksSize()); chunkBytesLimitErr != nil {
+				return nil, chunkBytesLimitErr
+			}
+
 			for _, series := range resp.Timeseries {
 				if limitErr := queryLimiter.AddSeries(series.Labels); limitErr != nil {
 					return nil, limitErr

pkg/ingester/client/custom.go

@@ -12,3 +12,18 @@ func (m *QueryStreamResponse) ChunksCount() int {
 	}
 	return count
 }
+
+// ChunksSize returns the size of all chunks in the response.
+func (m *QueryStreamResponse) ChunksSize() int {
+	if len(m.Chunkseries) == 0 {
+		return 0
+	}
+
+	size := 0
+	for _, entry := range m.Chunkseries {
+		for _, chunk := range entry.Chunks {
+			size += chunk.Size()
+		}
+	}
+	return size
+}

pkg/querier/blocks_store_queryable.go

@@ -626,6 +626,13 @@ func (q *blocksStoreQuerier) fetchSeriesFromStores(
 							return validation.LimitError(fmt.Sprintf(errMaxChunksPerQueryLimit, util.LabelMatchersToString(matchers), maxChunksLimit))
 						}
 					}
+					chunksSize := 0
+					for _, c := range s.Chunks {
+						chunksSize += c.Size()
+					}
+					if chunkBytesLimitErr := queryLimiter.AddChunkBytes(chunksSize); chunkBytesLimitErr != nil {
+						return chunkBytesLimitErr
+					}
 				}
 
 				if w := resp.GetWarning(); w != "" {