[Store Gateway] Token bucket limiter

CHANGELOG.md

@@ -4,6 +4,7 @@
 * [CHANGE] Upgrade Dockerfile Node version from 14x to 18x. #5906
 * [CHANGE] Ingester: Remove `-querier.query-store-for-labels-enabled` flag. Querying long-term store for labels is always enabled. #5984
 * [FEATURE] Ingester: Experimental: Enable native histogram ingestion via `-blocks-storage.tsdb.enable-native-histograms` flag. #5986
+* [FEATURE] Store Gateway: Token bucket limiter. #6016
 * [ENHANCEMENT] rulers: Add support to persist tokens in rulers. #5987
 * [ENHANCEMENT] Query Frontend/Querier: Added store gateway postings touched count and touched size in Querier stats and log in Query Frontend. #5892
 * [ENHANCEMENT] Query Frontend/Querier: Returns `warnings` on prometheus query responses. #5916

docs/blocks-storage/querier.md

@@ -1341,6 +1341,51 @@ blocks_storage:
     # CLI flag: -blocks-storage.bucket-store.series-batch-size
     [series_batch_size: <int> | default = 10000]
 
+    token_bucket_limiter:
+      # Whether token bucket limiter is enabled
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.enabled
+      [enabled: <boolean> | default = false]
+
+      # Whether the token bucket limiter is in dry run mode
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.dry-run
+      [dry_run: <boolean> | default = false]
+
+      # Pod token bucket size
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.pod-token-bucket-size
+      [pod_token_bucket_size: <int> | default = 859832320]
+
+      # User token bucket size
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.user-token-bucket-size
+      [user_token_bucket_size: <int> | default = 644874240]
+
+      # Request token bucket size
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.request-token-bucket-size
+      [request_token_bucket_size: <int> | default = 4194304]
+
+      # Multiplication factor used for fetched postings token
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-postings-token-factor
+      [fetched_postings_token_factor: <float> | default = 2]
+
+      # Multiplication factor used for touched postings token
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-postings-token-factor
+      [touched_postings_token_factor: <float> | default = 2]
+
+      # Multiplication factor used for fetched series token
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-series-token-factor
+      [fetched_series_token_factor: <float> | default = 2.5]
+
+      # Multiplication factor used for touched series token
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-series-token-factor
+      [touched_series_token_factor: <float> | default = 10]
+
+      # Multiplication factor used for fetched chunks token
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-chunks-token-factor
+      [fetched_chunks_token_factor: <float> | default = 0.5]
+
+      # Multiplication factor used for touched chunks token
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-chunks-token-factor
+      [touched_chunks_token_factor: <float> | default = 0.5]
+
   tsdb:
     # Local directory to store TSDBs in the ingesters.
     # CLI flag: -blocks-storage.tsdb.dir

docs/blocks-storage/store-gateway.md

@@ -1466,6 +1466,51 @@ blocks_storage:
     # CLI flag: -blocks-storage.bucket-store.series-batch-size
     [series_batch_size: <int> | default = 10000]
 
+    token_bucket_limiter:
+      # Whether token bucket limiter is enabled
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.enabled
+      [enabled: <boolean> | default = false]
+
+      # Whether the token bucket limiter is in dry run mode
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.dry-run
+      [dry_run: <boolean> | default = false]
+
+      # Pod token bucket size
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.pod-token-bucket-size
+      [pod_token_bucket_size: <int> | default = 859832320]
+
+      # User token bucket size
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.user-token-bucket-size
+      [user_token_bucket_size: <int> | default = 644874240]
+
+      # Request token bucket size
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.request-token-bucket-size
+      [request_token_bucket_size: <int> | default = 4194304]
+
+      # Multiplication factor used for fetched postings token
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-postings-token-factor
+      [fetched_postings_token_factor: <float> | default = 2]
+
+      # Multiplication factor used for touched postings token
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-postings-token-factor
+      [touched_postings_token_factor: <float> | default = 2]
+
+      # Multiplication factor used for fetched series token
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-series-token-factor
+      [fetched_series_token_factor: <float> | default = 2.5]
+
+      # Multiplication factor used for touched series token
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-series-token-factor
+      [touched_series_token_factor: <float> | default = 10]
+
+      # Multiplication factor used for fetched chunks token
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-chunks-token-factor
+      [fetched_chunks_token_factor: <float> | default = 0.5]
+
+      # Multiplication factor used for touched chunks token
+      # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-chunks-token-factor
+      [touched_chunks_token_factor: <float> | default = 0.5]
+
   tsdb:
     # Local directory to store TSDBs in the ingesters.
     # CLI flag: -blocks-storage.tsdb.dir

docs/configuration/config-file-reference.md

@@ -1899,6 +1899,51 @@ bucket_store:
   # CLI flag: -blocks-storage.bucket-store.series-batch-size
   [series_batch_size: <int> | default = 10000]
 
+  token_bucket_limiter:
+    # Whether token bucket limiter is enabled
+    # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.enabled
+    [enabled: <boolean> | default = false]
+
+    # Whether the token bucket limiter is in dry run mode
+    # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.dry-run
+    [dry_run: <boolean> | default = false]
+
+    # Pod token bucket size
+    # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.pod-token-bucket-size
+    [pod_token_bucket_size: <int> | default = 859832320]
+
+    # User token bucket size
+    # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.user-token-bucket-size
+    [user_token_bucket_size: <int> | default = 644874240]
+
+    # Request token bucket size
+    # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.request-token-bucket-size
+    [request_token_bucket_size: <int> | default = 4194304]
+
+    # Multiplication factor used for fetched postings token
+    # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-postings-token-factor
+    [fetched_postings_token_factor: <float> | default = 2]
+
+    # Multiplication factor used for touched postings token
+    # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-postings-token-factor
+    [touched_postings_token_factor: <float> | default = 2]
+
+    # Multiplication factor used for fetched series token
+    # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-series-token-factor
+    [fetched_series_token_factor: <float> | default = 2.5]
+
+    # Multiplication factor used for touched series token
+    # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-series-token-factor
+    [touched_series_token_factor: <float> | default = 10]
+
+    # Multiplication factor used for fetched chunks token
+    # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-chunks-token-factor
+    [fetched_chunks_token_factor: <float> | default = 0.5]
+
+    # Multiplication factor used for touched chunks token
+    # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-chunks-token-factor
+    [touched_chunks_token_factor: <float> | default = 0.5]
+
 tsdb:
   # Local directory to store TSDBs in the ingesters.
   # CLI flag: -blocks-storage.tsdb.dir

pkg/storage/tsdb/config.go

@@ -292,6 +292,23 @@ type BucketStoreConfig struct {
 
 	// Controls how many series to fetch per batch in Store Gateway. Default value is 10000.
 	SeriesBatchSize int `yaml:"series_batch_size"`
+
+	// Token bucket configs
+	TokenBucketLimiter TokenBucketLimiterConfig `yaml:"token_bucket_limiter"`
+}
+
+type TokenBucketLimiterConfig struct {
+	Enabled                    bool    `yaml:"enabled"`
+	DryRun                     bool    `yaml:"dry_run"`
+	PodTokenBucketSize         int64   `yaml:"pod_token_bucket_size"`
+	UserTokenBucketSize        int64   `yaml:"user_token_bucket_size"`
+	RequestTokenBucketSize     int64   `yaml:"request_token_bucket_size"`
+	FetchedPostingsTokenFactor float64 `yaml:"fetched_postings_token_factor"`
+	TouchedPostingsTokenFactor float64 `yaml:"touched_postings_token_factor"`
+	FetchedSeriesTokenFactor   float64 `yaml:"fetched_series_token_factor"`
+	TouchedSeriesTokenFactor   float64 `yaml:"touched_series_token_factor"`
+	FetchedChunksTokenFactor   float64 `yaml:"fetched_chunks_token_factor"`
+	TouchedChunksTokenFactor   float64 `yaml:"touched_chunks_token_factor"`
 }
 
 // RegisterFlags registers the BucketStore flags
@@ -325,6 +342,17 @@ func (cfg *BucketStoreConfig) RegisterFlags(f *flag.FlagSet) {
 	f.BoolVar(&cfg.LazyExpandedPostingsEnabled, "blocks-storage.bucket-store.lazy-expanded-postings-enabled", false, "If true, Store Gateway will estimate postings size and try to lazily expand postings if it downloads less data than expanding all postings.")
 	f.IntVar(&cfg.SeriesBatchSize, "blocks-storage.bucket-store.series-batch-size", store.SeriesBatchSize, "Controls how many series to fetch per batch in Store Gateway. Default value is 10000.")
 	f.StringVar(&cfg.BlockDiscoveryStrategy, "blocks-storage.bucket-store.block-discovery-strategy", string(ConcurrentDiscovery), "One of "+strings.Join(supportedBlockDiscoveryStrategies, ", ")+". When set to concurrent, stores will concurrently issue one call per directory to discover active blocks in the bucket. The recursive strategy iterates through all objects in the bucket, recursively traversing into each directory. This avoids N+1 calls at the expense of having slower bucket iterations. bucket_index strategy can be used in Compactor only and utilizes the existing bucket index to fetch block IDs to sync. This avoids iterating the bucket but can be impacted by delays of cleaner creating bucket index.")
+	f.BoolVar(&cfg.TokenBucketLimiter.Enabled, "blocks-storage.bucket-store.token-bucket-limiter.enabled", false, "Whether token bucket limiter is enabled")
+	f.BoolVar(&cfg.TokenBucketLimiter.DryRun, "blocks-storage.bucket-store.token-bucket-limiter.dry-run", false, "Whether the token bucket limiter is in dry run mode")
+	f.Int64Var(&cfg.TokenBucketLimiter.PodTokenBucketSize, "blocks-storage.bucket-store.token-bucket-limiter.pod-token-bucket-size", int64(820*units.Mebibyte), "Pod token bucket size")
+	f.Int64Var(&cfg.TokenBucketLimiter.UserTokenBucketSize, "blocks-storage.bucket-store.token-bucket-limiter.user-token-bucket-size", int64(615*units.Mebibyte), "User token bucket size")
+	f.Int64Var(&cfg.TokenBucketLimiter.RequestTokenBucketSize, "blocks-storage.bucket-store.token-bucket-limiter.request-token-bucket-size", int64(4*units.Mebibyte), "Request token bucket size")
+	f.Float64Var(&cfg.TokenBucketLimiter.FetchedPostingsTokenFactor, "blocks-storage.bucket-store.token-bucket-limiter.fetched-postings-token-factor", 2, "Multiplication factor used for fetched postings token")
+	f.Float64Var(&cfg.TokenBucketLimiter.TouchedPostingsTokenFactor, "blocks-storage.bucket-store.token-bucket-limiter.touched-postings-token-factor", 2, "Multiplication factor used for touched postings token")
+	f.Float64Var(&cfg.TokenBucketLimiter.FetchedSeriesTokenFactor, "blocks-storage.bucket-store.token-bucket-limiter.fetched-series-token-factor", 2.5, "Multiplication factor used for fetched series token")
+	f.Float64Var(&cfg.TokenBucketLimiter.TouchedSeriesTokenFactor, "blocks-storage.bucket-store.token-bucket-limiter.touched-series-token-factor", 10, "Multiplication factor used for touched series token")
+	f.Float64Var(&cfg.TokenBucketLimiter.FetchedChunksTokenFactor, "blocks-storage.bucket-store.token-bucket-limiter.fetched-chunks-token-factor", 0.5, "Multiplication factor used for fetched chunks token")
+	f.Float64Var(&cfg.TokenBucketLimiter.TouchedChunksTokenFactor, "blocks-storage.bucket-store.token-bucket-limiter.touched-chunks-token-factor", 0.5, "Multiplication factor used for touched chunks token")
 }
 
 // Validate the config.

pkg/storage/tsdb/inmemory_index_cache.go

@@ -103,7 +103,7 @@ func (c *InMemoryIndexCache) set(typ string, key storecache.CacheKey, val []byte
 
 	size := uint64(len(k) + len(val))
 	if size > c.maxItemSizeBytes {
-		level.Info(c.logger).Log(
+		level.Debug(c.logger).Log(
 			"msg", "item bigger than maxItemSizeBytes. Ignoring..",
 			"maxItemSizeBytes", c.maxItemSizeBytes,
 			"cacheType", typ,

pkg/storegateway/bucket_stores.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 	"math"
-	"net/http"
 	"os"
 	"path/filepath"
 	"strings"
@@ -35,6 +34,7 @@ import (
 
 	"github.com/cortexproject/cortex/pkg/storage/bucket"
 	"github.com/cortexproject/cortex/pkg/storage/tsdb"
+	"github.com/cortexproject/cortex/pkg/util"
 	"github.com/cortexproject/cortex/pkg/util/backoff"
 	cortex_errors "github.com/cortexproject/cortex/pkg/util/errors"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
@@ -73,6 +73,11 @@ type BucketStores struct {
 	storesErrorsMu sync.RWMutex
 	storesErrors   map[string]error
 
+	podTokenBucket *util.TokenBucket
+
+	userTokenBucketsMu sync.RWMutex
+	userTokenBuckets   map[string]*util.TokenBucket
+
 	// Keeps number of inflight requests
 	inflightRequestCnt int
 	inflightRequestMu  sync.RWMutex
@@ -115,6 +120,7 @@ func NewBucketStores(cfg tsdb.BlocksStorageConfig, shardingStrategy ShardingStra
 		metaFetcherMetrics: NewMetadataFetcherMetrics(),
 		queryGate:          queryGate,
 		partitioner:        newGapBasedPartitioner(cfg.BucketStore.PartitionerMaxGapBytes, reg),
+		userTokenBuckets:   make(map[string]*util.TokenBucket),
 		syncTimes: promauto.With(reg).NewHistogram(prometheus.HistogramOpts{
 			Name:    "cortex_bucket_stores_blocks_sync_seconds",
 			Help:    "The total time it takes to perform a sync stores",
@@ -144,6 +150,13 @@ func NewBucketStores(cfg tsdb.BlocksStorageConfig, shardingStrategy ShardingStra
 		return nil, errors.Wrap(err, "create chunks bytes pool")
 	}
 
+	if u.cfg.BucketStore.TokenBucketLimiter.Enabled {
+		u.podTokenBucket = util.NewTokenBucket(cfg.BucketStore.TokenBucketLimiter.PodTokenBucketSize, cfg.BucketStore.TokenBucketLimiter.PodTokenBucketSize, promauto.With(reg).NewGauge(prometheus.GaugeOpts{
+			Name: "cortex_bucket_stores_pod_token_bucket_remaining",
+			Help: "Number of tokens left in pod token bucket.",
+		}))
+	}
+
 	if reg != nil {
 		reg.MustRegister(u.bucketStoreMetrics, u.metaFetcherMetrics)
 	}
@@ -475,6 +488,12 @@ func (u *BucketStores) closeEmptyBucketStore(userID string) error {
 	unlockInDefer = false
 	u.storesMu.Unlock()
 
+	if u.cfg.BucketStore.TokenBucketLimiter.Enabled {
+		u.userTokenBucketsMu.Lock()
+		delete(u.userTokenBuckets, userID)
+		u.userTokenBucketsMu.Unlock()
+	}
+
 	u.metaFetcherMetrics.RemoveUserRegistry(userID)
 	u.bucketStoreMetrics.RemoveUserRegistry(userID)
 	return bs.Close()
@@ -612,13 +631,19 @@ func (u *BucketStores) getOrCreateStore(userID string) (*store.BucketStore, erro
 		bucketStoreOpts = append(bucketStoreOpts, store.WithDebugLogging())
 	}
 
+	if u.cfg.BucketStore.TokenBucketLimiter.Enabled {
+		u.userTokenBucketsMu.Lock()
+		u.userTokenBuckets[userID] = util.NewTokenBucket(u.cfg.BucketStore.TokenBucketLimiter.UserTokenBucketSize, u.cfg.BucketStore.TokenBucketLimiter.UserTokenBucketSize, nil)
+		u.userTokenBucketsMu.Unlock()
+	}
+
 	bs, err := store.NewBucketStore(
 		userBkt,
 		fetcher,
 		u.syncDirForUser(userID),
 		newChunksLimiterFactory(u.limits, userID),
 		newSeriesLimiterFactory(u.limits, userID),
-		newBytesLimiterFactory(u.limits, userID),
+		newBytesLimiterFactory(u.limits, userID, u.podTokenBucket, u.getUserTokenBucket(userID), u.cfg.BucketStore.TokenBucketLimiter, u.getTokensToRetrieve),
 		u.partitioner,
 		u.cfg.BucketStore.BlockSyncConcurrency,
 		false, // No need to enable backward compatibility with Thanos pre 0.8.0 queriers
@@ -680,6 +705,31 @@ func (u *BucketStores) deleteLocalFilesForExcludedTenants(includeUserIDs map[str
 	}
 }
 
+func (u *BucketStores) getUserTokenBucket(userID string) *util.TokenBucket {
+	u.userTokenBucketsMu.RLock()
+	defer u.userTokenBucketsMu.RUnlock()
+	return u.userTokenBuckets[userID]
+}
+
+func (u *BucketStores) getTokensToRetrieve(tokens uint64, dataType store.StoreDataType) int64 {
+	tokensToRetrieve := float64(tokens)
+	switch dataType {
+	case store.PostingsFetched:
+		tokensToRetrieve *= u.cfg.BucketStore.TokenBucketLimiter.FetchedPostingsTokenFactor
+	case store.PostingsTouched:
+		tokensToRetrieve *= u.cfg.BucketStore.TokenBucketLimiter.TouchedPostingsTokenFactor
+	case store.SeriesFetched:
+		tokensToRetrieve *= u.cfg.BucketStore.TokenBucketLimiter.FetchedSeriesTokenFactor
+	case store.SeriesTouched:
+		tokensToRetrieve *= u.cfg.BucketStore.TokenBucketLimiter.TouchedSeriesTokenFactor
+	case store.ChunksFetched:
+		tokensToRetrieve *= u.cfg.BucketStore.TokenBucketLimiter.FetchedChunksTokenFactor
+	case store.ChunksTouched:
+		tokensToRetrieve *= u.cfg.BucketStore.TokenBucketLimiter.TouchedChunksTokenFactor
+	}
+	return int64(tokensToRetrieve)
+}
+
 func getUserIDFromGRPCContext(ctx context.Context) string {
 	meta, ok := metadata.FromIncomingContext(ctx)
 	if !ok {
@@ -730,50 +780,3 @@ type spanSeriesServer struct {
 func (s spanSeriesServer) Context() context.Context {
 	return s.ctx
 }
-
-type limiter struct {
-	limiter *store.Limiter
-}
-
-func (c *limiter) Reserve(num uint64) error {
-	return c.ReserveWithType(num, 0)
-}
-
-func (c *limiter) ReserveWithType(num uint64, _ store.StoreDataType) error {
-	err := c.limiter.Reserve(num)
-	if err != nil {
-		return httpgrpc.Errorf(http.StatusUnprocessableEntity, err.Error())
-	}
-
-	return nil
-}
-
-func newChunksLimiterFactory(limits *validation.Overrides, userID string) store.ChunksLimiterFactory {
-	return func(failedCounter prometheus.Counter) store.ChunksLimiter {
-		// 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.MaxChunksPerQueryFromStore(userID)), failedCounter),
-		}
-	}
-}
-
-func newSeriesLimiterFactory(limits *validation.Overrides, userID string) store.SeriesLimiterFactory {
-	return func(failedCounter prometheus.Counter) store.SeriesLimiter {
-		// 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.MaxFetchedSeriesPerQuery(userID)), failedCounter),
-		}
-	}
-}
-
-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),
-		}
-	}
-}