add limit for store gateway downloaded bytes

Signed-off-by: Ben Ye <[email protected]>

Author: yeya24

pkg/storegateway/bucket_stores.go

@@ -492,7 +492,7 @@ func (u *BucketStores) getOrCreateStore(userID string) (*store.BucketStore, erro
 		u.syncDirForUser(userID),
 		newChunksLimiterFactory(u.limits, userID),
 		newSeriesLimiterFactory(u.limits, userID),
-		store.NewBytesLimiterFactory(0),
+		newBytesLimiterFactory(u.limits, userID),
 		u.partitioner,
 		u.cfg.BucketStore.BlockSyncConcurrency,
 		false, // No need to enable backward compatibility with Thanos pre 0.8.0 queriers
@@ -637,3 +637,13 @@ func newSeriesLimiterFactory(limits *validation.Overrides, userID string) store.
 		}
 	}
 }
+
+func newBytesLimiterFactory(limits *validation.Overrides, userID string) store.BytesLimiterFactory {
+	return func(failedCounter prometheus.Counter) store.BytesLimiter {
+		// Since limit overrides could be live reloaded, we have to get the current user's limit
+		// each time a new limiter is instantiated.
+		return &limiter{
+			limiter: store.NewLimiter(uint64(limits.MaxDownloadedBytesPerRequest(userID)), failedCounter),
+		}
+	}
+}

pkg/util/validation/limits.go

@@ -93,7 +93,8 @@ type Limits struct {
 	RulerMaxRuleGroupsPerTenant int            `yaml:"ruler_max_rule_groups_per_tenant" json:"ruler_max_rule_groups_per_tenant"`
 
 	// Store-gateway.
-	StoreGatewayTenantShardSize int `yaml:"store_gateway_tenant_shard_size" json:"store_gateway_tenant_shard_size"`
+	StoreGatewayTenantShardSize  int `yaml:"store_gateway_tenant_shard_size" json:"store_gateway_tenant_shard_size"`
+	MaxDownloadedBytesPerRequest int `yaml:"max_downloaded_bytes_per_request" json:"max_downloaded_bytes_per_request"`
 
 	// Compactor.
 	CompactorBlocksRetentionPeriod model.Duration `yaml:"compactor_blocks_retention_period" json:"compactor_blocks_retention_period"`
@@ -179,6 +180,7 @@ func (l *Limits) RegisterFlags(f *flag.FlagSet) {
 
 	// Store-gateway.
 	f.IntVar(&l.StoreGatewayTenantShardSize, "store-gateway.tenant-shard-size", 0, "The default tenant's shard size when the shuffle-sharding strategy is used. Must be set when the store-gateway sharding is enabled with the shuffle-sharding strategy. When this setting is specified in the per-tenant overrides, a value of 0 disables shuffle sharding for the tenant.")
+	f.IntVar(&l.MaxDownloadedBytesPerRequest, "store-gateway.max-downloaded-bytes-per-request", 0, "The maximum number of data bytes to download per gRPC request in Store Gateway, including Series/LabelNames/LabelValues requests. 0 to disable.")
 
 	// Alertmanager.
 	f.Var(&l.AlertmanagerReceiversBlockCIDRNetworks, "alertmanager.receivers-firewall-block-cidr-networks", "Comma-separated list of network CIDRs to block in Alertmanager receiver integrations.")
@@ -422,6 +424,12 @@ func (o *Overrides) MaxFetchedDataBytesPerQuery(userID string) int {
 	return o.GetOverridesForUser(userID).MaxFetchedDataBytesPerQuery
 }
 
+// MaxDownloadedBytesPerRequest returns the maximum number of bytes to download for each gRPC request in Store Gateway,
+// including any data fetched from cache or object storage.
+func (o *Overrides) MaxDownloadedBytesPerRequest(userID string) int {
+	return o.GetOverridesForUser(userID).MaxDownloadedBytesPerRequest
+}
+
 // MaxQueryLookback returns the max lookback period of queries.
 func (o *Overrides) MaxQueryLookback(userID string) time.Duration {
 	return time.Duration(o.GetOverridesForUser(userID).MaxQueryLookback)

pkg/util/validation/limits_test.go

@@ -564,3 +564,35 @@ tenant2:
 	require.Equal(t, 3, ov.MaxExemplars("tenant2"))
 	require.Equal(t, 5, ov.MaxExemplars("tenant3"))
 }
+
+func TestMaxDownloadedBytesPerRequestOverridesPerTenant(t *testing.T) {
+	SetDefaultLimitsForYAMLUnmarshalling(Limits{
+		MaxLabelNameLength: 100,
+	})
+
+	baseYAML := `
+max_downloaded_bytes_per_request: 5`
+	overridesYAML := `
+tenant1:
+  max_downloaded_bytes_per_request: 1
+tenant2:
+  max_downloaded_bytes_per_request: 3
+`
+
+	l := Limits{}
+	err := yaml.UnmarshalStrict([]byte(baseYAML), &l)
+	require.NoError(t, err)
+
+	overrides := map[string]*Limits{}
+	err = yaml.Unmarshal([]byte(overridesYAML), &overrides)
+	require.NoError(t, err, "parsing overrides")
+
+	tl := newMockTenantLimits(overrides)
+
+	ov, err := NewOverrides(l, tl)
+	require.NoError(t, err)
+
+	require.Equal(t, 1, ov.MaxDownloadedBytesPerRequest("tenant1"))
+	require.Equal(t, 3, ov.MaxDownloadedBytesPerRequest("tenant2"))
+	require.Equal(t, 5, ov.MaxDownloadedBytesPerRequest("tenant3"))
+}