From 2ad976a8bc5f75aa435d9570dbc7de67761f184e Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Tue, 4 Jun 2024 16:30:11 -0700 Subject: [PATCH 01/26] Create TokenBucket Signed-off-by: Justin Jung --- pkg/storage/tsdb/config.go | 22 ++++++ pkg/storage/tsdb/inmemory_index_cache.go | 2 +- pkg/util/token_bucket.go | 86 ++++++++++++++++++++++++ 3 files changed, 109 insertions(+), 1 deletion(-) create mode 100644 pkg/util/token_bucket.go diff --git a/pkg/storage/tsdb/config.go b/pkg/storage/tsdb/config.go index 5c9f0c23cb..10c50bb8fc 100644 --- a/pkg/storage/tsdb/config.go +++ b/pkg/storage/tsdb/config.go @@ -292,6 +292,18 @@ 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 + PodDataBytesRateLimit int64 `yaml:"pod_data_bytes_rate_limit"` + 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 +337,16 @@ 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.Int64Var(&cfg.PodDataBytesRateLimit, "blocks-storage.bucket-store.pod-data-bytes-rate-limit", int64(1*units.Gibibyte), "Overall data bytes rate limit for a pod") + f.Int64Var(&cfg.PodTokenBucketSize, "blocks-storage.bucket-store.pod-token-bucket-size", int64(820*units.Mebibyte), "Pod token bucket size") + f.Int64Var(&cfg.UserTokenBucketSize, "blocks-storage.bucket-store.user-token-bucket-size", int64(615*units.Mebibyte), "User token bucket size") + f.Int64Var(&cfg.RequestTokenBucketSize, "blocks-storage.bucket-store.request-token-bucket-size", int64(4*units.Mebibyte), "Request token bucket size") + f.Float64Var(&cfg.FetchedPostingsTokenFactor, "blocks-storage.bucket-store.fetched-postings-token-factor", 2, "Multiplication factor used for fetched postings token") + f.Float64Var(&cfg.TouchedPostingsTokenFactor, "blocks-storage.bucket-store.touched-postings-token-factor", 2, "Multiplication factor used for touched postings token") + f.Float64Var(&cfg.FetchedSeriesTokenFactor, "blocks-storage.bucket-store.fetched-series-token-factor", 2.5, "Multiplication factor used for fetched series token") + f.Float64Var(&cfg.TouchedSeriesTokenFactor, "blocks-storage.bucket-store.touched-series-token-factor", 10, "Multiplication factor used for touched series token") + f.Float64Var(&cfg.FetchedChunksTokenFactor, "blocks-storage.bucket-store.fetched-chunks-token-factor", 0, "Multiplication factor used for fetched chunks token") + f.Float64Var(&cfg.TouchedChunksTokenFactor, "blocks-storage.bucket-store.touched-chunks-token-factor", 0, "Multiplication factor used for touched chunks token") } // Validate the config. diff --git a/pkg/storage/tsdb/inmemory_index_cache.go b/pkg/storage/tsdb/inmemory_index_cache.go index 1530bf99f8..981e1bf6c4 100644 --- a/pkg/storage/tsdb/inmemory_index_cache.go +++ b/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, diff --git a/pkg/util/token_bucket.go b/pkg/util/token_bucket.go new file mode 100644 index 0000000000..b5d2d1ede4 --- /dev/null +++ b/pkg/util/token_bucket.go @@ -0,0 +1,86 @@ +package util + +import ( + "sync" + "time" + + "github.com/prometheus/client_golang/prometheus" +) + +type TokenBucket struct { + remainingTokens int64 + maxCapacity int64 + refillRate int64 + lastRefill time.Time + mu sync.Mutex + + remainingTokensMetric prometheus.Gauge +} + +func NewTokenBucket(maxCapacity, refillRate int64, remainingTokensMetric prometheus.Gauge) *TokenBucket { + if remainingTokensMetric != nil { + remainingTokensMetric.Set(float64(maxCapacity)) + } + + return &TokenBucket{ + remainingTokens: maxCapacity, + maxCapacity: maxCapacity, + refillRate: refillRate, + lastRefill: time.Now(), + remainingTokensMetric: remainingTokensMetric, + } +} + +func (t *TokenBucket) Retrieve(amount int64) bool { + t.mu.Lock() + defer t.mu.Unlock() + + t.updateTokens() + + if t.remainingTokens < amount { + return false + } + + t.remainingTokens -= amount + if t.remainingTokensMetric != nil { + t.remainingTokensMetric.Set(float64(t.remainingTokens)) + } + + return true +} + +func (t *TokenBucket) ForceRetrieve(amount int64) { + t.mu.Lock() + defer t.mu.Unlock() + + t.updateTokens() + + t.remainingTokens -= amount + if t.remainingTokensMetric != nil { + t.remainingTokensMetric.Set(float64(t.remainingTokens)) + } +} + +func (t *TokenBucket) Refund(amount int64) { + t.mu.Lock() + defer t.mu.Unlock() + + t.updateTokens() + + t.remainingTokens += amount + + if t.remainingTokens > t.maxCapacity { + t.remainingTokens = t.maxCapacity + } +} + +func (t *TokenBucket) updateTokens() { + now := time.Now() + refilledTokens := int64(now.Sub(t.lastRefill).Seconds() * float64(t.refillRate)) + t.remainingTokens += refilledTokens + t.lastRefill = now + + if t.remainingTokens > t.maxCapacity { + t.remainingTokens = t.maxCapacity + } +} From 917d2e9cae518943b1e40892f6fc2d5060868a05 Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Thu, 13 Jun 2024 09:24:48 -0700 Subject: [PATCH 02/26] Update bucket stores to pass token bucket Signed-off-by: Justin Jung --- pkg/storegateway/bucket_stores.go | 114 ++++++++++++++++++++++++++++-- 1 file changed, 110 insertions(+), 4 deletions(-) diff --git a/pkg/storegateway/bucket_stores.go b/pkg/storegateway/bucket_stores.go index 5d2c2d0ec1..63bf3bd98f 100644 --- a/pkg/storegateway/bucket_stores.go +++ b/pkg/storegateway/bucket_stores.go @@ -35,6 +35,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 +74,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 +121,11 @@ func NewBucketStores(cfg tsdb.BlocksStorageConfig, shardingStrategy ShardingStra metaFetcherMetrics: NewMetadataFetcherMetrics(), queryGate: queryGate, partitioner: newGapBasedPartitioner(cfg.BucketStore.PartitionerMaxGapBytes, reg), + podTokenBucket: util.NewTokenBucket(cfg.BucketStore.PodTokenBucketSize, cfg.BucketStore.PodTokenBucketSize, promauto.With(reg).NewGauge(prometheus.GaugeOpts{ + Name: "cortex_bucket_stores_pod_token_bucket_remaining", + Help: "Number of tokens left in pod token bucket.", + })), + 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", @@ -475,6 +486,10 @@ func (u *BucketStores) closeEmptyBucketStore(userID string) error { unlockInDefer = false u.storesMu.Unlock() + u.userTokenBucketsMu.Lock() + delete(u.userTokenBuckets, userID) + u.userTokenBucketsMu.Unlock() + u.metaFetcherMetrics.RemoveUserRegistry(userID) u.bucketStoreMetrics.RemoveUserRegistry(userID) return bs.Close() @@ -618,7 +633,7 @@ func (u *BucketStores) getOrCreateStore(userID string) (*store.BucketStore, erro 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.getTokensToRetrieve, u.cfg.BucketStore.RequestTokenBucketSize), u.partitioner, u.cfg.BucketStore.BlockSyncConcurrency, false, // No need to enable backward compatibility with Thanos pre 0.8.0 queriers @@ -632,6 +647,10 @@ func (u *BucketStores) getOrCreateStore(userID string) (*store.BucketStore, erro return nil, err } + u.userTokenBucketsMu.Lock() + u.userTokenBuckets[userID] = util.NewTokenBucket(u.cfg.BucketStore.UserTokenBucketSize, u.cfg.BucketStore.UserTokenBucketSize, nil) + u.userTokenBucketsMu.Unlock() + u.stores[userID] = bs u.metaFetcherMetrics.AddUserRegistry(userID, fetcherReg) u.bucketStoreMetrics.AddUserRegistry(userID, bucketStoreReg) @@ -680,6 +699,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 := tokens + switch dataType { + case store.PostingsFetched: + tokensToRetrieve *= uint64(u.cfg.BucketStore.FetchedPostingsTokenFactor) + case store.PostingsTouched: + tokensToRetrieve *= uint64(u.cfg.BucketStore.TouchedPostingsTokenFactor) + case store.SeriesFetched: + tokensToRetrieve *= uint64(u.cfg.BucketStore.FetchedSeriesTokenFactor) + case store.SeriesTouched: + tokensToRetrieve *= uint64(u.cfg.BucketStore.TouchedSeriesTokenFactor) + case store.ChunksFetched: + tokensToRetrieve *= uint64(u.cfg.BucketStore.FetchedChunksTokenFactor) + case store.ChunksTouched: + tokensToRetrieve *= uint64(u.cfg.BucketStore.TouchedChunksTokenFactor) + } + return int64(tokensToRetrieve) +} + func getUserIDFromGRPCContext(ctx context.Context) string { meta, ok := metadata.FromIncomingContext(ctx) if !ok { @@ -748,6 +792,65 @@ func (c *limiter) ReserveWithType(num uint64, _ store.StoreDataType) error { return nil } +type compositeLimiter struct { + limiters []store.BytesLimiter +} + +func (c *compositeLimiter) ReserveWithType(num uint64, dataType store.StoreDataType) error { + for _, l := range c.limiters { + if err := l.ReserveWithType(num, dataType); err != nil { + return err + } + } + return nil +} + +type tokenBucketLimiter struct { + podTokenBucket *util.TokenBucket + userTokenBucket *util.TokenBucket + requestTokenBucket *util.TokenBucket + getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64 +} + +func (t *tokenBucketLimiter) Reserve(_ uint64) error { + return nil +} + +func (t *tokenBucketLimiter) ReserveWithType(num uint64, dataType store.StoreDataType) error { + tokensToRetrieve := t.getTokensToRetrieve(num, dataType) + + // check provisioned bucket + retrieved := t.requestTokenBucket.Retrieve(tokensToRetrieve) + if retrieved { + t.userTokenBucket.ForceRetrieve(tokensToRetrieve) + t.podTokenBucket.ForceRetrieve(tokensToRetrieve) + return nil + } + + // if provisioned bucket is not enough, check burst buckets + retrieved = t.userTokenBucket.Retrieve(tokensToRetrieve) + if !retrieved { + return fmt.Errorf("not enough tokens in user bucket") + } + + retrieved = t.podTokenBucket.Retrieve(tokensToRetrieve) + if !retrieved { + t.userTokenBucket.Refund(tokensToRetrieve) + return fmt.Errorf("not enough tokens in pod bucket") + } + + return nil +} + +func newTokenBucketLimiter(podTokenBucket, userTokenBucket, requestTokenBucket *util.TokenBucket, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64) *tokenBucketLimiter { + return &tokenBucketLimiter{ + podTokenBucket: podTokenBucket, + userTokenBucket: userTokenBucket, + requestTokenBucket: requestTokenBucket, + getTokensToRetrieve: getTokensToRetrieve, + } +} + 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 @@ -768,12 +871,15 @@ func newSeriesLimiterFactory(limits *validation.Overrides, userID string) store. } } -func newBytesLimiterFactory(limits *validation.Overrides, userID string) store.BytesLimiterFactory { +func newBytesLimiterFactory(limits *validation.Overrides, userID string, podTokenBucket, userTokenBucket *util.TokenBucket, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64, requestTokenBucketSize int64) store.BytesLimiterFactory { return func(failedCounter prometheus.Counter) store.BytesLimiter { + limiters := []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), + limiters = append(limiters, store.NewLimiter(uint64(limits.MaxDownloadedBytesPerRequest(userID)), failedCounter)) + limiters = append(limiters, newTokenBucketLimiter(podTokenBucket, userTokenBucket, util.NewTokenBucket(requestTokenBucketSize, requestTokenBucketSize, nil), getTokensToRetrieve)) + return &compositeLimiter{ + limiters: limiters, } } } From 6700b08a26934dfe5f9e6b244b0b4fcbec9d836d Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Thu, 13 Jun 2024 13:09:35 -0700 Subject: [PATCH 03/26] Move limiters to a new file Signed-off-by: Justin Jung --- pkg/storegateway/bucket_stores.go | 110 --------------------------- pkg/storegateway/limiter.go | 121 ++++++++++++++++++++++++++++++ 2 files changed, 121 insertions(+), 110 deletions(-) create mode 100644 pkg/storegateway/limiter.go diff --git a/pkg/storegateway/bucket_stores.go b/pkg/storegateway/bucket_stores.go index 63bf3bd98f..18d0536494 100644 --- a/pkg/storegateway/bucket_stores.go +++ b/pkg/storegateway/bucket_stores.go @@ -4,7 +4,6 @@ import ( "context" "fmt" "math" - "net/http" "os" "path/filepath" "strings" @@ -774,112 +773,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 -} - -type compositeLimiter struct { - limiters []store.BytesLimiter -} - -func (c *compositeLimiter) ReserveWithType(num uint64, dataType store.StoreDataType) error { - for _, l := range c.limiters { - if err := l.ReserveWithType(num, dataType); err != nil { - return err - } - } - return nil -} - -type tokenBucketLimiter struct { - podTokenBucket *util.TokenBucket - userTokenBucket *util.TokenBucket - requestTokenBucket *util.TokenBucket - getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64 -} - -func (t *tokenBucketLimiter) Reserve(_ uint64) error { - return nil -} - -func (t *tokenBucketLimiter) ReserveWithType(num uint64, dataType store.StoreDataType) error { - tokensToRetrieve := t.getTokensToRetrieve(num, dataType) - - // check provisioned bucket - retrieved := t.requestTokenBucket.Retrieve(tokensToRetrieve) - if retrieved { - t.userTokenBucket.ForceRetrieve(tokensToRetrieve) - t.podTokenBucket.ForceRetrieve(tokensToRetrieve) - return nil - } - - // if provisioned bucket is not enough, check burst buckets - retrieved = t.userTokenBucket.Retrieve(tokensToRetrieve) - if !retrieved { - return fmt.Errorf("not enough tokens in user bucket") - } - - retrieved = t.podTokenBucket.Retrieve(tokensToRetrieve) - if !retrieved { - t.userTokenBucket.Refund(tokensToRetrieve) - return fmt.Errorf("not enough tokens in pod bucket") - } - - return nil -} - -func newTokenBucketLimiter(podTokenBucket, userTokenBucket, requestTokenBucket *util.TokenBucket, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64) *tokenBucketLimiter { - return &tokenBucketLimiter{ - podTokenBucket: podTokenBucket, - userTokenBucket: userTokenBucket, - requestTokenBucket: requestTokenBucket, - getTokensToRetrieve: getTokensToRetrieve, - } -} - -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, podTokenBucket, userTokenBucket *util.TokenBucket, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64, requestTokenBucketSize int64) store.BytesLimiterFactory { - return func(failedCounter prometheus.Counter) store.BytesLimiter { - limiters := []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. - limiters = append(limiters, store.NewLimiter(uint64(limits.MaxDownloadedBytesPerRequest(userID)), failedCounter)) - limiters = append(limiters, newTokenBucketLimiter(podTokenBucket, userTokenBucket, util.NewTokenBucket(requestTokenBucketSize, requestTokenBucketSize, nil), getTokensToRetrieve)) - return &compositeLimiter{ - limiters: limiters, - } - } -} diff --git a/pkg/storegateway/limiter.go b/pkg/storegateway/limiter.go new file mode 100644 index 0000000000..a1bf24e76b --- /dev/null +++ b/pkg/storegateway/limiter.go @@ -0,0 +1,121 @@ +package storegateway + +import ( + "fmt" + "net/http" + + "github.com/cortexproject/cortex/pkg/util" + "github.com/cortexproject/cortex/pkg/util/validation" + "github.com/prometheus/client_golang/prometheus" + "github.com/thanos-io/thanos/pkg/store" + "github.com/weaveworks/common/httpgrpc" +) + +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 +} + +type compositeLimiter struct { + limiters []store.BytesLimiter +} + +func (c *compositeLimiter) ReserveWithType(num uint64, dataType store.StoreDataType) error { + for _, l := range c.limiters { + if err := l.ReserveWithType(num, dataType); err != nil { + return err + } + } + return nil +} + +type tokenBucketLimiter struct { + podTokenBucket *util.TokenBucket + userTokenBucket *util.TokenBucket + requestTokenBucket *util.TokenBucket + getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64 +} + +func (t *tokenBucketLimiter) Reserve(_ uint64) error { + return nil +} + +func (t *tokenBucketLimiter) ReserveWithType(num uint64, dataType store.StoreDataType) error { + tokensToRetrieve := t.getTokensToRetrieve(num, dataType) + + // check provisioned bucket + retrieved := t.requestTokenBucket.Retrieve(tokensToRetrieve) + if retrieved { + t.userTokenBucket.ForceRetrieve(tokensToRetrieve) + t.podTokenBucket.ForceRetrieve(tokensToRetrieve) + return nil + } + + // if provisioned bucket is not enough, check burst buckets + retrieved = t.userTokenBucket.Retrieve(tokensToRetrieve) + if !retrieved { + return fmt.Errorf("not enough tokens in user bucket") + } + + retrieved = t.podTokenBucket.Retrieve(tokensToRetrieve) + if !retrieved { + t.userTokenBucket.Refund(tokensToRetrieve) + return fmt.Errorf("not enough tokens in pod bucket") + } + + return nil +} + +func newTokenBucketLimiter(podTokenBucket, userTokenBucket, requestTokenBucket *util.TokenBucket, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64) *tokenBucketLimiter { + return &tokenBucketLimiter{ + podTokenBucket: podTokenBucket, + userTokenBucket: userTokenBucket, + requestTokenBucket: requestTokenBucket, + getTokensToRetrieve: getTokensToRetrieve, + } +} + +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, podTokenBucket, userTokenBucket *util.TokenBucket, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64, requestTokenBucketSize int64) store.BytesLimiterFactory { + return func(failedCounter prometheus.Counter) store.BytesLimiter { + limiters := []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. + limiters = append(limiters, store.NewLimiter(uint64(limits.MaxDownloadedBytesPerRequest(userID)), failedCounter)) + limiters = append(limiters, newTokenBucketLimiter(podTokenBucket, userTokenBucket, util.NewTokenBucket(requestTokenBucketSize, requestTokenBucketSize, nil), getTokensToRetrieve)) + return &compositeLimiter{ + limiters: limiters, + } + } +} From 5e78511a8a8491a381d0329f8bab865adb09f7d2 Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Thu, 13 Jun 2024 15:28:28 -0700 Subject: [PATCH 04/26] Added tests for limiters and token bucket Signed-off-by: Justin Jung --- pkg/storegateway/limiter.go | 8 ++-- pkg/storegateway/limiter_test.go | 73 ++++++++++++++++++++++++++++++++ pkg/util/token_bucket_test.go | 34 +++++++++++++++ 3 files changed, 111 insertions(+), 4 deletions(-) create mode 100644 pkg/storegateway/limiter_test.go create mode 100644 pkg/util/token_bucket_test.go diff --git a/pkg/storegateway/limiter.go b/pkg/storegateway/limiter.go index a1bf24e76b..a5cba607ee 100644 --- a/pkg/storegateway/limiter.go +++ b/pkg/storegateway/limiter.go @@ -55,7 +55,7 @@ func (t *tokenBucketLimiter) Reserve(_ uint64) error { func (t *tokenBucketLimiter) ReserveWithType(num uint64, dataType store.StoreDataType) error { tokensToRetrieve := t.getTokensToRetrieve(num, dataType) - // check provisioned bucket + // check request bucket retrieved := t.requestTokenBucket.Retrieve(tokensToRetrieve) if retrieved { t.userTokenBucket.ForceRetrieve(tokensToRetrieve) @@ -63,16 +63,16 @@ func (t *tokenBucketLimiter) ReserveWithType(num uint64, dataType store.StoreDat return nil } - // if provisioned bucket is not enough, check burst buckets + // if request bucket is running low, check shared buckets retrieved = t.userTokenBucket.Retrieve(tokensToRetrieve) if !retrieved { - return fmt.Errorf("not enough tokens in user bucket") + return fmt.Errorf("not enough tokens in user token bucket") } retrieved = t.podTokenBucket.Retrieve(tokensToRetrieve) if !retrieved { t.userTokenBucket.Refund(tokensToRetrieve) - return fmt.Errorf("not enough tokens in pod bucket") + return fmt.Errorf("not enough tokens in pod token bucket") } return nil diff --git a/pkg/storegateway/limiter_test.go b/pkg/storegateway/limiter_test.go new file mode 100644 index 0000000000..3898595fb0 --- /dev/null +++ b/pkg/storegateway/limiter_test.go @@ -0,0 +1,73 @@ +package storegateway + +import ( + "testing" + + "github.com/cortexproject/cortex/pkg/util" + "github.com/prometheus/client_golang/prometheus" + "github.com/stretchr/testify/assert" + "github.com/thanos-io/thanos/pkg/store" +) + +func TestLimiter(t *testing.T) { + l := &limiter{ + limiter: store.NewLimiter(2, prometheus.NewCounter(prometheus.CounterOpts{})), + } + + assert.NoError(t, l.Reserve(1)) + assert.NoError(t, l.ReserveWithType(1, store.PostingsFetched)) + assert.Error(t, l.Reserve(1)) + assert.Error(t, l.ReserveWithType(1, store.PostingsFetched)) +} + +func TestCompositeLimiter(t *testing.T) { + l := &compositeLimiter{ + limiters: []store.BytesLimiter{ + store.NewLimiter(2, prometheus.NewCounter(prometheus.CounterOpts{})), + store.NewLimiter(1, prometheus.NewCounter(prometheus.CounterOpts{})), + }, + } + + assert.NoError(t, l.ReserveWithType(1, store.PostingsFetched)) + assert.Error(t, l.ReserveWithType(1, store.PostingsFetched)) +} + +func TestNewTokenBucketLimiter(t *testing.T) { + podTokenBucket := util.NewTokenBucket(3, 3, nil) + userTokenBucket := util.NewTokenBucket(2, 2, nil) + requestTokenBucket := util.NewTokenBucket(1, 1, nil) + l := newTokenBucketLimiter(podTokenBucket, userTokenBucket, requestTokenBucket, func(tokens uint64, dataType store.StoreDataType) int64 { + if dataType == store.SeriesFetched { + return int64(tokens) * 5 + } + return int64(tokens) + }) + + // should force retrieve tokens from all buckets upon succeeding + assert.NoError(t, l.ReserveWithType(2, store.PostingsFetched)) + assert.False(t, podTokenBucket.Retrieve(2)) + + // should fail if user token bucket is running low + podTokenBucket.Refund(3) + userTokenBucket.Refund(2) + requestTokenBucket.Refund(1) + assert.ErrorContains(t, l.ReserveWithType(3, store.PostingsFetched), "not enough tokens in user token bucket") + + // should fail if pod token bucket is running low + podTokenBucket.Refund(3) + userTokenBucket.Refund(2) + requestTokenBucket.Refund(1) + podTokenBucket.ForceRetrieve(3) + assert.ErrorContains(t, l.ReserveWithType(2, store.PostingsFetched), "not enough tokens in pod token bucket") + + // should retrieve different amount of tokens based on data type + podTokenBucket.Refund(3) + userTokenBucket.Refund(2) + requestTokenBucket.Refund(1) + assert.ErrorContains(t, l.ReserveWithType(1, store.SeriesFetched), "not enough tokens in user token bucket") + + // should always succeed if retrieve token bucket has enough tokens, although shared buckets are empty + podTokenBucket.ForceRetrieve(3) + userTokenBucket.ForceRetrieve(2) + assert.NoError(t, l.ReserveWithType(1, store.PostingsFetched)) +} diff --git a/pkg/util/token_bucket_test.go b/pkg/util/token_bucket_test.go new file mode 100644 index 0000000000..b74fdc2065 --- /dev/null +++ b/pkg/util/token_bucket_test.go @@ -0,0 +1,34 @@ +package util + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +func TestTokenBucket_Retrieve(t *testing.T) { + bucket := NewTokenBucket(10, 600, nil) + + assert.True(t, bucket.Retrieve(10)) + assert.False(t, bucket.Retrieve(10)) + time.Sleep(time.Second) + assert.True(t, bucket.Retrieve(10)) +} + +func TestTokenBucket_ForceRetrieve(t *testing.T) { + bucket := NewTokenBucket(10, 600, nil) + + bucket.ForceRetrieve(20) + assert.False(t, bucket.Retrieve(10)) + time.Sleep(time.Second) + assert.True(t, bucket.Retrieve(10)) +} + +func TestTokenBucket_Refund(t *testing.T) { + bucket := NewTokenBucket(10, 600, nil) + + bucket.ForceRetrieve(100) + bucket.Refund(100) + assert.True(t, bucket.Retrieve(10)) +} From 8ff710e82abc18f4bfa0f5628b7013eebd4a05d5 Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Thu, 13 Jun 2024 17:18:17 -0700 Subject: [PATCH 05/26] Add more tests Signed-off-by: Justin Jung --- pkg/storage/tsdb/config.go | 4 ++-- pkg/storegateway/bucket_stores.go | 22 +++++++++--------- pkg/storegateway/bucket_stores_test.go | 31 ++++++++++++++++++++++++++ 3 files changed, 44 insertions(+), 13 deletions(-) diff --git a/pkg/storage/tsdb/config.go b/pkg/storage/tsdb/config.go index 10c50bb8fc..53f142e3bb 100644 --- a/pkg/storage/tsdb/config.go +++ b/pkg/storage/tsdb/config.go @@ -345,8 +345,8 @@ func (cfg *BucketStoreConfig) RegisterFlags(f *flag.FlagSet) { f.Float64Var(&cfg.TouchedPostingsTokenFactor, "blocks-storage.bucket-store.touched-postings-token-factor", 2, "Multiplication factor used for touched postings token") f.Float64Var(&cfg.FetchedSeriesTokenFactor, "blocks-storage.bucket-store.fetched-series-token-factor", 2.5, "Multiplication factor used for fetched series token") f.Float64Var(&cfg.TouchedSeriesTokenFactor, "blocks-storage.bucket-store.touched-series-token-factor", 10, "Multiplication factor used for touched series token") - f.Float64Var(&cfg.FetchedChunksTokenFactor, "blocks-storage.bucket-store.fetched-chunks-token-factor", 0, "Multiplication factor used for fetched chunks token") - f.Float64Var(&cfg.TouchedChunksTokenFactor, "blocks-storage.bucket-store.touched-chunks-token-factor", 0, "Multiplication factor used for touched chunks token") + f.Float64Var(&cfg.FetchedChunksTokenFactor, "blocks-storage.bucket-store.fetched-chunks-token-factor", 0.5, "Multiplication factor used for fetched chunks token") + f.Float64Var(&cfg.TouchedChunksTokenFactor, "blocks-storage.bucket-store.touched-chunks-token-factor", 0.5, "Multiplication factor used for touched chunks token") } // Validate the config. diff --git a/pkg/storegateway/bucket_stores.go b/pkg/storegateway/bucket_stores.go index 18d0536494..283382b6b3 100644 --- a/pkg/storegateway/bucket_stores.go +++ b/pkg/storegateway/bucket_stores.go @@ -626,6 +626,10 @@ func (u *BucketStores) getOrCreateStore(userID string) (*store.BucketStore, erro bucketStoreOpts = append(bucketStoreOpts, store.WithDebugLogging()) } + u.userTokenBucketsMu.Lock() + u.userTokenBuckets[userID] = util.NewTokenBucket(u.cfg.BucketStore.UserTokenBucketSize, u.cfg.BucketStore.UserTokenBucketSize, nil) + u.userTokenBucketsMu.Unlock() + bs, err := store.NewBucketStore( userBkt, fetcher, @@ -646,10 +650,6 @@ func (u *BucketStores) getOrCreateStore(userID string) (*store.BucketStore, erro return nil, err } - u.userTokenBucketsMu.Lock() - u.userTokenBuckets[userID] = util.NewTokenBucket(u.cfg.BucketStore.UserTokenBucketSize, u.cfg.BucketStore.UserTokenBucketSize, nil) - u.userTokenBucketsMu.Unlock() - u.stores[userID] = bs u.metaFetcherMetrics.AddUserRegistry(userID, fetcherReg) u.bucketStoreMetrics.AddUserRegistry(userID, bucketStoreReg) @@ -705,20 +705,20 @@ func (u *BucketStores) getUserTokenBucket(userID string) *util.TokenBucket { } func (u *BucketStores) getTokensToRetrieve(tokens uint64, dataType store.StoreDataType) int64 { - tokensToRetrieve := tokens + tokensToRetrieve := float64(tokens) switch dataType { case store.PostingsFetched: - tokensToRetrieve *= uint64(u.cfg.BucketStore.FetchedPostingsTokenFactor) + tokensToRetrieve *= u.cfg.BucketStore.FetchedPostingsTokenFactor case store.PostingsTouched: - tokensToRetrieve *= uint64(u.cfg.BucketStore.TouchedPostingsTokenFactor) + tokensToRetrieve *= u.cfg.BucketStore.TouchedPostingsTokenFactor case store.SeriesFetched: - tokensToRetrieve *= uint64(u.cfg.BucketStore.FetchedSeriesTokenFactor) + tokensToRetrieve *= u.cfg.BucketStore.FetchedSeriesTokenFactor case store.SeriesTouched: - tokensToRetrieve *= uint64(u.cfg.BucketStore.TouchedSeriesTokenFactor) + tokensToRetrieve *= u.cfg.BucketStore.TouchedSeriesTokenFactor case store.ChunksFetched: - tokensToRetrieve *= uint64(u.cfg.BucketStore.FetchedChunksTokenFactor) + tokensToRetrieve *= u.cfg.BucketStore.FetchedChunksTokenFactor case store.ChunksTouched: - tokensToRetrieve *= uint64(u.cfg.BucketStore.TouchedChunksTokenFactor) + tokensToRetrieve *= u.cfg.BucketStore.TouchedChunksTokenFactor } return int64(tokensToRetrieve) } diff --git a/pkg/storegateway/bucket_stores_test.go b/pkg/storegateway/bucket_stores_test.go index a10a4599fb..ea9cff7211 100644 --- a/pkg/storegateway/bucket_stores_test.go +++ b/pkg/storegateway/bucket_stores_test.go @@ -212,6 +212,8 @@ func TestBucketStores_InitialSync(t *testing.T) { } require.NoError(t, stores.InitialSync(ctx)) + assert.NotNil(t, stores.getUserTokenBucket("user-1")) + assert.NotNil(t, stores.getUserTokenBucket("user-2")) // Query series after the initial sync. for userID, metricName := range userToMetric { @@ -718,6 +720,8 @@ func TestBucketStores_deleteLocalFilesForExcludedTenants(t *testing.T) { sharding.users = []string{user1} require.NoError(t, stores.SyncBlocks(ctx)) require.Equal(t, []string{user1}, getUsersInDir(t, cfg.BucketStore.SyncDir)) + assert.NotNil(t, stores.getUserTokenBucket(user1)) + assert.Nil(t, stores.getUserTokenBucket(user2)) require.NoError(t, testutil.GatherAndCompare(reg, strings.NewReader(` # HELP cortex_bucket_store_block_drops_total Total number of local blocks that were dropped. @@ -735,6 +739,8 @@ func TestBucketStores_deleteLocalFilesForExcludedTenants(t *testing.T) { sharding.users = nil require.NoError(t, stores.SyncBlocks(ctx)) require.Equal(t, []string(nil), getUsersInDir(t, cfg.BucketStore.SyncDir)) + assert.Nil(t, stores.getUserTokenBucket(user1)) + assert.Nil(t, stores.getUserTokenBucket(user2)) require.NoError(t, testutil.GatherAndCompare(reg, strings.NewReader(` # HELP cortex_bucket_store_block_drops_total Total number of local blocks that were dropped. @@ -763,6 +769,31 @@ func TestBucketStores_deleteLocalFilesForExcludedTenants(t *testing.T) { `), metricNames...)) } +func TestBucketStores_getTokensToRetrieve(t *testing.T) { + cfg := prepareStorageConfig(t) + cfg.BucketStore.FetchedPostingsTokenFactor = 1 + cfg.BucketStore.TouchedPostingsTokenFactor = 2 + cfg.BucketStore.FetchedSeriesTokenFactor = 3 + cfg.BucketStore.TouchedSeriesTokenFactor = 4 + cfg.BucketStore.FetchedChunksTokenFactor = 0 + cfg.BucketStore.TouchedChunksTokenFactor = 0.5 + + storageDir := t.TempDir() + bucket, err := filesystem.NewBucketClient(filesystem.Config{Directory: storageDir}) + require.NoError(t, err) + + reg := prometheus.NewPedanticRegistry() + stores, err := NewBucketStores(cfg, NewNoShardingStrategy(log.NewNopLogger(), nil), objstore.WithNoopInstr(bucket), defaultLimitsOverrides(t), mockLoggingLevel(), log.NewNopLogger(), reg) + require.NoError(t, err) + + assert.Equal(t, int64(2), stores.getTokensToRetrieve(2, store.PostingsFetched)) + assert.Equal(t, int64(4), stores.getTokensToRetrieve(2, store.PostingsTouched)) + assert.Equal(t, int64(6), stores.getTokensToRetrieve(2, store.SeriesFetched)) + assert.Equal(t, int64(8), stores.getTokensToRetrieve(2, store.SeriesTouched)) + assert.Equal(t, int64(0), stores.getTokensToRetrieve(2, store.ChunksFetched)) + assert.Equal(t, int64(1), stores.getTokensToRetrieve(2, store.ChunksTouched)) +} + func getUsersInDir(t *testing.T, dir string) []string { fs, err := os.ReadDir(dir) require.NoError(t, err) From 6c15c92f3d65239d3245619662ff50d07fea838a Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Thu, 13 Jun 2024 22:08:41 -0700 Subject: [PATCH 06/26] Added enable flag Signed-off-by: Justin Jung --- pkg/storage/tsdb/config.go | 28 ++++++---- pkg/storegateway/bucket_stores.go | 32 ++++++----- pkg/storegateway/bucket_stores_test.go | 76 +++++++++++++++++++++----- pkg/storegateway/limiter.go | 14 ++++- pkg/storegateway/limiter_test.go | 3 +- 5 files changed, 110 insertions(+), 43 deletions(-) diff --git a/pkg/storage/tsdb/config.go b/pkg/storage/tsdb/config.go index 53f142e3bb..4b6fcb7e5e 100644 --- a/pkg/storage/tsdb/config.go +++ b/pkg/storage/tsdb/config.go @@ -294,7 +294,12 @@ type BucketStoreConfig struct { SeriesBatchSize int `yaml:"series_batch_size"` // Token bucket configs - PodDataBytesRateLimit int64 `yaml:"pod_data_bytes_rate_limit"` + 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"` @@ -337,16 +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.Int64Var(&cfg.PodDataBytesRateLimit, "blocks-storage.bucket-store.pod-data-bytes-rate-limit", int64(1*units.Gibibyte), "Overall data bytes rate limit for a pod") - f.Int64Var(&cfg.PodTokenBucketSize, "blocks-storage.bucket-store.pod-token-bucket-size", int64(820*units.Mebibyte), "Pod token bucket size") - f.Int64Var(&cfg.UserTokenBucketSize, "blocks-storage.bucket-store.user-token-bucket-size", int64(615*units.Mebibyte), "User token bucket size") - f.Int64Var(&cfg.RequestTokenBucketSize, "blocks-storage.bucket-store.request-token-bucket-size", int64(4*units.Mebibyte), "Request token bucket size") - f.Float64Var(&cfg.FetchedPostingsTokenFactor, "blocks-storage.bucket-store.fetched-postings-token-factor", 2, "Multiplication factor used for fetched postings token") - f.Float64Var(&cfg.TouchedPostingsTokenFactor, "blocks-storage.bucket-store.touched-postings-token-factor", 2, "Multiplication factor used for touched postings token") - f.Float64Var(&cfg.FetchedSeriesTokenFactor, "blocks-storage.bucket-store.fetched-series-token-factor", 2.5, "Multiplication factor used for fetched series token") - f.Float64Var(&cfg.TouchedSeriesTokenFactor, "blocks-storage.bucket-store.touched-series-token-factor", 10, "Multiplication factor used for touched series token") - f.Float64Var(&cfg.FetchedChunksTokenFactor, "blocks-storage.bucket-store.fetched-chunks-token-factor", 0.5, "Multiplication factor used for fetched chunks token") - f.Float64Var(&cfg.TouchedChunksTokenFactor, "blocks-storage.bucket-store.touched-chunks-token-factor", 0.5, "Multiplication factor used for touched chunks token") + 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. diff --git a/pkg/storegateway/bucket_stores.go b/pkg/storegateway/bucket_stores.go index 283382b6b3..9ec08251f1 100644 --- a/pkg/storegateway/bucket_stores.go +++ b/pkg/storegateway/bucket_stores.go @@ -120,7 +120,7 @@ func NewBucketStores(cfg tsdb.BlocksStorageConfig, shardingStrategy ShardingStra metaFetcherMetrics: NewMetadataFetcherMetrics(), queryGate: queryGate, partitioner: newGapBasedPartitioner(cfg.BucketStore.PartitionerMaxGapBytes, reg), - podTokenBucket: util.NewTokenBucket(cfg.BucketStore.PodTokenBucketSize, cfg.BucketStore.PodTokenBucketSize, promauto.With(reg).NewGauge(prometheus.GaugeOpts{ + 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.", })), @@ -485,9 +485,11 @@ func (u *BucketStores) closeEmptyBucketStore(userID string) error { unlockInDefer = false u.storesMu.Unlock() - u.userTokenBucketsMu.Lock() - delete(u.userTokenBuckets, userID) - u.userTokenBucketsMu.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) @@ -626,9 +628,11 @@ func (u *BucketStores) getOrCreateStore(userID string) (*store.BucketStore, erro bucketStoreOpts = append(bucketStoreOpts, store.WithDebugLogging()) } - u.userTokenBucketsMu.Lock() - u.userTokenBuckets[userID] = util.NewTokenBucket(u.cfg.BucketStore.UserTokenBucketSize, u.cfg.BucketStore.UserTokenBucketSize, nil) - u.userTokenBucketsMu.Unlock() + 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, @@ -636,7 +640,7 @@ func (u *BucketStores) getOrCreateStore(userID string) (*store.BucketStore, erro u.syncDirForUser(userID), newChunksLimiterFactory(u.limits, userID), newSeriesLimiterFactory(u.limits, userID), - newBytesLimiterFactory(u.limits, userID, u.podTokenBucket, u.getUserTokenBucket(userID), u.getTokensToRetrieve, u.cfg.BucketStore.RequestTokenBucketSize), + 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 @@ -708,17 +712,17 @@ func (u *BucketStores) getTokensToRetrieve(tokens uint64, dataType store.StoreDa tokensToRetrieve := float64(tokens) switch dataType { case store.PostingsFetched: - tokensToRetrieve *= u.cfg.BucketStore.FetchedPostingsTokenFactor + tokensToRetrieve *= u.cfg.BucketStore.TokenBucketLimiter.FetchedPostingsTokenFactor case store.PostingsTouched: - tokensToRetrieve *= u.cfg.BucketStore.TouchedPostingsTokenFactor + tokensToRetrieve *= u.cfg.BucketStore.TokenBucketLimiter.TouchedPostingsTokenFactor case store.SeriesFetched: - tokensToRetrieve *= u.cfg.BucketStore.FetchedSeriesTokenFactor + tokensToRetrieve *= u.cfg.BucketStore.TokenBucketLimiter.FetchedSeriesTokenFactor case store.SeriesTouched: - tokensToRetrieve *= u.cfg.BucketStore.TouchedSeriesTokenFactor + tokensToRetrieve *= u.cfg.BucketStore.TokenBucketLimiter.TouchedSeriesTokenFactor case store.ChunksFetched: - tokensToRetrieve *= u.cfg.BucketStore.FetchedChunksTokenFactor + tokensToRetrieve *= u.cfg.BucketStore.TokenBucketLimiter.FetchedChunksTokenFactor case store.ChunksTouched: - tokensToRetrieve *= u.cfg.BucketStore.TouchedChunksTokenFactor + tokensToRetrieve *= u.cfg.BucketStore.TokenBucketLimiter.TouchedChunksTokenFactor } return int64(tokensToRetrieve) } diff --git a/pkg/storegateway/bucket_stores_test.go b/pkg/storegateway/bucket_stores_test.go index ea9cff7211..0a37b005c6 100644 --- a/pkg/storegateway/bucket_stores_test.go +++ b/pkg/storegateway/bucket_stores_test.go @@ -212,8 +212,6 @@ func TestBucketStores_InitialSync(t *testing.T) { } require.NoError(t, stores.InitialSync(ctx)) - assert.NotNil(t, stores.getUserTokenBucket("user-1")) - assert.NotNil(t, stores.getUserTokenBucket("user-2")) // Query series after the initial sync. for userID, metricName := range userToMetric { @@ -720,8 +718,6 @@ func TestBucketStores_deleteLocalFilesForExcludedTenants(t *testing.T) { sharding.users = []string{user1} require.NoError(t, stores.SyncBlocks(ctx)) require.Equal(t, []string{user1}, getUsersInDir(t, cfg.BucketStore.SyncDir)) - assert.NotNil(t, stores.getUserTokenBucket(user1)) - assert.Nil(t, stores.getUserTokenBucket(user2)) require.NoError(t, testutil.GatherAndCompare(reg, strings.NewReader(` # HELP cortex_bucket_store_block_drops_total Total number of local blocks that were dropped. @@ -739,8 +735,6 @@ func TestBucketStores_deleteLocalFilesForExcludedTenants(t *testing.T) { sharding.users = nil require.NoError(t, stores.SyncBlocks(ctx)) require.Equal(t, []string(nil), getUsersInDir(t, cfg.BucketStore.SyncDir)) - assert.Nil(t, stores.getUserTokenBucket(user1)) - assert.Nil(t, stores.getUserTokenBucket(user2)) require.NoError(t, testutil.GatherAndCompare(reg, strings.NewReader(` # HELP cortex_bucket_store_block_drops_total Total number of local blocks that were dropped. @@ -769,22 +763,76 @@ func TestBucketStores_deleteLocalFilesForExcludedTenants(t *testing.T) { `), metricNames...)) } +func TestBucketStores_getUserTokenBucket(t *testing.T) { + const ( + user1 = "user-1" + user2 = "user-2" + ) + + ctx := context.Background() + cfg := prepareStorageConfig(t) + cfg.BucketStore.TokenBucketLimiter.Enabled = true + + storageDir := t.TempDir() + userToMetric := map[string]string{ + user1: "series_1", + user2: "series_2", + } + for userID, metricName := range userToMetric { + generateStorageBlock(t, storageDir, userID, metricName, 10, 100, 15) + } + + sharding := userShardingStrategy{} + sharding.users = []string{user1, user2} + + bucket, err := filesystem.NewBucketClient(filesystem.Config{Directory: storageDir}) + assert.NoError(t, err) + + reg := prometheus.NewPedanticRegistry() + stores, err := NewBucketStores(cfg, &sharding, objstore.WithNoopInstr(bucket), defaultLimitsOverrides(t), mockLoggingLevel(), log.NewNopLogger(), reg) + assert.NoError(t, err) + + assert.NoError(t, stores.InitialSync(ctx)) + assert.NotNil(t, stores.getUserTokenBucket("user-1")) + assert.NotNil(t, stores.getUserTokenBucket("user-2")) + + sharding.users = []string{user1} + assert.NoError(t, stores.SyncBlocks(ctx)) + assert.NotNil(t, stores.getUserTokenBucket("user-1")) + assert.Nil(t, stores.getUserTokenBucket("user-2")) + + sharding.users = []string{} + assert.NoError(t, stores.SyncBlocks(ctx)) + assert.Nil(t, stores.getUserTokenBucket("user-1")) + assert.Nil(t, stores.getUserTokenBucket("user-2")) + + cfg.BucketStore.TokenBucketLimiter.Enabled = false + sharding.users = []string{user1, user2} + reg = prometheus.NewPedanticRegistry() + stores, err = NewBucketStores(cfg, &sharding, objstore.WithNoopInstr(bucket), defaultLimitsOverrides(t), mockLoggingLevel(), log.NewNopLogger(), reg) + assert.NoError(t, err) + + assert.NoError(t, stores.InitialSync(ctx)) + assert.Nil(t, stores.getUserTokenBucket("user-1")) + assert.Nil(t, stores.getUserTokenBucket("user-2")) +} + func TestBucketStores_getTokensToRetrieve(t *testing.T) { cfg := prepareStorageConfig(t) - cfg.BucketStore.FetchedPostingsTokenFactor = 1 - cfg.BucketStore.TouchedPostingsTokenFactor = 2 - cfg.BucketStore.FetchedSeriesTokenFactor = 3 - cfg.BucketStore.TouchedSeriesTokenFactor = 4 - cfg.BucketStore.FetchedChunksTokenFactor = 0 - cfg.BucketStore.TouchedChunksTokenFactor = 0.5 + cfg.BucketStore.TokenBucketLimiter.FetchedPostingsTokenFactor = 1 + cfg.BucketStore.TokenBucketLimiter.TouchedPostingsTokenFactor = 2 + cfg.BucketStore.TokenBucketLimiter.FetchedSeriesTokenFactor = 3 + cfg.BucketStore.TokenBucketLimiter.TouchedSeriesTokenFactor = 4 + cfg.BucketStore.TokenBucketLimiter.FetchedChunksTokenFactor = 0 + cfg.BucketStore.TokenBucketLimiter.TouchedChunksTokenFactor = 0.5 storageDir := t.TempDir() bucket, err := filesystem.NewBucketClient(filesystem.Config{Directory: storageDir}) - require.NoError(t, err) + assert.NoError(t, err) reg := prometheus.NewPedanticRegistry() stores, err := NewBucketStores(cfg, NewNoShardingStrategy(log.NewNopLogger(), nil), objstore.WithNoopInstr(bucket), defaultLimitsOverrides(t), mockLoggingLevel(), log.NewNopLogger(), reg) - require.NoError(t, err) + assert.NoError(t, err) assert.Equal(t, int64(2), stores.getTokensToRetrieve(2, store.PostingsFetched)) assert.Equal(t, int64(4), stores.getTokensToRetrieve(2, store.PostingsTouched)) diff --git a/pkg/storegateway/limiter.go b/pkg/storegateway/limiter.go index a5cba607ee..717a409307 100644 --- a/pkg/storegateway/limiter.go +++ b/pkg/storegateway/limiter.go @@ -4,6 +4,7 @@ import ( "fmt" "net/http" + "github.com/cortexproject/cortex/pkg/storage/tsdb" "github.com/cortexproject/cortex/pkg/util" "github.com/cortexproject/cortex/pkg/util/validation" "github.com/prometheus/client_golang/prometheus" @@ -46,6 +47,7 @@ type tokenBucketLimiter struct { userTokenBucket *util.TokenBucket requestTokenBucket *util.TokenBucket getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64 + dryRun bool } func (t *tokenBucketLimiter) Reserve(_ uint64) error { @@ -78,12 +80,13 @@ func (t *tokenBucketLimiter) ReserveWithType(num uint64, dataType store.StoreDat return nil } -func newTokenBucketLimiter(podTokenBucket, userTokenBucket, requestTokenBucket *util.TokenBucket, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64) *tokenBucketLimiter { +func newTokenBucketLimiter(podTokenBucket, userTokenBucket, requestTokenBucket *util.TokenBucket, dryRun bool, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64) *tokenBucketLimiter { return &tokenBucketLimiter{ podTokenBucket: podTokenBucket, userTokenBucket: userTokenBucket, requestTokenBucket: requestTokenBucket, getTokensToRetrieve: getTokensToRetrieve, + dryRun: dryRun, } } @@ -107,13 +110,18 @@ func newSeriesLimiterFactory(limits *validation.Overrides, userID string) store. } } -func newBytesLimiterFactory(limits *validation.Overrides, userID string, podTokenBucket, userTokenBucket *util.TokenBucket, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64, requestTokenBucketSize int64) store.BytesLimiterFactory { +func newBytesLimiterFactory(limits *validation.Overrides, userID string, podTokenBucket, userTokenBucket *util.TokenBucket, tokenBucketLimiterCfg tsdb.TokenBucketLimiterConfig, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64) store.BytesLimiterFactory { return func(failedCounter prometheus.Counter) store.BytesLimiter { limiters := []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. limiters = append(limiters, store.NewLimiter(uint64(limits.MaxDownloadedBytesPerRequest(userID)), failedCounter)) - limiters = append(limiters, newTokenBucketLimiter(podTokenBucket, userTokenBucket, util.NewTokenBucket(requestTokenBucketSize, requestTokenBucketSize, nil), getTokensToRetrieve)) + + if tokenBucketLimiterCfg.Enabled { + requestTokenBucket := util.NewTokenBucket(tokenBucketLimiterCfg.RequestTokenBucketSize, tokenBucketLimiterCfg.RequestTokenBucketSize, nil) + limiters = append(limiters, newTokenBucketLimiter(podTokenBucket, userTokenBucket, requestTokenBucket, tokenBucketLimiterCfg.DryRun, getTokensToRetrieve)) + } + return &compositeLimiter{ limiters: limiters, } diff --git a/pkg/storegateway/limiter_test.go b/pkg/storegateway/limiter_test.go index 3898595fb0..fcf6e84baa 100644 --- a/pkg/storegateway/limiter_test.go +++ b/pkg/storegateway/limiter_test.go @@ -36,7 +36,7 @@ func TestNewTokenBucketLimiter(t *testing.T) { podTokenBucket := util.NewTokenBucket(3, 3, nil) userTokenBucket := util.NewTokenBucket(2, 2, nil) requestTokenBucket := util.NewTokenBucket(1, 1, nil) - l := newTokenBucketLimiter(podTokenBucket, userTokenBucket, requestTokenBucket, func(tokens uint64, dataType store.StoreDataType) int64 { + l := newTokenBucketLimiter(podTokenBucket, userTokenBucket, requestTokenBucket, true, func(tokens uint64, dataType store.StoreDataType) int64 { if dataType == store.SeriesFetched { return int64(tokens) * 5 } @@ -71,3 +71,4 @@ func TestNewTokenBucketLimiter(t *testing.T) { userTokenBucket.ForceRetrieve(2) assert.NoError(t, l.ReserveWithType(1, store.PostingsFetched)) } + From 6ae585acea396be0d1e887f8db0373d2090c9385 Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Thu, 13 Jun 2024 22:52:43 -0700 Subject: [PATCH 07/26] Add dryrun feature Signed-off-by: Justin Jung --- pkg/storegateway/limiter.go | 22 +++++++++++- pkg/storegateway/limiter_test.go | 58 ++++++++++++++++++++++++++------ pkg/util/token_bucket.go | 9 +++++ pkg/util/token_bucket_test.go | 8 +++-- 4 files changed, 83 insertions(+), 14 deletions(-) diff --git a/pkg/storegateway/limiter.go b/pkg/storegateway/limiter.go index 717a409307..0f07e057aa 100644 --- a/pkg/storegateway/limiter.go +++ b/pkg/storegateway/limiter.go @@ -6,7 +6,9 @@ import ( "github.com/cortexproject/cortex/pkg/storage/tsdb" "github.com/cortexproject/cortex/pkg/util" + util_log "github.com/cortexproject/cortex/pkg/util/log" "github.com/cortexproject/cortex/pkg/util/validation" + "github.com/go-kit/log/level" "github.com/prometheus/client_golang/prometheus" "github.com/thanos-io/thanos/pkg/store" "github.com/weaveworks/common/httpgrpc" @@ -65,18 +67,36 @@ func (t *tokenBucketLimiter) ReserveWithType(num uint64, dataType store.StoreDat return nil } - // if request bucket is running low, check shared buckets + // if we can't retrieve from request bucket, check shared buckets retrieved = t.userTokenBucket.Retrieve(tokensToRetrieve) if !retrieved { + // if dry run, force retrieve all tokens and return nil + if t.dryRun { + t.requestTokenBucket.ForceRetrieve(tokensToRetrieve) + t.userTokenBucket.ForceRetrieve(tokensToRetrieve) + t.podTokenBucket.ForceRetrieve(tokensToRetrieve) + level.Warn(util_log.Logger).Log("msg", "not enough tokens in user token bucket", "dataType", dataType, "dataSize", num, "tokens", tokensToRetrieve) + return nil + } return fmt.Errorf("not enough tokens in user token bucket") } retrieved = t.podTokenBucket.Retrieve(tokensToRetrieve) if !retrieved { t.userTokenBucket.Refund(tokensToRetrieve) + + // if dry run, force retrieve all tokens and return nil + if t.dryRun { + // user bucket is already retrieved + t.requestTokenBucket.ForceRetrieve(tokensToRetrieve) + t.podTokenBucket.ForceRetrieve(tokensToRetrieve) + level.Warn(util_log.Logger).Log("msg", "not enough tokens in pod token bucket", "dataType", dataType, "dataSize", num, "tokens", tokensToRetrieve) + return nil + } return fmt.Errorf("not enough tokens in pod token bucket") } + t.requestTokenBucket.ForceRetrieve(tokensToRetrieve) return nil } diff --git a/pkg/storegateway/limiter_test.go b/pkg/storegateway/limiter_test.go index fcf6e84baa..82345117f5 100644 --- a/pkg/storegateway/limiter_test.go +++ b/pkg/storegateway/limiter_test.go @@ -36,7 +36,7 @@ func TestNewTokenBucketLimiter(t *testing.T) { podTokenBucket := util.NewTokenBucket(3, 3, nil) userTokenBucket := util.NewTokenBucket(2, 2, nil) requestTokenBucket := util.NewTokenBucket(1, 1, nil) - l := newTokenBucketLimiter(podTokenBucket, userTokenBucket, requestTokenBucket, true, func(tokens uint64, dataType store.StoreDataType) int64 { + l := newTokenBucketLimiter(podTokenBucket, userTokenBucket, requestTokenBucket, false, func(tokens uint64, dataType store.StoreDataType) int64 { if dataType == store.SeriesFetched { return int64(tokens) * 5 } @@ -45,30 +45,66 @@ func TestNewTokenBucketLimiter(t *testing.T) { // should force retrieve tokens from all buckets upon succeeding assert.NoError(t, l.ReserveWithType(2, store.PostingsFetched)) - assert.False(t, podTokenBucket.Retrieve(2)) + assert.Equal(t, int64(1), podTokenBucket.RemainingTokens()) + assert.Equal(t, int64(0), userTokenBucket.RemainingTokens()) + assert.Equal(t, int64(-1), requestTokenBucket.RemainingTokens()) // should fail if user token bucket is running low - podTokenBucket.Refund(3) + podTokenBucket.Refund(2) userTokenBucket.Refund(2) - requestTokenBucket.Refund(1) + requestTokenBucket.Refund(2) assert.ErrorContains(t, l.ReserveWithType(3, store.PostingsFetched), "not enough tokens in user token bucket") + assert.Equal(t, int64(3), podTokenBucket.RemainingTokens()) + assert.Equal(t, int64(2), userTokenBucket.RemainingTokens()) + assert.Equal(t, int64(1), requestTokenBucket.RemainingTokens()) // should fail if pod token bucket is running low - podTokenBucket.Refund(3) - userTokenBucket.Refund(2) - requestTokenBucket.Refund(1) - podTokenBucket.ForceRetrieve(3) + podTokenBucket.ForceRetrieve(2) assert.ErrorContains(t, l.ReserveWithType(2, store.PostingsFetched), "not enough tokens in pod token bucket") + assert.Equal(t, int64(1), podTokenBucket.RemainingTokens()) + assert.Equal(t, int64(2), userTokenBucket.RemainingTokens()) + assert.Equal(t, int64(1), requestTokenBucket.RemainingTokens()) // should retrieve different amount of tokens based on data type - podTokenBucket.Refund(3) - userTokenBucket.Refund(2) - requestTokenBucket.Refund(1) + podTokenBucket.Refund(2) assert.ErrorContains(t, l.ReserveWithType(1, store.SeriesFetched), "not enough tokens in user token bucket") + assert.Equal(t, int64(3), podTokenBucket.RemainingTokens()) + assert.Equal(t, int64(2), userTokenBucket.RemainingTokens()) + assert.Equal(t, int64(1), requestTokenBucket.RemainingTokens()) // should always succeed if retrieve token bucket has enough tokens, although shared buckets are empty podTokenBucket.ForceRetrieve(3) userTokenBucket.ForceRetrieve(2) assert.NoError(t, l.ReserveWithType(1, store.PostingsFetched)) + assert.Equal(t, int64(-1), podTokenBucket.RemainingTokens()) + assert.Equal(t, int64(-1), userTokenBucket.RemainingTokens()) + assert.Equal(t, int64(0), requestTokenBucket.RemainingTokens()) } +func TestNewTokenBucketLimter_DryRun(t *testing.T) { + podTokenBucket := util.NewTokenBucket(3, 3, nil) + userTokenBucket := util.NewTokenBucket(2, 2, nil) + requestTokenBucket := util.NewTokenBucket(1, 1, nil) + l := newTokenBucketLimiter(podTokenBucket, userTokenBucket, requestTokenBucket, true, func(tokens uint64, dataType store.StoreDataType) int64 { + if dataType == store.SeriesFetched { + return int64(tokens) * 5 + } + return int64(tokens) + }) + + // should force retrieve tokens from all buckets upon succeeding + assert.NoError(t, l.ReserveWithType(2, store.PostingsFetched)) + assert.False(t, podTokenBucket.Retrieve(2)) + assert.Equal(t, int64(1), podTokenBucket.RemainingTokens()) + assert.Equal(t, int64(0), userTokenBucket.RemainingTokens()) + assert.Equal(t, int64(-1), requestTokenBucket.RemainingTokens()) + + // should not fail even if tokens are not enough + podTokenBucket.Refund(2) + userTokenBucket.Refund(2) + requestTokenBucket.Refund(2) + assert.NoError(t, l.ReserveWithType(5, store.PostingsFetched)) + assert.Equal(t, int64(-2), podTokenBucket.RemainingTokens()) + assert.Equal(t, int64(-3), userTokenBucket.RemainingTokens()) + assert.Equal(t, int64(-4), requestTokenBucket.RemainingTokens()) +} diff --git a/pkg/util/token_bucket.go b/pkg/util/token_bucket.go index b5d2d1ede4..5186859135 100644 --- a/pkg/util/token_bucket.go +++ b/pkg/util/token_bucket.go @@ -74,6 +74,15 @@ func (t *TokenBucket) Refund(amount int64) { } } +func (t *TokenBucket) RemainingTokens() int64 { + t.mu.Lock() + defer t.mu.Unlock() + + t.updateTokens() + + return t.remainingTokens +} + func (t *TokenBucket) updateTokens() { now := time.Now() refilledTokens := int64(now.Sub(t.lastRefill).Seconds() * float64(t.refillRate)) diff --git a/pkg/util/token_bucket_test.go b/pkg/util/token_bucket_test.go index b74fdc2065..db67b9632b 100644 --- a/pkg/util/token_bucket_test.go +++ b/pkg/util/token_bucket_test.go @@ -14,21 +14,25 @@ func TestTokenBucket_Retrieve(t *testing.T) { assert.False(t, bucket.Retrieve(10)) time.Sleep(time.Second) assert.True(t, bucket.Retrieve(10)) + assert.Equal(t, int64(0), bucket.RemainingTokens()) } func TestTokenBucket_ForceRetrieve(t *testing.T) { bucket := NewTokenBucket(10, 600, nil) bucket.ForceRetrieve(20) + assert.Equal(t, int64(-10), bucket.RemainingTokens()) assert.False(t, bucket.Retrieve(10)) time.Sleep(time.Second) assert.True(t, bucket.Retrieve(10)) + assert.Equal(t, int64(0), bucket.RemainingTokens()) } func TestTokenBucket_Refund(t *testing.T) { bucket := NewTokenBucket(10, 600, nil) - bucket.ForceRetrieve(100) - bucket.Refund(100) + bucket.ForceRetrieve(10) + bucket.Refund(20) assert.True(t, bucket.Retrieve(10)) + assert.Equal(t, int64(0), bucket.RemainingTokens()) } From bd8b9e79255b84287aee42cf8c54145e1075cd22 Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Thu, 13 Jun 2024 23:07:59 -0700 Subject: [PATCH 08/26] Add doc Signed-off-by: Justin Jung --- docs/blocks-storage/store-gateway.md | 45 ++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/docs/blocks-storage/store-gateway.md b/docs/blocks-storage/store-gateway.md index 4b7e504fa4..87cc74ebeb 100644 --- a/docs/blocks-storage/store-gateway.md +++ b/docs/blocks-storage/store-gateway.md @@ -1466,6 +1466,51 @@ blocks_storage: # CLI flag: -blocks-storage.bucket-store.series-batch-size [series_batch_size: | default = 10000] + token_bucket_limiter: + # Whether token bucket limiter is enabled + # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.enabled + [enabled: | 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: | default = false] + + # Pod token bucket size + # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.pod-token-bucket-size + [pod_token_bucket_size: | default = 859832320] + + # User token bucket size + # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.user-token-bucket-size + [user_token_bucket_size: | default = 644874240] + + # Request token bucket size + # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.request-token-bucket-size + [request_token_bucket_size: | 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: | 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: | 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: | 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: | 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: | 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: | default = 0.5] + tsdb: # Local directory to store TSDBs in the ingesters. # CLI flag: -blocks-storage.tsdb.dir From f66aa50cb13eb429945ef748b3aedfeec5d416db Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Thu, 13 Jun 2024 23:18:43 -0700 Subject: [PATCH 09/26] Add changelog Signed-off-by: Justin Jung --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8328ea318a..23180cbc3d 100644 --- a/CHANGELOG.md +++ b/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 From 89408238f4641d908f13de3d1e47a578c4e1e078 Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Fri, 14 Jun 2024 09:54:59 -0700 Subject: [PATCH 10/26] Lint Signed-off-by: Justin Jung --- pkg/storegateway/limiter.go | 9 +++++---- pkg/storegateway/limiter_test.go | 3 ++- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/pkg/storegateway/limiter.go b/pkg/storegateway/limiter.go index 0f07e057aa..232d48c12a 100644 --- a/pkg/storegateway/limiter.go +++ b/pkg/storegateway/limiter.go @@ -4,14 +4,15 @@ import ( "fmt" "net/http" - "github.com/cortexproject/cortex/pkg/storage/tsdb" - "github.com/cortexproject/cortex/pkg/util" - util_log "github.com/cortexproject/cortex/pkg/util/log" - "github.com/cortexproject/cortex/pkg/util/validation" "github.com/go-kit/log/level" "github.com/prometheus/client_golang/prometheus" "github.com/thanos-io/thanos/pkg/store" "github.com/weaveworks/common/httpgrpc" + + "github.com/cortexproject/cortex/pkg/storage/tsdb" + "github.com/cortexproject/cortex/pkg/util" + util_log "github.com/cortexproject/cortex/pkg/util/log" + "github.com/cortexproject/cortex/pkg/util/validation" ) type limiter struct { diff --git a/pkg/storegateway/limiter_test.go b/pkg/storegateway/limiter_test.go index 82345117f5..785d40d977 100644 --- a/pkg/storegateway/limiter_test.go +++ b/pkg/storegateway/limiter_test.go @@ -3,10 +3,11 @@ package storegateway import ( "testing" - "github.com/cortexproject/cortex/pkg/util" "github.com/prometheus/client_golang/prometheus" "github.com/stretchr/testify/assert" "github.com/thanos-io/thanos/pkg/store" + + "github.com/cortexproject/cortex/pkg/util" ) func TestLimiter(t *testing.T) { From 705477867d64f99e917552911841e687babf3cc8 Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Fri, 14 Jun 2024 10:05:19 -0700 Subject: [PATCH 11/26] Do not create pod token bucket if the feature is not enabled Signed-off-by: Justin Jung --- pkg/storegateway/bucket_stores.go | 13 ++++++++----- pkg/storegateway/bucket_stores_test.go | 4 +++- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/pkg/storegateway/bucket_stores.go b/pkg/storegateway/bucket_stores.go index 9ec08251f1..702b2ca29d 100644 --- a/pkg/storegateway/bucket_stores.go +++ b/pkg/storegateway/bucket_stores.go @@ -120,11 +120,7 @@ func NewBucketStores(cfg tsdb.BlocksStorageConfig, shardingStrategy ShardingStra metaFetcherMetrics: NewMetadataFetcherMetrics(), queryGate: queryGate, partitioner: newGapBasedPartitioner(cfg.BucketStore.PartitionerMaxGapBytes, reg), - 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.", - })), - userTokenBuckets: make(map[string]*util.TokenBucket), + 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", @@ -154,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) } diff --git a/pkg/storegateway/bucket_stores_test.go b/pkg/storegateway/bucket_stores_test.go index 0a37b005c6..9e1f831f4d 100644 --- a/pkg/storegateway/bucket_stores_test.go +++ b/pkg/storegateway/bucket_stores_test.go @@ -763,7 +763,7 @@ func TestBucketStores_deleteLocalFilesForExcludedTenants(t *testing.T) { `), metricNames...)) } -func TestBucketStores_getUserTokenBucket(t *testing.T) { +func TestBucketStores_tokenBuckets(t *testing.T) { const ( user1 = "user-1" user2 = "user-2" @@ -791,6 +791,7 @@ func TestBucketStores_getUserTokenBucket(t *testing.T) { reg := prometheus.NewPedanticRegistry() stores, err := NewBucketStores(cfg, &sharding, objstore.WithNoopInstr(bucket), defaultLimitsOverrides(t), mockLoggingLevel(), log.NewNopLogger(), reg) assert.NoError(t, err) + assert.NotNil(t, stores.podTokenBucket) assert.NoError(t, stores.InitialSync(ctx)) assert.NotNil(t, stores.getUserTokenBucket("user-1")) @@ -813,6 +814,7 @@ func TestBucketStores_getUserTokenBucket(t *testing.T) { assert.NoError(t, err) assert.NoError(t, stores.InitialSync(ctx)) + assert.Nil(t, stores.podTokenBucket) assert.Nil(t, stores.getUserTokenBucket("user-1")) assert.Nil(t, stores.getUserTokenBucket("user-2")) } From cf39c550983e6a069b144f0c952ad6097ddc0318 Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Fri, 14 Jun 2024 10:25:30 -0700 Subject: [PATCH 12/26] More docs Signed-off-by: Justin Jung --- docs/blocks-storage/querier.md | 45 +++++++++++++++++++++ docs/configuration/config-file-reference.md | 45 +++++++++++++++++++++ 2 files changed, 90 insertions(+) diff --git a/docs/blocks-storage/querier.md b/docs/blocks-storage/querier.md index 1ca05c1b05..9c3015ce59 100644 --- a/docs/blocks-storage/querier.md +++ b/docs/blocks-storage/querier.md @@ -1341,6 +1341,51 @@ blocks_storage: # CLI flag: -blocks-storage.bucket-store.series-batch-size [series_batch_size: | default = 10000] + token_bucket_limiter: + # Whether token bucket limiter is enabled + # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.enabled + [enabled: | 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: | default = false] + + # Pod token bucket size + # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.pod-token-bucket-size + [pod_token_bucket_size: | default = 859832320] + + # User token bucket size + # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.user-token-bucket-size + [user_token_bucket_size: | default = 644874240] + + # Request token bucket size + # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.request-token-bucket-size + [request_token_bucket_size: | 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: | 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: | 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: | 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: | 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: | 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: | default = 0.5] + tsdb: # Local directory to store TSDBs in the ingesters. # CLI flag: -blocks-storage.tsdb.dir diff --git a/docs/configuration/config-file-reference.md b/docs/configuration/config-file-reference.md index ab0570c67e..8ee20abf6f 100644 --- a/docs/configuration/config-file-reference.md +++ b/docs/configuration/config-file-reference.md @@ -1899,6 +1899,51 @@ bucket_store: # CLI flag: -blocks-storage.bucket-store.series-batch-size [series_batch_size: | default = 10000] + token_bucket_limiter: + # Whether token bucket limiter is enabled + # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.enabled + [enabled: | 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: | default = false] + + # Pod token bucket size + # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.pod-token-bucket-size + [pod_token_bucket_size: | default = 859832320] + + # User token bucket size + # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.user-token-bucket-size + [user_token_bucket_size: | default = 644874240] + + # Request token bucket size + # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.request-token-bucket-size + [request_token_bucket_size: | 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: | 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: | 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: | 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: | 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: | 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: | default = 0.5] + tsdb: # Local directory to store TSDBs in the ingesters. # CLI flag: -blocks-storage.tsdb.dir From c9444748581a38b88ebc6b370a4304a3d40b851a Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Wed, 19 Jun 2024 17:22:15 -0700 Subject: [PATCH 13/26] Address comments Signed-off-by: Justin Jung --- pkg/storage/tsdb/config.go | 4 ++-- pkg/storage/tsdb/inmemory_index_cache.go | 2 +- pkg/storegateway/bucket_stores.go | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pkg/storage/tsdb/config.go b/pkg/storage/tsdb/config.go index 4b6fcb7e5e..203d290b6e 100644 --- a/pkg/storage/tsdb/config.go +++ b/pkg/storage/tsdb/config.go @@ -300,7 +300,7 @@ type BucketStoreConfig struct { type TokenBucketLimiterConfig struct { Enabled bool `yaml:"enabled"` DryRun bool `yaml:"dry_run"` - PodTokenBucketSize int64 `yaml:"pod_token_bucket_size"` + InstanceTokenBucketSize int64 `yaml:"instance_token_bucket_size"` UserTokenBucketSize int64 `yaml:"user_token_bucket_size"` RequestTokenBucketSize int64 `yaml:"request_token_bucket_size"` FetchedPostingsTokenFactor float64 `yaml:"fetched_postings_token_factor"` @@ -344,7 +344,7 @@ func (cfg *BucketStoreConfig) RegisterFlags(f *flag.FlagSet) { 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.InstanceTokenBucketSize, "blocks-storage.bucket-store.token-bucket-limiter.instance-token-bucket-size", int64(820*units.Mebibyte), "Instance 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") diff --git a/pkg/storage/tsdb/inmemory_index_cache.go b/pkg/storage/tsdb/inmemory_index_cache.go index 981e1bf6c4..1530bf99f8 100644 --- a/pkg/storage/tsdb/inmemory_index_cache.go +++ b/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.Debug(c.logger).Log( + level.Info(c.logger).Log( "msg", "item bigger than maxItemSizeBytes. Ignoring..", "maxItemSizeBytes", c.maxItemSizeBytes, "cacheType", typ, diff --git a/pkg/storegateway/bucket_stores.go b/pkg/storegateway/bucket_stores.go index 702b2ca29d..9c278dfbad 100644 --- a/pkg/storegateway/bucket_stores.go +++ b/pkg/storegateway/bucket_stores.go @@ -151,7 +151,7 @@ func NewBucketStores(cfg tsdb.BlocksStorageConfig, shardingStrategy ShardingStra } if u.cfg.BucketStore.TokenBucketLimiter.Enabled { - u.podTokenBucket = util.NewTokenBucket(cfg.BucketStore.TokenBucketLimiter.PodTokenBucketSize, cfg.BucketStore.TokenBucketLimiter.PodTokenBucketSize, promauto.With(reg).NewGauge(prometheus.GaugeOpts{ + u.podTokenBucket = util.NewTokenBucket(cfg.BucketStore.TokenBucketLimiter.InstanceTokenBucketSize, cfg.BucketStore.TokenBucketLimiter.InstanceTokenBucketSize, promauto.With(reg).NewGauge(prometheus.GaugeOpts{ Name: "cortex_bucket_stores_pod_token_bucket_remaining", Help: "Number of tokens left in pod token bucket.", })) From 4ed7bcb8f2c7676b8b821b491e523ecb5039f78f Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Wed, 19 Jun 2024 17:34:53 -0700 Subject: [PATCH 14/26] Rename podTokenBucket to instanceTokenBucket Signed-off-by: Justin Jung --- docs/blocks-storage/querier.md | 6 ++-- docs/blocks-storage/store-gateway.md | 6 ++-- docs/configuration/config-file-reference.md | 6 ++-- pkg/storegateway/bucket_stores.go | 10 +++--- pkg/storegateway/bucket_stores_test.go | 4 +-- pkg/storegateway/limiter.go | 18 +++++------ pkg/storegateway/limiter_test.go | 34 ++++++++++----------- 7 files changed, 42 insertions(+), 42 deletions(-) diff --git a/docs/blocks-storage/querier.md b/docs/blocks-storage/querier.md index 9c3015ce59..80e034a226 100644 --- a/docs/blocks-storage/querier.md +++ b/docs/blocks-storage/querier.md @@ -1350,9 +1350,9 @@ blocks_storage: # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.dry-run [dry_run: | default = false] - # Pod token bucket size - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.pod-token-bucket-size - [pod_token_bucket_size: | default = 859832320] + # Instance token bucket size + # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.instance-token-bucket-size + [instance_token_bucket_size: | default = 859832320] # User token bucket size # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.user-token-bucket-size diff --git a/docs/blocks-storage/store-gateway.md b/docs/blocks-storage/store-gateway.md index 87cc74ebeb..f459838c6d 100644 --- a/docs/blocks-storage/store-gateway.md +++ b/docs/blocks-storage/store-gateway.md @@ -1475,9 +1475,9 @@ blocks_storage: # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.dry-run [dry_run: | default = false] - # Pod token bucket size - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.pod-token-bucket-size - [pod_token_bucket_size: | default = 859832320] + # Instance token bucket size + # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.instance-token-bucket-size + [instance_token_bucket_size: | default = 859832320] # User token bucket size # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.user-token-bucket-size diff --git a/docs/configuration/config-file-reference.md b/docs/configuration/config-file-reference.md index 8ee20abf6f..145393cb9f 100644 --- a/docs/configuration/config-file-reference.md +++ b/docs/configuration/config-file-reference.md @@ -1908,9 +1908,9 @@ bucket_store: # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.dry-run [dry_run: | default = false] - # Pod token bucket size - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.pod-token-bucket-size - [pod_token_bucket_size: | default = 859832320] + # Instance token bucket size + # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.instance-token-bucket-size + [instance_token_bucket_size: | default = 859832320] # User token bucket size # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.user-token-bucket-size diff --git a/pkg/storegateway/bucket_stores.go b/pkg/storegateway/bucket_stores.go index 9c278dfbad..a7a8345a1b 100644 --- a/pkg/storegateway/bucket_stores.go +++ b/pkg/storegateway/bucket_stores.go @@ -73,7 +73,7 @@ type BucketStores struct { storesErrorsMu sync.RWMutex storesErrors map[string]error - podTokenBucket *util.TokenBucket + instanceTokenBucket *util.TokenBucket userTokenBucketsMu sync.RWMutex userTokenBuckets map[string]*util.TokenBucket @@ -151,9 +151,9 @@ func NewBucketStores(cfg tsdb.BlocksStorageConfig, shardingStrategy ShardingStra } if u.cfg.BucketStore.TokenBucketLimiter.Enabled { - u.podTokenBucket = util.NewTokenBucket(cfg.BucketStore.TokenBucketLimiter.InstanceTokenBucketSize, cfg.BucketStore.TokenBucketLimiter.InstanceTokenBucketSize, promauto.With(reg).NewGauge(prometheus.GaugeOpts{ - Name: "cortex_bucket_stores_pod_token_bucket_remaining", - Help: "Number of tokens left in pod token bucket.", + u.instanceTokenBucket = util.NewTokenBucket(cfg.BucketStore.TokenBucketLimiter.InstanceTokenBucketSize, cfg.BucketStore.TokenBucketLimiter.InstanceTokenBucketSize, promauto.With(reg).NewGauge(prometheus.GaugeOpts{ + Name: "cortex_bucket_stores_instance_token_bucket_remaining", + Help: "Number of tokens left in instance token bucket.", })) } @@ -643,7 +643,7 @@ func (u *BucketStores) getOrCreateStore(userID string) (*store.BucketStore, erro u.syncDirForUser(userID), newChunksLimiterFactory(u.limits, userID), newSeriesLimiterFactory(u.limits, userID), - newBytesLimiterFactory(u.limits, userID, u.podTokenBucket, u.getUserTokenBucket(userID), u.cfg.BucketStore.TokenBucketLimiter, u.getTokensToRetrieve), + newBytesLimiterFactory(u.limits, userID, u.instanceTokenBucket, 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 diff --git a/pkg/storegateway/bucket_stores_test.go b/pkg/storegateway/bucket_stores_test.go index 9e1f831f4d..6d25b6f3cc 100644 --- a/pkg/storegateway/bucket_stores_test.go +++ b/pkg/storegateway/bucket_stores_test.go @@ -791,7 +791,7 @@ func TestBucketStores_tokenBuckets(t *testing.T) { reg := prometheus.NewPedanticRegistry() stores, err := NewBucketStores(cfg, &sharding, objstore.WithNoopInstr(bucket), defaultLimitsOverrides(t), mockLoggingLevel(), log.NewNopLogger(), reg) assert.NoError(t, err) - assert.NotNil(t, stores.podTokenBucket) + assert.NotNil(t, stores.instanceTokenBucket) assert.NoError(t, stores.InitialSync(ctx)) assert.NotNil(t, stores.getUserTokenBucket("user-1")) @@ -814,7 +814,7 @@ func TestBucketStores_tokenBuckets(t *testing.T) { assert.NoError(t, err) assert.NoError(t, stores.InitialSync(ctx)) - assert.Nil(t, stores.podTokenBucket) + assert.Nil(t, stores.instanceTokenBucket) assert.Nil(t, stores.getUserTokenBucket("user-1")) assert.Nil(t, stores.getUserTokenBucket("user-2")) } diff --git a/pkg/storegateway/limiter.go b/pkg/storegateway/limiter.go index 232d48c12a..c9f94707cf 100644 --- a/pkg/storegateway/limiter.go +++ b/pkg/storegateway/limiter.go @@ -46,7 +46,7 @@ func (c *compositeLimiter) ReserveWithType(num uint64, dataType store.StoreDataT } type tokenBucketLimiter struct { - podTokenBucket *util.TokenBucket + instanceTokenBucket *util.TokenBucket userTokenBucket *util.TokenBucket requestTokenBucket *util.TokenBucket getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64 @@ -64,7 +64,7 @@ func (t *tokenBucketLimiter) ReserveWithType(num uint64, dataType store.StoreDat retrieved := t.requestTokenBucket.Retrieve(tokensToRetrieve) if retrieved { t.userTokenBucket.ForceRetrieve(tokensToRetrieve) - t.podTokenBucket.ForceRetrieve(tokensToRetrieve) + t.instanceTokenBucket.ForceRetrieve(tokensToRetrieve) return nil } @@ -75,14 +75,14 @@ func (t *tokenBucketLimiter) ReserveWithType(num uint64, dataType store.StoreDat if t.dryRun { t.requestTokenBucket.ForceRetrieve(tokensToRetrieve) t.userTokenBucket.ForceRetrieve(tokensToRetrieve) - t.podTokenBucket.ForceRetrieve(tokensToRetrieve) + t.instanceTokenBucket.ForceRetrieve(tokensToRetrieve) level.Warn(util_log.Logger).Log("msg", "not enough tokens in user token bucket", "dataType", dataType, "dataSize", num, "tokens", tokensToRetrieve) return nil } return fmt.Errorf("not enough tokens in user token bucket") } - retrieved = t.podTokenBucket.Retrieve(tokensToRetrieve) + retrieved = t.instanceTokenBucket.Retrieve(tokensToRetrieve) if !retrieved { t.userTokenBucket.Refund(tokensToRetrieve) @@ -90,7 +90,7 @@ func (t *tokenBucketLimiter) ReserveWithType(num uint64, dataType store.StoreDat if t.dryRun { // user bucket is already retrieved t.requestTokenBucket.ForceRetrieve(tokensToRetrieve) - t.podTokenBucket.ForceRetrieve(tokensToRetrieve) + t.instanceTokenBucket.ForceRetrieve(tokensToRetrieve) level.Warn(util_log.Logger).Log("msg", "not enough tokens in pod token bucket", "dataType", dataType, "dataSize", num, "tokens", tokensToRetrieve) return nil } @@ -101,9 +101,9 @@ func (t *tokenBucketLimiter) ReserveWithType(num uint64, dataType store.StoreDat return nil } -func newTokenBucketLimiter(podTokenBucket, userTokenBucket, requestTokenBucket *util.TokenBucket, dryRun bool, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64) *tokenBucketLimiter { +func newTokenBucketLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket *util.TokenBucket, dryRun bool, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64) *tokenBucketLimiter { return &tokenBucketLimiter{ - podTokenBucket: podTokenBucket, + instanceTokenBucket: instanceTokenBucket, userTokenBucket: userTokenBucket, requestTokenBucket: requestTokenBucket, getTokensToRetrieve: getTokensToRetrieve, @@ -131,7 +131,7 @@ func newSeriesLimiterFactory(limits *validation.Overrides, userID string) store. } } -func newBytesLimiterFactory(limits *validation.Overrides, userID string, podTokenBucket, userTokenBucket *util.TokenBucket, tokenBucketLimiterCfg tsdb.TokenBucketLimiterConfig, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64) store.BytesLimiterFactory { +func newBytesLimiterFactory(limits *validation.Overrides, userID string, instanceTokenBucket, userTokenBucket *util.TokenBucket, tokenBucketLimiterCfg tsdb.TokenBucketLimiterConfig, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64) store.BytesLimiterFactory { return func(failedCounter prometheus.Counter) store.BytesLimiter { limiters := []store.BytesLimiter{} // Since limit overrides could be live reloaded, we have to get the current user's limit @@ -140,7 +140,7 @@ func newBytesLimiterFactory(limits *validation.Overrides, userID string, podToke if tokenBucketLimiterCfg.Enabled { requestTokenBucket := util.NewTokenBucket(tokenBucketLimiterCfg.RequestTokenBucketSize, tokenBucketLimiterCfg.RequestTokenBucketSize, nil) - limiters = append(limiters, newTokenBucketLimiter(podTokenBucket, userTokenBucket, requestTokenBucket, tokenBucketLimiterCfg.DryRun, getTokensToRetrieve)) + limiters = append(limiters, newTokenBucketLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket, tokenBucketLimiterCfg.DryRun, getTokensToRetrieve)) } return &compositeLimiter{ diff --git a/pkg/storegateway/limiter_test.go b/pkg/storegateway/limiter_test.go index 785d40d977..c10a377d45 100644 --- a/pkg/storegateway/limiter_test.go +++ b/pkg/storegateway/limiter_test.go @@ -34,10 +34,10 @@ func TestCompositeLimiter(t *testing.T) { } func TestNewTokenBucketLimiter(t *testing.T) { - podTokenBucket := util.NewTokenBucket(3, 3, nil) + instanceTokenBucket := util.NewTokenBucket(3, 3, nil) userTokenBucket := util.NewTokenBucket(2, 2, nil) requestTokenBucket := util.NewTokenBucket(1, 1, nil) - l := newTokenBucketLimiter(podTokenBucket, userTokenBucket, requestTokenBucket, false, func(tokens uint64, dataType store.StoreDataType) int64 { + l := newTokenBucketLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket, false, func(tokens uint64, dataType store.StoreDataType) int64 { if dataType == store.SeriesFetched { return int64(tokens) * 5 } @@ -46,47 +46,47 @@ func TestNewTokenBucketLimiter(t *testing.T) { // should force retrieve tokens from all buckets upon succeeding assert.NoError(t, l.ReserveWithType(2, store.PostingsFetched)) - assert.Equal(t, int64(1), podTokenBucket.RemainingTokens()) + assert.Equal(t, int64(1), instanceTokenBucket.RemainingTokens()) assert.Equal(t, int64(0), userTokenBucket.RemainingTokens()) assert.Equal(t, int64(-1), requestTokenBucket.RemainingTokens()) // should fail if user token bucket is running low - podTokenBucket.Refund(2) + instanceTokenBucket.Refund(2) userTokenBucket.Refund(2) requestTokenBucket.Refund(2) assert.ErrorContains(t, l.ReserveWithType(3, store.PostingsFetched), "not enough tokens in user token bucket") - assert.Equal(t, int64(3), podTokenBucket.RemainingTokens()) + assert.Equal(t, int64(3), instanceTokenBucket.RemainingTokens()) assert.Equal(t, int64(2), userTokenBucket.RemainingTokens()) assert.Equal(t, int64(1), requestTokenBucket.RemainingTokens()) // should fail if pod token bucket is running low - podTokenBucket.ForceRetrieve(2) + instanceTokenBucket.ForceRetrieve(2) assert.ErrorContains(t, l.ReserveWithType(2, store.PostingsFetched), "not enough tokens in pod token bucket") - assert.Equal(t, int64(1), podTokenBucket.RemainingTokens()) + assert.Equal(t, int64(1), instanceTokenBucket.RemainingTokens()) assert.Equal(t, int64(2), userTokenBucket.RemainingTokens()) assert.Equal(t, int64(1), requestTokenBucket.RemainingTokens()) // should retrieve different amount of tokens based on data type - podTokenBucket.Refund(2) + instanceTokenBucket.Refund(2) assert.ErrorContains(t, l.ReserveWithType(1, store.SeriesFetched), "not enough tokens in user token bucket") - assert.Equal(t, int64(3), podTokenBucket.RemainingTokens()) + assert.Equal(t, int64(3), instanceTokenBucket.RemainingTokens()) assert.Equal(t, int64(2), userTokenBucket.RemainingTokens()) assert.Equal(t, int64(1), requestTokenBucket.RemainingTokens()) // should always succeed if retrieve token bucket has enough tokens, although shared buckets are empty - podTokenBucket.ForceRetrieve(3) + instanceTokenBucket.ForceRetrieve(3) userTokenBucket.ForceRetrieve(2) assert.NoError(t, l.ReserveWithType(1, store.PostingsFetched)) - assert.Equal(t, int64(-1), podTokenBucket.RemainingTokens()) + assert.Equal(t, int64(-1), instanceTokenBucket.RemainingTokens()) assert.Equal(t, int64(-1), userTokenBucket.RemainingTokens()) assert.Equal(t, int64(0), requestTokenBucket.RemainingTokens()) } func TestNewTokenBucketLimter_DryRun(t *testing.T) { - podTokenBucket := util.NewTokenBucket(3, 3, nil) + instanceTokenBucket := util.NewTokenBucket(3, 3, nil) userTokenBucket := util.NewTokenBucket(2, 2, nil) requestTokenBucket := util.NewTokenBucket(1, 1, nil) - l := newTokenBucketLimiter(podTokenBucket, userTokenBucket, requestTokenBucket, true, func(tokens uint64, dataType store.StoreDataType) int64 { + l := newTokenBucketLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket, true, func(tokens uint64, dataType store.StoreDataType) int64 { if dataType == store.SeriesFetched { return int64(tokens) * 5 } @@ -95,17 +95,17 @@ func TestNewTokenBucketLimter_DryRun(t *testing.T) { // should force retrieve tokens from all buckets upon succeeding assert.NoError(t, l.ReserveWithType(2, store.PostingsFetched)) - assert.False(t, podTokenBucket.Retrieve(2)) - assert.Equal(t, int64(1), podTokenBucket.RemainingTokens()) + assert.False(t, instanceTokenBucket.Retrieve(2)) + assert.Equal(t, int64(1), instanceTokenBucket.RemainingTokens()) assert.Equal(t, int64(0), userTokenBucket.RemainingTokens()) assert.Equal(t, int64(-1), requestTokenBucket.RemainingTokens()) // should not fail even if tokens are not enough - podTokenBucket.Refund(2) + instanceTokenBucket.Refund(2) userTokenBucket.Refund(2) requestTokenBucket.Refund(2) assert.NoError(t, l.ReserveWithType(5, store.PostingsFetched)) - assert.Equal(t, int64(-2), podTokenBucket.RemainingTokens()) + assert.Equal(t, int64(-2), instanceTokenBucket.RemainingTokens()) assert.Equal(t, int64(-3), userTokenBucket.RemainingTokens()) assert.Equal(t, int64(-4), requestTokenBucket.RemainingTokens()) } From 3219a44a027621cc26c35d4f0c31bbf4bdfdbb67 Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Wed, 26 Jun 2024 08:19:53 -0700 Subject: [PATCH 15/26] Updated default values Signed-off-by: Justin Jung --- docs/blocks-storage/querier.md | 12 ++++++------ docs/blocks-storage/store-gateway.md | 12 ++++++------ docs/configuration/config-file-reference.md | 12 ++++++------ pkg/storage/tsdb/config.go | 12 ++++++------ 4 files changed, 24 insertions(+), 24 deletions(-) diff --git a/docs/blocks-storage/querier.md b/docs/blocks-storage/querier.md index 80e034a226..6695a2e091 100644 --- a/docs/blocks-storage/querier.md +++ b/docs/blocks-storage/querier.md @@ -1364,27 +1364,27 @@ blocks_storage: # Multiplication factor used for fetched postings token # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-postings-token-factor - [fetched_postings_token_factor: | default = 2] + [fetched_postings_token_factor: | default = 0] # Multiplication factor used for touched postings token # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-postings-token-factor - [touched_postings_token_factor: | default = 2] + [touched_postings_token_factor: | default = 5] # Multiplication factor used for fetched series token # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-series-token-factor - [fetched_series_token_factor: | default = 2.5] + [fetched_series_token_factor: | default = 0] # Multiplication factor used for touched series token # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-series-token-factor - [touched_series_token_factor: | default = 10] + [touched_series_token_factor: | default = 25] # Multiplication factor used for fetched chunks token # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-chunks-token-factor - [fetched_chunks_token_factor: | default = 0.5] + [fetched_chunks_token_factor: | default = 0] # Multiplication factor used for touched chunks token # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-chunks-token-factor - [touched_chunks_token_factor: | default = 0.5] + [touched_chunks_token_factor: | default = 1] tsdb: # Local directory to store TSDBs in the ingesters. diff --git a/docs/blocks-storage/store-gateway.md b/docs/blocks-storage/store-gateway.md index f459838c6d..6ea708eb28 100644 --- a/docs/blocks-storage/store-gateway.md +++ b/docs/blocks-storage/store-gateway.md @@ -1489,27 +1489,27 @@ blocks_storage: # Multiplication factor used for fetched postings token # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-postings-token-factor - [fetched_postings_token_factor: | default = 2] + [fetched_postings_token_factor: | default = 0] # Multiplication factor used for touched postings token # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-postings-token-factor - [touched_postings_token_factor: | default = 2] + [touched_postings_token_factor: | default = 5] # Multiplication factor used for fetched series token # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-series-token-factor - [fetched_series_token_factor: | default = 2.5] + [fetched_series_token_factor: | default = 0] # Multiplication factor used for touched series token # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-series-token-factor - [touched_series_token_factor: | default = 10] + [touched_series_token_factor: | default = 25] # Multiplication factor used for fetched chunks token # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-chunks-token-factor - [fetched_chunks_token_factor: | default = 0.5] + [fetched_chunks_token_factor: | default = 0] # Multiplication factor used for touched chunks token # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-chunks-token-factor - [touched_chunks_token_factor: | default = 0.5] + [touched_chunks_token_factor: | default = 1] tsdb: # Local directory to store TSDBs in the ingesters. diff --git a/docs/configuration/config-file-reference.md b/docs/configuration/config-file-reference.md index 145393cb9f..734540a4b5 100644 --- a/docs/configuration/config-file-reference.md +++ b/docs/configuration/config-file-reference.md @@ -1922,27 +1922,27 @@ bucket_store: # Multiplication factor used for fetched postings token # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-postings-token-factor - [fetched_postings_token_factor: | default = 2] + [fetched_postings_token_factor: | default = 0] # Multiplication factor used for touched postings token # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-postings-token-factor - [touched_postings_token_factor: | default = 2] + [touched_postings_token_factor: | default = 5] # Multiplication factor used for fetched series token # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-series-token-factor - [fetched_series_token_factor: | default = 2.5] + [fetched_series_token_factor: | default = 0] # Multiplication factor used for touched series token # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-series-token-factor - [touched_series_token_factor: | default = 10] + [touched_series_token_factor: | default = 25] # Multiplication factor used for fetched chunks token # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-chunks-token-factor - [fetched_chunks_token_factor: | default = 0.5] + [fetched_chunks_token_factor: | default = 0] # Multiplication factor used for touched chunks token # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-chunks-token-factor - [touched_chunks_token_factor: | default = 0.5] + [touched_chunks_token_factor: | default = 1] tsdb: # Local directory to store TSDBs in the ingesters. diff --git a/pkg/storage/tsdb/config.go b/pkg/storage/tsdb/config.go index 203d290b6e..45542632ad 100644 --- a/pkg/storage/tsdb/config.go +++ b/pkg/storage/tsdb/config.go @@ -347,12 +347,12 @@ func (cfg *BucketStoreConfig) RegisterFlags(f *flag.FlagSet) { f.Int64Var(&cfg.TokenBucketLimiter.InstanceTokenBucketSize, "blocks-storage.bucket-store.token-bucket-limiter.instance-token-bucket-size", int64(820*units.Mebibyte), "Instance 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") + f.Float64Var(&cfg.TokenBucketLimiter.FetchedPostingsTokenFactor, "blocks-storage.bucket-store.token-bucket-limiter.fetched-postings-token-factor", 0, "Multiplication factor used for fetched postings token") + f.Float64Var(&cfg.TokenBucketLimiter.TouchedPostingsTokenFactor, "blocks-storage.bucket-store.token-bucket-limiter.touched-postings-token-factor", 5, "Multiplication factor used for touched postings token") + f.Float64Var(&cfg.TokenBucketLimiter.FetchedSeriesTokenFactor, "blocks-storage.bucket-store.token-bucket-limiter.fetched-series-token-factor", 0, "Multiplication factor used for fetched series token") + f.Float64Var(&cfg.TokenBucketLimiter.TouchedSeriesTokenFactor, "blocks-storage.bucket-store.token-bucket-limiter.touched-series-token-factor", 25, "Multiplication factor used for touched series token") + f.Float64Var(&cfg.TokenBucketLimiter.FetchedChunksTokenFactor, "blocks-storage.bucket-store.token-bucket-limiter.fetched-chunks-token-factor", 0, "Multiplication factor used for fetched chunks token") + f.Float64Var(&cfg.TokenBucketLimiter.TouchedChunksTokenFactor, "blocks-storage.bucket-store.token-bucket-limiter.touched-chunks-token-factor", 1, "Multiplication factor used for touched chunks token") } // Validate the config. From 52518a1abfd8011574ca0547f0d4866f8651a97d Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Thu, 27 Jun 2024 16:28:44 -0700 Subject: [PATCH 16/26] Rename TokenBucketLimiter to TokenBucketBytesLimiter Signed-off-by: Justin Jung --- docs/blocks-storage/querier.md | 24 +++++++++---------- docs/blocks-storage/store-gateway.md | 24 +++++++++---------- docs/configuration/config-file-reference.md | 24 +++++++++---------- pkg/storage/tsdb/config.go | 26 ++++++++++----------- pkg/storegateway/bucket_stores.go | 24 +++++++++---------- pkg/storegateway/bucket_stores_test.go | 16 ++++++------- pkg/storegateway/limiter.go | 18 +++++++------- pkg/storegateway/limiter_test.go | 6 ++--- 8 files changed, 81 insertions(+), 81 deletions(-) diff --git a/docs/blocks-storage/querier.md b/docs/blocks-storage/querier.md index 6695a2e091..f98f402447 100644 --- a/docs/blocks-storage/querier.md +++ b/docs/blocks-storage/querier.md @@ -1341,49 +1341,49 @@ blocks_storage: # CLI flag: -blocks-storage.bucket-store.series-batch-size [series_batch_size: | default = 10000] - token_bucket_limiter: + token_bucket_bytes_limiter: # Whether token bucket limiter is enabled - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.enabled + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.enabled [enabled: | default = false] # Whether the token bucket limiter is in dry run mode - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.dry-run + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.dry-run [dry_run: | default = false] # Instance token bucket size - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.instance-token-bucket-size + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.instance-token-bucket-size [instance_token_bucket_size: | default = 859832320] # User token bucket size - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.user-token-bucket-size + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.user-token-bucket-size [user_token_bucket_size: | default = 644874240] # Request token bucket size - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.request-token-bucket-size + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.request-token-bucket-size [request_token_bucket_size: | default = 4194304] # Multiplication factor used for fetched postings token - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-postings-token-factor + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-postings-token-factor [fetched_postings_token_factor: | default = 0] # Multiplication factor used for touched postings token - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-postings-token-factor + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-postings-token-factor [touched_postings_token_factor: | default = 5] # Multiplication factor used for fetched series token - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-series-token-factor + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-series-token-factor [fetched_series_token_factor: | default = 0] # Multiplication factor used for touched series token - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-series-token-factor + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-series-token-factor [touched_series_token_factor: | default = 25] # Multiplication factor used for fetched chunks token - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-chunks-token-factor + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-chunks-token-factor [fetched_chunks_token_factor: | default = 0] # Multiplication factor used for touched chunks token - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-chunks-token-factor + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-chunks-token-factor [touched_chunks_token_factor: | default = 1] tsdb: diff --git a/docs/blocks-storage/store-gateway.md b/docs/blocks-storage/store-gateway.md index 6ea708eb28..24232bf966 100644 --- a/docs/blocks-storage/store-gateway.md +++ b/docs/blocks-storage/store-gateway.md @@ -1466,49 +1466,49 @@ blocks_storage: # CLI flag: -blocks-storage.bucket-store.series-batch-size [series_batch_size: | default = 10000] - token_bucket_limiter: + token_bucket_bytes_limiter: # Whether token bucket limiter is enabled - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.enabled + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.enabled [enabled: | default = false] # Whether the token bucket limiter is in dry run mode - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.dry-run + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.dry-run [dry_run: | default = false] # Instance token bucket size - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.instance-token-bucket-size + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.instance-token-bucket-size [instance_token_bucket_size: | default = 859832320] # User token bucket size - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.user-token-bucket-size + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.user-token-bucket-size [user_token_bucket_size: | default = 644874240] # Request token bucket size - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.request-token-bucket-size + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.request-token-bucket-size [request_token_bucket_size: | default = 4194304] # Multiplication factor used for fetched postings token - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-postings-token-factor + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-postings-token-factor [fetched_postings_token_factor: | default = 0] # Multiplication factor used for touched postings token - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-postings-token-factor + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-postings-token-factor [touched_postings_token_factor: | default = 5] # Multiplication factor used for fetched series token - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-series-token-factor + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-series-token-factor [fetched_series_token_factor: | default = 0] # Multiplication factor used for touched series token - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-series-token-factor + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-series-token-factor [touched_series_token_factor: | default = 25] # Multiplication factor used for fetched chunks token - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-chunks-token-factor + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-chunks-token-factor [fetched_chunks_token_factor: | default = 0] # Multiplication factor used for touched chunks token - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-chunks-token-factor + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-chunks-token-factor [touched_chunks_token_factor: | default = 1] tsdb: diff --git a/docs/configuration/config-file-reference.md b/docs/configuration/config-file-reference.md index 734540a4b5..bc47adb96e 100644 --- a/docs/configuration/config-file-reference.md +++ b/docs/configuration/config-file-reference.md @@ -1899,49 +1899,49 @@ bucket_store: # CLI flag: -blocks-storage.bucket-store.series-batch-size [series_batch_size: | default = 10000] - token_bucket_limiter: + token_bucket_bytes_limiter: # Whether token bucket limiter is enabled - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.enabled + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.enabled [enabled: | default = false] # Whether the token bucket limiter is in dry run mode - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.dry-run + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.dry-run [dry_run: | default = false] # Instance token bucket size - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.instance-token-bucket-size + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.instance-token-bucket-size [instance_token_bucket_size: | default = 859832320] # User token bucket size - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.user-token-bucket-size + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.user-token-bucket-size [user_token_bucket_size: | default = 644874240] # Request token bucket size - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.request-token-bucket-size + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.request-token-bucket-size [request_token_bucket_size: | default = 4194304] # Multiplication factor used for fetched postings token - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-postings-token-factor + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-postings-token-factor [fetched_postings_token_factor: | default = 0] # Multiplication factor used for touched postings token - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-postings-token-factor + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-postings-token-factor [touched_postings_token_factor: | default = 5] # Multiplication factor used for fetched series token - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-series-token-factor + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-series-token-factor [fetched_series_token_factor: | default = 0] # Multiplication factor used for touched series token - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-series-token-factor + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-series-token-factor [touched_series_token_factor: | default = 25] # Multiplication factor used for fetched chunks token - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.fetched-chunks-token-factor + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-chunks-token-factor [fetched_chunks_token_factor: | default = 0] # Multiplication factor used for touched chunks token - # CLI flag: -blocks-storage.bucket-store.token-bucket-limiter.touched-chunks-token-factor + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-chunks-token-factor [touched_chunks_token_factor: | default = 1] tsdb: diff --git a/pkg/storage/tsdb/config.go b/pkg/storage/tsdb/config.go index 45542632ad..6a7720fae1 100644 --- a/pkg/storage/tsdb/config.go +++ b/pkg/storage/tsdb/config.go @@ -294,10 +294,10 @@ type BucketStoreConfig struct { SeriesBatchSize int `yaml:"series_batch_size"` // Token bucket configs - TokenBucketLimiter TokenBucketLimiterConfig `yaml:"token_bucket_limiter"` + TokenBucketBytesLimiter TokenBucketBytesLimiterConfig `yaml:"token_bucket_bytes_limiter"` } -type TokenBucketLimiterConfig struct { +type TokenBucketBytesLimiterConfig struct { Enabled bool `yaml:"enabled"` DryRun bool `yaml:"dry_run"` InstanceTokenBucketSize int64 `yaml:"instance_token_bucket_size"` @@ -342,17 +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.InstanceTokenBucketSize, "blocks-storage.bucket-store.token-bucket-limiter.instance-token-bucket-size", int64(820*units.Mebibyte), "Instance 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", 0, "Multiplication factor used for fetched postings token") - f.Float64Var(&cfg.TokenBucketLimiter.TouchedPostingsTokenFactor, "blocks-storage.bucket-store.token-bucket-limiter.touched-postings-token-factor", 5, "Multiplication factor used for touched postings token") - f.Float64Var(&cfg.TokenBucketLimiter.FetchedSeriesTokenFactor, "blocks-storage.bucket-store.token-bucket-limiter.fetched-series-token-factor", 0, "Multiplication factor used for fetched series token") - f.Float64Var(&cfg.TokenBucketLimiter.TouchedSeriesTokenFactor, "blocks-storage.bucket-store.token-bucket-limiter.touched-series-token-factor", 25, "Multiplication factor used for touched series token") - f.Float64Var(&cfg.TokenBucketLimiter.FetchedChunksTokenFactor, "blocks-storage.bucket-store.token-bucket-limiter.fetched-chunks-token-factor", 0, "Multiplication factor used for fetched chunks token") - f.Float64Var(&cfg.TokenBucketLimiter.TouchedChunksTokenFactor, "blocks-storage.bucket-store.token-bucket-limiter.touched-chunks-token-factor", 1, "Multiplication factor used for touched chunks token") + f.BoolVar(&cfg.TokenBucketBytesLimiter.Enabled, "blocks-storage.bucket-store.token-bucket-bytes-limiter.enabled", false, "Whether token bucket limiter is enabled") + f.BoolVar(&cfg.TokenBucketBytesLimiter.DryRun, "blocks-storage.bucket-store.token-bucket-bytes-limiter.dry-run", false, "Whether the token bucket limiter is in dry run mode") + f.Int64Var(&cfg.TokenBucketBytesLimiter.InstanceTokenBucketSize, "blocks-storage.bucket-store.token-bucket-bytes-limiter.instance-token-bucket-size", int64(820*units.Mebibyte), "Instance token bucket size") + f.Int64Var(&cfg.TokenBucketBytesLimiter.UserTokenBucketSize, "blocks-storage.bucket-store.token-bucket-bytes-limiter.user-token-bucket-size", int64(615*units.Mebibyte), "User token bucket size") + f.Int64Var(&cfg.TokenBucketBytesLimiter.RequestTokenBucketSize, "blocks-storage.bucket-store.token-bucket-bytes-limiter.request-token-bucket-size", int64(4*units.Mebibyte), "Request token bucket size") + f.Float64Var(&cfg.TokenBucketBytesLimiter.FetchedPostingsTokenFactor, "blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-postings-token-factor", 0, "Multiplication factor used for fetched postings token") + f.Float64Var(&cfg.TokenBucketBytesLimiter.TouchedPostingsTokenFactor, "blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-postings-token-factor", 5, "Multiplication factor used for touched postings token") + f.Float64Var(&cfg.TokenBucketBytesLimiter.FetchedSeriesTokenFactor, "blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-series-token-factor", 0, "Multiplication factor used for fetched series token") + f.Float64Var(&cfg.TokenBucketBytesLimiter.TouchedSeriesTokenFactor, "blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-series-token-factor", 25, "Multiplication factor used for touched series token") + f.Float64Var(&cfg.TokenBucketBytesLimiter.FetchedChunksTokenFactor, "blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-chunks-token-factor", 0, "Multiplication factor used for fetched chunks token") + f.Float64Var(&cfg.TokenBucketBytesLimiter.TouchedChunksTokenFactor, "blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-chunks-token-factor", 1, "Multiplication factor used for touched chunks token") } // Validate the config. diff --git a/pkg/storegateway/bucket_stores.go b/pkg/storegateway/bucket_stores.go index a7a8345a1b..dbad9e3b16 100644 --- a/pkg/storegateway/bucket_stores.go +++ b/pkg/storegateway/bucket_stores.go @@ -150,8 +150,8 @@ func NewBucketStores(cfg tsdb.BlocksStorageConfig, shardingStrategy ShardingStra return nil, errors.Wrap(err, "create chunks bytes pool") } - if u.cfg.BucketStore.TokenBucketLimiter.Enabled { - u.instanceTokenBucket = util.NewTokenBucket(cfg.BucketStore.TokenBucketLimiter.InstanceTokenBucketSize, cfg.BucketStore.TokenBucketLimiter.InstanceTokenBucketSize, promauto.With(reg).NewGauge(prometheus.GaugeOpts{ + if u.cfg.BucketStore.TokenBucketBytesLimiter.Enabled { + u.instanceTokenBucket = util.NewTokenBucket(cfg.BucketStore.TokenBucketBytesLimiter.InstanceTokenBucketSize, cfg.BucketStore.TokenBucketBytesLimiter.InstanceTokenBucketSize, promauto.With(reg).NewGauge(prometheus.GaugeOpts{ Name: "cortex_bucket_stores_instance_token_bucket_remaining", Help: "Number of tokens left in instance token bucket.", })) @@ -488,7 +488,7 @@ func (u *BucketStores) closeEmptyBucketStore(userID string) error { unlockInDefer = false u.storesMu.Unlock() - if u.cfg.BucketStore.TokenBucketLimiter.Enabled { + if u.cfg.BucketStore.TokenBucketBytesLimiter.Enabled { u.userTokenBucketsMu.Lock() delete(u.userTokenBuckets, userID) u.userTokenBucketsMu.Unlock() @@ -631,9 +631,9 @@ func (u *BucketStores) getOrCreateStore(userID string) (*store.BucketStore, erro bucketStoreOpts = append(bucketStoreOpts, store.WithDebugLogging()) } - if u.cfg.BucketStore.TokenBucketLimiter.Enabled { + if u.cfg.BucketStore.TokenBucketBytesLimiter.Enabled { u.userTokenBucketsMu.Lock() - u.userTokenBuckets[userID] = util.NewTokenBucket(u.cfg.BucketStore.TokenBucketLimiter.UserTokenBucketSize, u.cfg.BucketStore.TokenBucketLimiter.UserTokenBucketSize, nil) + u.userTokenBuckets[userID] = util.NewTokenBucket(u.cfg.BucketStore.TokenBucketBytesLimiter.UserTokenBucketSize, u.cfg.BucketStore.TokenBucketBytesLimiter.UserTokenBucketSize, nil) u.userTokenBucketsMu.Unlock() } @@ -643,7 +643,7 @@ func (u *BucketStores) getOrCreateStore(userID string) (*store.BucketStore, erro u.syncDirForUser(userID), newChunksLimiterFactory(u.limits, userID), newSeriesLimiterFactory(u.limits, userID), - newBytesLimiterFactory(u.limits, userID, u.instanceTokenBucket, u.getUserTokenBucket(userID), u.cfg.BucketStore.TokenBucketLimiter, u.getTokensToRetrieve), + newBytesLimiterFactory(u.limits, userID, u.instanceTokenBucket, u.getUserTokenBucket(userID), u.cfg.BucketStore.TokenBucketBytesLimiter, u.getTokensToRetrieve), u.partitioner, u.cfg.BucketStore.BlockSyncConcurrency, false, // No need to enable backward compatibility with Thanos pre 0.8.0 queriers @@ -715,17 +715,17 @@ func (u *BucketStores) getTokensToRetrieve(tokens uint64, dataType store.StoreDa tokensToRetrieve := float64(tokens) switch dataType { case store.PostingsFetched: - tokensToRetrieve *= u.cfg.BucketStore.TokenBucketLimiter.FetchedPostingsTokenFactor + tokensToRetrieve *= u.cfg.BucketStore.TokenBucketBytesLimiter.FetchedPostingsTokenFactor case store.PostingsTouched: - tokensToRetrieve *= u.cfg.BucketStore.TokenBucketLimiter.TouchedPostingsTokenFactor + tokensToRetrieve *= u.cfg.BucketStore.TokenBucketBytesLimiter.TouchedPostingsTokenFactor case store.SeriesFetched: - tokensToRetrieve *= u.cfg.BucketStore.TokenBucketLimiter.FetchedSeriesTokenFactor + tokensToRetrieve *= u.cfg.BucketStore.TokenBucketBytesLimiter.FetchedSeriesTokenFactor case store.SeriesTouched: - tokensToRetrieve *= u.cfg.BucketStore.TokenBucketLimiter.TouchedSeriesTokenFactor + tokensToRetrieve *= u.cfg.BucketStore.TokenBucketBytesLimiter.TouchedSeriesTokenFactor case store.ChunksFetched: - tokensToRetrieve *= u.cfg.BucketStore.TokenBucketLimiter.FetchedChunksTokenFactor + tokensToRetrieve *= u.cfg.BucketStore.TokenBucketBytesLimiter.FetchedChunksTokenFactor case store.ChunksTouched: - tokensToRetrieve *= u.cfg.BucketStore.TokenBucketLimiter.TouchedChunksTokenFactor + tokensToRetrieve *= u.cfg.BucketStore.TokenBucketBytesLimiter.TouchedChunksTokenFactor } return int64(tokensToRetrieve) } diff --git a/pkg/storegateway/bucket_stores_test.go b/pkg/storegateway/bucket_stores_test.go index 6d25b6f3cc..e24581c06e 100644 --- a/pkg/storegateway/bucket_stores_test.go +++ b/pkg/storegateway/bucket_stores_test.go @@ -771,7 +771,7 @@ func TestBucketStores_tokenBuckets(t *testing.T) { ctx := context.Background() cfg := prepareStorageConfig(t) - cfg.BucketStore.TokenBucketLimiter.Enabled = true + cfg.BucketStore.TokenBucketBytesLimiter.Enabled = true storageDir := t.TempDir() userToMetric := map[string]string{ @@ -807,7 +807,7 @@ func TestBucketStores_tokenBuckets(t *testing.T) { assert.Nil(t, stores.getUserTokenBucket("user-1")) assert.Nil(t, stores.getUserTokenBucket("user-2")) - cfg.BucketStore.TokenBucketLimiter.Enabled = false + cfg.BucketStore.TokenBucketBytesLimiter.Enabled = false sharding.users = []string{user1, user2} reg = prometheus.NewPedanticRegistry() stores, err = NewBucketStores(cfg, &sharding, objstore.WithNoopInstr(bucket), defaultLimitsOverrides(t), mockLoggingLevel(), log.NewNopLogger(), reg) @@ -821,12 +821,12 @@ func TestBucketStores_tokenBuckets(t *testing.T) { func TestBucketStores_getTokensToRetrieve(t *testing.T) { cfg := prepareStorageConfig(t) - cfg.BucketStore.TokenBucketLimiter.FetchedPostingsTokenFactor = 1 - cfg.BucketStore.TokenBucketLimiter.TouchedPostingsTokenFactor = 2 - cfg.BucketStore.TokenBucketLimiter.FetchedSeriesTokenFactor = 3 - cfg.BucketStore.TokenBucketLimiter.TouchedSeriesTokenFactor = 4 - cfg.BucketStore.TokenBucketLimiter.FetchedChunksTokenFactor = 0 - cfg.BucketStore.TokenBucketLimiter.TouchedChunksTokenFactor = 0.5 + cfg.BucketStore.TokenBucketBytesLimiter.FetchedPostingsTokenFactor = 1 + cfg.BucketStore.TokenBucketBytesLimiter.TouchedPostingsTokenFactor = 2 + cfg.BucketStore.TokenBucketBytesLimiter.FetchedSeriesTokenFactor = 3 + cfg.BucketStore.TokenBucketBytesLimiter.TouchedSeriesTokenFactor = 4 + cfg.BucketStore.TokenBucketBytesLimiter.FetchedChunksTokenFactor = 0 + cfg.BucketStore.TokenBucketBytesLimiter.TouchedChunksTokenFactor = 0.5 storageDir := t.TempDir() bucket, err := filesystem.NewBucketClient(filesystem.Config{Directory: storageDir}) diff --git a/pkg/storegateway/limiter.go b/pkg/storegateway/limiter.go index c9f94707cf..d377e9cccb 100644 --- a/pkg/storegateway/limiter.go +++ b/pkg/storegateway/limiter.go @@ -45,7 +45,7 @@ func (c *compositeLimiter) ReserveWithType(num uint64, dataType store.StoreDataT return nil } -type tokenBucketLimiter struct { +type tokenBucketBytesLimiter struct { instanceTokenBucket *util.TokenBucket userTokenBucket *util.TokenBucket requestTokenBucket *util.TokenBucket @@ -53,11 +53,11 @@ type tokenBucketLimiter struct { dryRun bool } -func (t *tokenBucketLimiter) Reserve(_ uint64) error { +func (t *tokenBucketBytesLimiter) Reserve(_ uint64) error { return nil } -func (t *tokenBucketLimiter) ReserveWithType(num uint64, dataType store.StoreDataType) error { +func (t *tokenBucketBytesLimiter) ReserveWithType(num uint64, dataType store.StoreDataType) error { tokensToRetrieve := t.getTokensToRetrieve(num, dataType) // check request bucket @@ -101,8 +101,8 @@ func (t *tokenBucketLimiter) ReserveWithType(num uint64, dataType store.StoreDat return nil } -func newTokenBucketLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket *util.TokenBucket, dryRun bool, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64) *tokenBucketLimiter { - return &tokenBucketLimiter{ +func newTokenBucketBytesLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket *util.TokenBucket, dryRun bool, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64) *tokenBucketBytesLimiter { + return &tokenBucketBytesLimiter{ instanceTokenBucket: instanceTokenBucket, userTokenBucket: userTokenBucket, requestTokenBucket: requestTokenBucket, @@ -131,16 +131,16 @@ func newSeriesLimiterFactory(limits *validation.Overrides, userID string) store. } } -func newBytesLimiterFactory(limits *validation.Overrides, userID string, instanceTokenBucket, userTokenBucket *util.TokenBucket, tokenBucketLimiterCfg tsdb.TokenBucketLimiterConfig, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64) store.BytesLimiterFactory { +func newBytesLimiterFactory(limits *validation.Overrides, userID string, instanceTokenBucket, userTokenBucket *util.TokenBucket, tokenBucketBytesLimiterCfg tsdb.TokenBucketBytesLimiterConfig, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64) store.BytesLimiterFactory { return func(failedCounter prometheus.Counter) store.BytesLimiter { limiters := []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. limiters = append(limiters, store.NewLimiter(uint64(limits.MaxDownloadedBytesPerRequest(userID)), failedCounter)) - if tokenBucketLimiterCfg.Enabled { - requestTokenBucket := util.NewTokenBucket(tokenBucketLimiterCfg.RequestTokenBucketSize, tokenBucketLimiterCfg.RequestTokenBucketSize, nil) - limiters = append(limiters, newTokenBucketLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket, tokenBucketLimiterCfg.DryRun, getTokensToRetrieve)) + if tokenBucketBytesLimiterCfg.Enabled { + requestTokenBucket := util.NewTokenBucket(tokenBucketBytesLimiterCfg.RequestTokenBucketSize, tokenBucketBytesLimiterCfg.RequestTokenBucketSize, nil) + limiters = append(limiters, newTokenBucketBytesLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket, tokenBucketBytesLimiterCfg.DryRun, getTokensToRetrieve)) } return &compositeLimiter{ diff --git a/pkg/storegateway/limiter_test.go b/pkg/storegateway/limiter_test.go index c10a377d45..33c3439787 100644 --- a/pkg/storegateway/limiter_test.go +++ b/pkg/storegateway/limiter_test.go @@ -33,11 +33,11 @@ func TestCompositeLimiter(t *testing.T) { assert.Error(t, l.ReserveWithType(1, store.PostingsFetched)) } -func TestNewTokenBucketLimiter(t *testing.T) { +func TestNewTokenBucketBytesLimiter(t *testing.T) { instanceTokenBucket := util.NewTokenBucket(3, 3, nil) userTokenBucket := util.NewTokenBucket(2, 2, nil) requestTokenBucket := util.NewTokenBucket(1, 1, nil) - l := newTokenBucketLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket, false, func(tokens uint64, dataType store.StoreDataType) int64 { + l := newTokenBucketBytesLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket, false, func(tokens uint64, dataType store.StoreDataType) int64 { if dataType == store.SeriesFetched { return int64(tokens) * 5 } @@ -86,7 +86,7 @@ func TestNewTokenBucketLimter_DryRun(t *testing.T) { instanceTokenBucket := util.NewTokenBucket(3, 3, nil) userTokenBucket := util.NewTokenBucket(2, 2, nil) requestTokenBucket := util.NewTokenBucket(1, 1, nil) - l := newTokenBucketLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket, true, func(tokens uint64, dataType store.StoreDataType) int64 { + l := newTokenBucketBytesLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket, true, func(tokens uint64, dataType store.StoreDataType) int64 { if dataType == store.SeriesFetched { return int64(tokens) * 5 } From c5e6968578a7683185ae43b591204d2476fccdac Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Thu, 27 Jun 2024 16:42:35 -0700 Subject: [PATCH 17/26] Changed error to httpgrpc Signed-off-by: Justin Jung --- pkg/storegateway/limiter.go | 2 +- pkg/storegateway/limiter_test.go | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pkg/storegateway/limiter.go b/pkg/storegateway/limiter.go index d377e9cccb..56e4f0a828 100644 --- a/pkg/storegateway/limiter.go +++ b/pkg/storegateway/limiter.go @@ -39,7 +39,7 @@ type compositeLimiter struct { func (c *compositeLimiter) ReserveWithType(num uint64, dataType store.StoreDataType) error { for _, l := range c.limiters { if err := l.ReserveWithType(num, dataType); err != nil { - return err + return httpgrpc.Errorf(http.StatusUnprocessableEntity, err.Error()) } } return nil diff --git a/pkg/storegateway/limiter_test.go b/pkg/storegateway/limiter_test.go index 33c3439787..01c33c41c6 100644 --- a/pkg/storegateway/limiter_test.go +++ b/pkg/storegateway/limiter_test.go @@ -30,7 +30,7 @@ func TestCompositeLimiter(t *testing.T) { } assert.NoError(t, l.ReserveWithType(1, store.PostingsFetched)) - assert.Error(t, l.ReserveWithType(1, store.PostingsFetched)) + assert.ErrorContains(t, l.ReserveWithType(1, store.PostingsFetched), "(422)") } func TestNewTokenBucketBytesLimiter(t *testing.T) { From d0c4fbd3e9d8c9428965a213e17114384cc365fd Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Tue, 2 Jul 2024 14:43:29 -0700 Subject: [PATCH 18/26] Nit Signed-off-by: Justin Jung --- pkg/util/token_bucket.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pkg/util/token_bucket.go b/pkg/util/token_bucket.go index 5186859135..dc54bfb613 100644 --- a/pkg/util/token_bucket.go +++ b/pkg/util/token_bucket.go @@ -72,6 +72,10 @@ func (t *TokenBucket) Refund(amount int64) { if t.remainingTokens > t.maxCapacity { t.remainingTokens = t.maxCapacity } + + if t.remainingTokensMetric != nil { + t.remainingTokensMetric.Set(float64(t.remainingTokens)) + } } func (t *TokenBucket) RemainingTokens() int64 { From dff09ce07dbd5e305f43d7d29f565157be4c513f Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Tue, 2 Jul 2024 15:01:44 -0700 Subject: [PATCH 19/26] Increment failure metric when token bucket returns error Signed-off-by: Justin Jung --- pkg/storegateway/limiter.go | 16 +++++++++++++--- pkg/storegateway/limiter_test.go | 4 ++-- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/pkg/storegateway/limiter.go b/pkg/storegateway/limiter.go index 56e4f0a828..f02e8b50de 100644 --- a/pkg/storegateway/limiter.go +++ b/pkg/storegateway/limiter.go @@ -3,6 +3,7 @@ package storegateway import ( "fmt" "net/http" + "sync" "github.com/go-kit/log/level" "github.com/prometheus/client_golang/prometheus" @@ -51,6 +52,10 @@ type tokenBucketBytesLimiter struct { requestTokenBucket *util.TokenBucket getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64 dryRun bool + + // Counter metric which we will increase if limit is exceeded. + failedCounter prometheus.Counter + failedOnce sync.Once } func (t *tokenBucketBytesLimiter) Reserve(_ uint64) error { @@ -79,6 +84,8 @@ func (t *tokenBucketBytesLimiter) ReserveWithType(num uint64, dataType store.Sto level.Warn(util_log.Logger).Log("msg", "not enough tokens in user token bucket", "dataType", dataType, "dataSize", num, "tokens", tokensToRetrieve) return nil } + // We need to protect from the counter being incremented twice due to concurrency + t.failedOnce.Do(t.failedCounter.Inc) return fmt.Errorf("not enough tokens in user token bucket") } @@ -94,6 +101,8 @@ func (t *tokenBucketBytesLimiter) ReserveWithType(num uint64, dataType store.Sto level.Warn(util_log.Logger).Log("msg", "not enough tokens in pod token bucket", "dataType", dataType, "dataSize", num, "tokens", tokensToRetrieve) return nil } + // We need to protect from the counter being incremented twice due to concurrency + t.failedOnce.Do(t.failedCounter.Inc) return fmt.Errorf("not enough tokens in pod token bucket") } @@ -101,13 +110,14 @@ func (t *tokenBucketBytesLimiter) ReserveWithType(num uint64, dataType store.Sto return nil } -func newTokenBucketBytesLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket *util.TokenBucket, dryRun bool, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64) *tokenBucketBytesLimiter { +func newTokenBucketBytesLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket *util.TokenBucket, dryRun bool, failedCounter prometheus.Counter, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64) *tokenBucketBytesLimiter { return &tokenBucketBytesLimiter{ instanceTokenBucket: instanceTokenBucket, userTokenBucket: userTokenBucket, requestTokenBucket: requestTokenBucket, - getTokensToRetrieve: getTokensToRetrieve, dryRun: dryRun, + failedCounter: failedCounter, + getTokensToRetrieve: getTokensToRetrieve, } } @@ -140,7 +150,7 @@ func newBytesLimiterFactory(limits *validation.Overrides, userID string, instanc if tokenBucketBytesLimiterCfg.Enabled { requestTokenBucket := util.NewTokenBucket(tokenBucketBytesLimiterCfg.RequestTokenBucketSize, tokenBucketBytesLimiterCfg.RequestTokenBucketSize, nil) - limiters = append(limiters, newTokenBucketBytesLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket, tokenBucketBytesLimiterCfg.DryRun, getTokensToRetrieve)) + limiters = append(limiters, newTokenBucketBytesLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket, tokenBucketBytesLimiterCfg.DryRun, failedCounter, getTokensToRetrieve)) } return &compositeLimiter{ diff --git a/pkg/storegateway/limiter_test.go b/pkg/storegateway/limiter_test.go index 01c33c41c6..64336ce7fd 100644 --- a/pkg/storegateway/limiter_test.go +++ b/pkg/storegateway/limiter_test.go @@ -37,7 +37,7 @@ func TestNewTokenBucketBytesLimiter(t *testing.T) { instanceTokenBucket := util.NewTokenBucket(3, 3, nil) userTokenBucket := util.NewTokenBucket(2, 2, nil) requestTokenBucket := util.NewTokenBucket(1, 1, nil) - l := newTokenBucketBytesLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket, false, func(tokens uint64, dataType store.StoreDataType) int64 { + l := newTokenBucketBytesLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket, false, prometheus.NewCounter(prometheus.CounterOpts{}), func(tokens uint64, dataType store.StoreDataType) int64 { if dataType == store.SeriesFetched { return int64(tokens) * 5 } @@ -86,7 +86,7 @@ func TestNewTokenBucketLimter_DryRun(t *testing.T) { instanceTokenBucket := util.NewTokenBucket(3, 3, nil) userTokenBucket := util.NewTokenBucket(2, 2, nil) requestTokenBucket := util.NewTokenBucket(1, 1, nil) - l := newTokenBucketBytesLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket, true, func(tokens uint64, dataType store.StoreDataType) int64 { + l := newTokenBucketBytesLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket, true, prometheus.NewCounter(prometheus.CounterOpts{}), func(tokens uint64, dataType store.StoreDataType) int64 { if dataType == store.SeriesFetched { return int64(tokens) * 5 } From 846708701f5881327e2e3d48b2bd402c7dc0a285 Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Thu, 11 Jul 2024 13:40:54 -0700 Subject: [PATCH 20/26] Simplify token bucket by making Retrieve to always deduct token Signed-off-by: Justin Jung --- pkg/storegateway/bucket_stores.go | 6 +- pkg/storegateway/limiter.go | 58 ++++------- pkg/storegateway/limiter_test.go | 161 ++++++++++++++++-------------- pkg/util/token_bucket.go | 49 +-------- pkg/util/token_bucket_test.go | 29 +----- 5 files changed, 121 insertions(+), 182 deletions(-) diff --git a/pkg/storegateway/bucket_stores.go b/pkg/storegateway/bucket_stores.go index dbad9e3b16..ebc409cc29 100644 --- a/pkg/storegateway/bucket_stores.go +++ b/pkg/storegateway/bucket_stores.go @@ -151,7 +151,7 @@ func NewBucketStores(cfg tsdb.BlocksStorageConfig, shardingStrategy ShardingStra } if u.cfg.BucketStore.TokenBucketBytesLimiter.Enabled { - u.instanceTokenBucket = util.NewTokenBucket(cfg.BucketStore.TokenBucketBytesLimiter.InstanceTokenBucketSize, cfg.BucketStore.TokenBucketBytesLimiter.InstanceTokenBucketSize, promauto.With(reg).NewGauge(prometheus.GaugeOpts{ + u.instanceTokenBucket = util.NewTokenBucket(cfg.BucketStore.TokenBucketBytesLimiter.InstanceTokenBucketSize, promauto.With(reg).NewGauge(prometheus.GaugeOpts{ Name: "cortex_bucket_stores_instance_token_bucket_remaining", Help: "Number of tokens left in instance token bucket.", })) @@ -633,7 +633,7 @@ func (u *BucketStores) getOrCreateStore(userID string) (*store.BucketStore, erro if u.cfg.BucketStore.TokenBucketBytesLimiter.Enabled { u.userTokenBucketsMu.Lock() - u.userTokenBuckets[userID] = util.NewTokenBucket(u.cfg.BucketStore.TokenBucketBytesLimiter.UserTokenBucketSize, u.cfg.BucketStore.TokenBucketBytesLimiter.UserTokenBucketSize, nil) + u.userTokenBuckets[userID] = util.NewTokenBucket(u.cfg.BucketStore.TokenBucketBytesLimiter.UserTokenBucketSize, nil) u.userTokenBucketsMu.Unlock() } @@ -643,7 +643,7 @@ func (u *BucketStores) getOrCreateStore(userID string) (*store.BucketStore, erro u.syncDirForUser(userID), newChunksLimiterFactory(u.limits, userID), newSeriesLimiterFactory(u.limits, userID), - newBytesLimiterFactory(u.limits, userID, u.instanceTokenBucket, u.getUserTokenBucket(userID), u.cfg.BucketStore.TokenBucketBytesLimiter, u.getTokensToRetrieve), + newBytesLimiterFactory(u.limits, userID, u.getUserTokenBucket(userID), u.instanceTokenBucket, u.cfg.BucketStore.TokenBucketBytesLimiter, u.getTokensToRetrieve), u.partitioner, u.cfg.BucketStore.BlockSyncConcurrency, false, // No need to enable backward compatibility with Thanos pre 0.8.0 queriers diff --git a/pkg/storegateway/limiter.go b/pkg/storegateway/limiter.go index f02e8b50de..b9e77d6322 100644 --- a/pkg/storegateway/limiter.go +++ b/pkg/storegateway/limiter.go @@ -65,56 +65,42 @@ func (t *tokenBucketBytesLimiter) Reserve(_ uint64) error { func (t *tokenBucketBytesLimiter) ReserveWithType(num uint64, dataType store.StoreDataType) error { tokensToRetrieve := t.getTokensToRetrieve(num, dataType) - // check request bucket - retrieved := t.requestTokenBucket.Retrieve(tokensToRetrieve) - if retrieved { - t.userTokenBucket.ForceRetrieve(tokensToRetrieve) - t.instanceTokenBucket.ForceRetrieve(tokensToRetrieve) + requestTokenRemaining := t.requestTokenBucket.Retrieve(tokensToRetrieve) + userTokenRemaining := t.userTokenBucket.Retrieve(tokensToRetrieve) + instanceTokenRemaining := t.instanceTokenBucket.Retrieve(tokensToRetrieve) + + // if we can retrieve from request bucket, let the request go through + if requestTokenRemaining >= 0 { return nil } - // if we can't retrieve from request bucket, check shared buckets - retrieved = t.userTokenBucket.Retrieve(tokensToRetrieve) - if !retrieved { - // if dry run, force retrieve all tokens and return nil - if t.dryRun { - t.requestTokenBucket.ForceRetrieve(tokensToRetrieve) - t.userTokenBucket.ForceRetrieve(tokensToRetrieve) - t.instanceTokenBucket.ForceRetrieve(tokensToRetrieve) - level.Warn(util_log.Logger).Log("msg", "not enough tokens in user token bucket", "dataType", dataType, "dataSize", num, "tokens", tokensToRetrieve) - return nil - } - // We need to protect from the counter being incremented twice due to concurrency - t.failedOnce.Do(t.failedCounter.Inc) - return fmt.Errorf("not enough tokens in user token bucket") - } + errMsg := "" - retrieved = t.instanceTokenBucket.Retrieve(tokensToRetrieve) - if !retrieved { - t.userTokenBucket.Refund(tokensToRetrieve) + if userTokenRemaining < 0 { + errMsg = "not enough tokens in user token bucket" + } else if instanceTokenRemaining < 0 { + errMsg = "not enough tokens in instance token bucket" + } - // if dry run, force retrieve all tokens and return nil + if errMsg != "" { if t.dryRun { - // user bucket is already retrieved - t.requestTokenBucket.ForceRetrieve(tokensToRetrieve) - t.instanceTokenBucket.ForceRetrieve(tokensToRetrieve) - level.Warn(util_log.Logger).Log("msg", "not enough tokens in pod token bucket", "dataType", dataType, "dataSize", num, "tokens", tokensToRetrieve) + level.Warn(util_log.Logger).Log("msg", errMsg, "dataType", dataType, "dataSize", num, "tokens", tokensToRetrieve) return nil } + // We need to protect from the counter being incremented twice due to concurrency t.failedOnce.Do(t.failedCounter.Inc) - return fmt.Errorf("not enough tokens in pod token bucket") + return fmt.Errorf(errMsg) } - t.requestTokenBucket.ForceRetrieve(tokensToRetrieve) return nil } -func newTokenBucketBytesLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket *util.TokenBucket, dryRun bool, failedCounter prometheus.Counter, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64) *tokenBucketBytesLimiter { +func newTokenBucketBytesLimiter(requestTokenBucket, userTokenBucket, instanceTokenBucket *util.TokenBucket, dryRun bool, failedCounter prometheus.Counter, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64) *tokenBucketBytesLimiter { return &tokenBucketBytesLimiter{ - instanceTokenBucket: instanceTokenBucket, - userTokenBucket: userTokenBucket, requestTokenBucket: requestTokenBucket, + userTokenBucket: userTokenBucket, + instanceTokenBucket: instanceTokenBucket, dryRun: dryRun, failedCounter: failedCounter, getTokensToRetrieve: getTokensToRetrieve, @@ -141,7 +127,7 @@ func newSeriesLimiterFactory(limits *validation.Overrides, userID string) store. } } -func newBytesLimiterFactory(limits *validation.Overrides, userID string, instanceTokenBucket, userTokenBucket *util.TokenBucket, tokenBucketBytesLimiterCfg tsdb.TokenBucketBytesLimiterConfig, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64) store.BytesLimiterFactory { +func newBytesLimiterFactory(limits *validation.Overrides, userID string, userTokenBucket, instanceTokenBucket *util.TokenBucket, tokenBucketBytesLimiterCfg tsdb.TokenBucketBytesLimiterConfig, getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64) store.BytesLimiterFactory { return func(failedCounter prometheus.Counter) store.BytesLimiter { limiters := []store.BytesLimiter{} // Since limit overrides could be live reloaded, we have to get the current user's limit @@ -149,8 +135,8 @@ func newBytesLimiterFactory(limits *validation.Overrides, userID string, instanc limiters = append(limiters, store.NewLimiter(uint64(limits.MaxDownloadedBytesPerRequest(userID)), failedCounter)) if tokenBucketBytesLimiterCfg.Enabled { - requestTokenBucket := util.NewTokenBucket(tokenBucketBytesLimiterCfg.RequestTokenBucketSize, tokenBucketBytesLimiterCfg.RequestTokenBucketSize, nil) - limiters = append(limiters, newTokenBucketBytesLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket, tokenBucketBytesLimiterCfg.DryRun, failedCounter, getTokensToRetrieve)) + requestTokenBucket := util.NewTokenBucket(tokenBucketBytesLimiterCfg.RequestTokenBucketSize, nil) + limiters = append(limiters, newTokenBucketBytesLimiter(requestTokenBucket, userTokenBucket, instanceTokenBucket, tokenBucketBytesLimiterCfg.DryRun, failedCounter, getTokensToRetrieve)) } return &compositeLimiter{ diff --git a/pkg/storegateway/limiter_test.go b/pkg/storegateway/limiter_test.go index 64336ce7fd..c02d4e081e 100644 --- a/pkg/storegateway/limiter_test.go +++ b/pkg/storegateway/limiter_test.go @@ -34,78 +34,93 @@ func TestCompositeLimiter(t *testing.T) { } func TestNewTokenBucketBytesLimiter(t *testing.T) { - instanceTokenBucket := util.NewTokenBucket(3, 3, nil) - userTokenBucket := util.NewTokenBucket(2, 2, nil) - requestTokenBucket := util.NewTokenBucket(1, 1, nil) - l := newTokenBucketBytesLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket, false, prometheus.NewCounter(prometheus.CounterOpts{}), func(tokens uint64, dataType store.StoreDataType) int64 { - if dataType == store.SeriesFetched { - return int64(tokens) * 5 - } - return int64(tokens) - }) - - // should force retrieve tokens from all buckets upon succeeding - assert.NoError(t, l.ReserveWithType(2, store.PostingsFetched)) - assert.Equal(t, int64(1), instanceTokenBucket.RemainingTokens()) - assert.Equal(t, int64(0), userTokenBucket.RemainingTokens()) - assert.Equal(t, int64(-1), requestTokenBucket.RemainingTokens()) - - // should fail if user token bucket is running low - instanceTokenBucket.Refund(2) - userTokenBucket.Refund(2) - requestTokenBucket.Refund(2) - assert.ErrorContains(t, l.ReserveWithType(3, store.PostingsFetched), "not enough tokens in user token bucket") - assert.Equal(t, int64(3), instanceTokenBucket.RemainingTokens()) - assert.Equal(t, int64(2), userTokenBucket.RemainingTokens()) - assert.Equal(t, int64(1), requestTokenBucket.RemainingTokens()) - - // should fail if pod token bucket is running low - instanceTokenBucket.ForceRetrieve(2) - assert.ErrorContains(t, l.ReserveWithType(2, store.PostingsFetched), "not enough tokens in pod token bucket") - assert.Equal(t, int64(1), instanceTokenBucket.RemainingTokens()) - assert.Equal(t, int64(2), userTokenBucket.RemainingTokens()) - assert.Equal(t, int64(1), requestTokenBucket.RemainingTokens()) - - // should retrieve different amount of tokens based on data type - instanceTokenBucket.Refund(2) - assert.ErrorContains(t, l.ReserveWithType(1, store.SeriesFetched), "not enough tokens in user token bucket") - assert.Equal(t, int64(3), instanceTokenBucket.RemainingTokens()) - assert.Equal(t, int64(2), userTokenBucket.RemainingTokens()) - assert.Equal(t, int64(1), requestTokenBucket.RemainingTokens()) - - // should always succeed if retrieve token bucket has enough tokens, although shared buckets are empty - instanceTokenBucket.ForceRetrieve(3) - userTokenBucket.ForceRetrieve(2) - assert.NoError(t, l.ReserveWithType(1, store.PostingsFetched)) - assert.Equal(t, int64(-1), instanceTokenBucket.RemainingTokens()) - assert.Equal(t, int64(-1), userTokenBucket.RemainingTokens()) - assert.Equal(t, int64(0), requestTokenBucket.RemainingTokens()) -} - -func TestNewTokenBucketLimter_DryRun(t *testing.T) { - instanceTokenBucket := util.NewTokenBucket(3, 3, nil) - userTokenBucket := util.NewTokenBucket(2, 2, nil) - requestTokenBucket := util.NewTokenBucket(1, 1, nil) - l := newTokenBucketBytesLimiter(instanceTokenBucket, userTokenBucket, requestTokenBucket, true, prometheus.NewCounter(prometheus.CounterOpts{}), func(tokens uint64, dataType store.StoreDataType) int64 { - if dataType == store.SeriesFetched { - return int64(tokens) * 5 - } - return int64(tokens) - }) - - // should force retrieve tokens from all buckets upon succeeding - assert.NoError(t, l.ReserveWithType(2, store.PostingsFetched)) - assert.False(t, instanceTokenBucket.Retrieve(2)) - assert.Equal(t, int64(1), instanceTokenBucket.RemainingTokens()) - assert.Equal(t, int64(0), userTokenBucket.RemainingTokens()) - assert.Equal(t, int64(-1), requestTokenBucket.RemainingTokens()) + tests := map[string]struct { + tokenToRetrieve uint64 + requestTokenBucketSize int64 + userTokenBucketSize int64 + instanceTokenBucketSize int64 + expectedRequestTokenRemaining int64 + expectedUserTokenRemaining int64 + expectedInstanceTokenRemaining int64 + expectedErrStr string + getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64 + dryRun bool + }{ + "should retrieve buckets from all buckets": { + tokenToRetrieve: 1, + requestTokenBucketSize: 1, + userTokenBucketSize: 1, + instanceTokenBucketSize: 1, + }, + "should succeed if there is enough request token, regardless of user or instance bucket": { + tokenToRetrieve: 1, + requestTokenBucketSize: 1, + userTokenBucketSize: 0, + instanceTokenBucketSize: 0, + expectedUserTokenRemaining: -1, + expectedInstanceTokenRemaining: -1, + }, + "should fail if not enough user tokens remaining": { + tokenToRetrieve: 2, + requestTokenBucketSize: 1, + userTokenBucketSize: 1, + instanceTokenBucketSize: 2, + expectedErrStr: "not enough tokens in user token bucket", + expectedRequestTokenRemaining: -1, + expectedUserTokenRemaining: -1, + }, + "should fail if not enough instance tokens remaining": { + tokenToRetrieve: 2, + requestTokenBucketSize: 1, + userTokenBucketSize: 2, + instanceTokenBucketSize: 1, + expectedErrStr: "not enough tokens in instance token bucket", + expectedRequestTokenRemaining: -1, + expectedInstanceTokenRemaining: -1, + }, + "should use getTokensToRetrieve to calculate tokens": { + tokenToRetrieve: 1, + getTokensToRetrieve: func(tokens uint64, dataType store.StoreDataType) int64 { + if dataType == store.PostingsFetched { + return 0 + } + return 1 + }, + }, + "should not fail if dryRun": { + tokenToRetrieve: 1, + expectedRequestTokenRemaining: -1, + expectedUserTokenRemaining: -1, + expectedInstanceTokenRemaining: -1, + dryRun: true, + }, + } - // should not fail even if tokens are not enough - instanceTokenBucket.Refund(2) - userTokenBucket.Refund(2) - requestTokenBucket.Refund(2) - assert.NoError(t, l.ReserveWithType(5, store.PostingsFetched)) - assert.Equal(t, int64(-2), instanceTokenBucket.RemainingTokens()) - assert.Equal(t, int64(-3), userTokenBucket.RemainingTokens()) - assert.Equal(t, int64(-4), requestTokenBucket.RemainingTokens()) + for testName, testData := range tests { + t.Run(testName, func(t *testing.T) { + requestTokenBucket := util.NewTokenBucket(testData.requestTokenBucketSize, nil) + userTokenBucket := util.NewTokenBucket(testData.userTokenBucketSize, nil) + instanceTokenBucket := util.NewTokenBucket(testData.instanceTokenBucketSize, nil) + + getTokensToRetrieve := func(tokens uint64, dataType store.StoreDataType) int64 { + return int64(tokens) + } + if testData.getTokensToRetrieve != nil { + getTokensToRetrieve = testData.getTokensToRetrieve + } + l := newTokenBucketBytesLimiter(requestTokenBucket, userTokenBucket, instanceTokenBucket, testData.dryRun, prometheus.NewCounter(prometheus.CounterOpts{}), getTokensToRetrieve) + + err := l.ReserveWithType(testData.tokenToRetrieve, store.PostingsFetched) + + assert.Equal(t, testData.expectedRequestTokenRemaining, requestTokenBucket.Retrieve(0)) + assert.Equal(t, testData.expectedUserTokenRemaining, userTokenBucket.Retrieve(0)) + assert.Equal(t, testData.expectedInstanceTokenRemaining, instanceTokenBucket.Retrieve(0)) + + if testData.expectedErrStr != "" { + assert.ErrorContains(t, err, testData.expectedErrStr) + } else { + assert.NoError(t, err) + } + }) + } } diff --git a/pkg/util/token_bucket.go b/pkg/util/token_bucket.go index dc54bfb613..3c3fdcbb62 100644 --- a/pkg/util/token_bucket.go +++ b/pkg/util/token_bucket.go @@ -17,7 +17,7 @@ type TokenBucket struct { remainingTokensMetric prometheus.Gauge } -func NewTokenBucket(maxCapacity, refillRate int64, remainingTokensMetric prometheus.Gauge) *TokenBucket { +func NewTokenBucket(maxCapacity int64, remainingTokensMetric prometheus.Gauge) *TokenBucket { if remainingTokensMetric != nil { remainingTokensMetric.Set(float64(maxCapacity)) } @@ -25,64 +25,23 @@ func NewTokenBucket(maxCapacity, refillRate int64, remainingTokensMetric prometh return &TokenBucket{ remainingTokens: maxCapacity, maxCapacity: maxCapacity, - refillRate: refillRate, + refillRate: maxCapacity, lastRefill: time.Now(), remainingTokensMetric: remainingTokensMetric, } } -func (t *TokenBucket) Retrieve(amount int64) bool { +// Retrieve always deducts token, even if there is not enough remaining tokens. +func (t *TokenBucket) Retrieve(amount int64) int64 { t.mu.Lock() defer t.mu.Unlock() t.updateTokens() - - if t.remainingTokens < amount { - return false - } - t.remainingTokens -= amount - if t.remainingTokensMetric != nil { - t.remainingTokensMetric.Set(float64(t.remainingTokens)) - } - - return true -} - -func (t *TokenBucket) ForceRetrieve(amount int64) { - t.mu.Lock() - defer t.mu.Unlock() - t.updateTokens() - - t.remainingTokens -= amount if t.remainingTokensMetric != nil { t.remainingTokensMetric.Set(float64(t.remainingTokens)) } -} - -func (t *TokenBucket) Refund(amount int64) { - t.mu.Lock() - defer t.mu.Unlock() - - t.updateTokens() - - t.remainingTokens += amount - - if t.remainingTokens > t.maxCapacity { - t.remainingTokens = t.maxCapacity - } - - if t.remainingTokensMetric != nil { - t.remainingTokensMetric.Set(float64(t.remainingTokens)) - } -} - -func (t *TokenBucket) RemainingTokens() int64 { - t.mu.Lock() - defer t.mu.Unlock() - - t.updateTokens() return t.remainingTokens } diff --git a/pkg/util/token_bucket_test.go b/pkg/util/token_bucket_test.go index db67b9632b..a5a0bd932d 100644 --- a/pkg/util/token_bucket_test.go +++ b/pkg/util/token_bucket_test.go @@ -8,31 +8,10 @@ import ( ) func TestTokenBucket_Retrieve(t *testing.T) { - bucket := NewTokenBucket(10, 600, nil) + bucket := NewTokenBucket(10, nil) - assert.True(t, bucket.Retrieve(10)) - assert.False(t, bucket.Retrieve(10)) + assert.Equal(t, int64(0), bucket.Retrieve(10)) + assert.Negative(t, bucket.Retrieve(1)) time.Sleep(time.Second) - assert.True(t, bucket.Retrieve(10)) - assert.Equal(t, int64(0), bucket.RemainingTokens()) -} - -func TestTokenBucket_ForceRetrieve(t *testing.T) { - bucket := NewTokenBucket(10, 600, nil) - - bucket.ForceRetrieve(20) - assert.Equal(t, int64(-10), bucket.RemainingTokens()) - assert.False(t, bucket.Retrieve(10)) - time.Sleep(time.Second) - assert.True(t, bucket.Retrieve(10)) - assert.Equal(t, int64(0), bucket.RemainingTokens()) -} - -func TestTokenBucket_Refund(t *testing.T) { - bucket := NewTokenBucket(10, 600, nil) - - bucket.ForceRetrieve(10) - bucket.Refund(20) - assert.True(t, bucket.Retrieve(10)) - assert.Equal(t, int64(0), bucket.RemainingTokens()) + assert.Positive(t, bucket.Retrieve(5)) } From 54c4e491d9f65623f0f673eab72cf1bc06f0b8c2 Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Thu, 11 Jul 2024 15:53:47 -0700 Subject: [PATCH 21/26] Throw 429 and 422 for different failure scenarios Signed-off-by: Justin Jung --- pkg/storegateway/limiter.go | 21 +++++++------ pkg/storegateway/limiter_test.go | 52 +++++++++++++++++++++++--------- pkg/util/token_bucket.go | 4 +++ pkg/util/token_bucket_test.go | 6 ++++ 4 files changed, 59 insertions(+), 24 deletions(-) diff --git a/pkg/storegateway/limiter.go b/pkg/storegateway/limiter.go index b9e77d6322..37bc795d82 100644 --- a/pkg/storegateway/limiter.go +++ b/pkg/storegateway/limiter.go @@ -16,6 +16,8 @@ import ( "github.com/cortexproject/cortex/pkg/util/validation" ) +const tokenBucketBytesLimiterErrStr = "store gateway resource exhausted" + type limiter struct { limiter *store.Limiter } @@ -40,7 +42,7 @@ type compositeLimiter struct { func (c *compositeLimiter) ReserveWithType(num uint64, dataType store.StoreDataType) error { for _, l := range c.limiters { if err := l.ReserveWithType(num, dataType); err != nil { - return httpgrpc.Errorf(http.StatusUnprocessableEntity, err.Error()) + return err // nested limiters are expected to return httpgrpc error } } return nil @@ -74,23 +76,24 @@ func (t *tokenBucketBytesLimiter) ReserveWithType(num uint64, dataType store.Sto return nil } - errMsg := "" + errCode := 0 - if userTokenRemaining < 0 { - errMsg = "not enough tokens in user token bucket" - } else if instanceTokenRemaining < 0 { - errMsg = "not enough tokens in instance token bucket" + fmt.Printf("tokensToRetrieve: %d, maxCapacity: %d", tokensToRetrieve, t.userTokenBucket.MaxCapacity()) + if tokensToRetrieve > t.userTokenBucket.MaxCapacity() || tokensToRetrieve > t.instanceTokenBucket.MaxCapacity() { + errCode = http.StatusUnprocessableEntity + } else if userTokenRemaining < 0 || instanceTokenRemaining < 0 { + errCode = http.StatusTooManyRequests } - if errMsg != "" { + if errCode > 0 { if t.dryRun { - level.Warn(util_log.Logger).Log("msg", errMsg, "dataType", dataType, "dataSize", num, "tokens", tokensToRetrieve) + level.Warn(util_log.Logger).Log("msg", tokenBucketBytesLimiterErrStr, "dataType", dataType, "dataSize", num, "tokens", tokensToRetrieve, "errorCode", errCode) return nil } // We need to protect from the counter being incremented twice due to concurrency t.failedOnce.Do(t.failedCounter.Inc) - return fmt.Errorf(errMsg) + return httpgrpc.Errorf(errCode, tokenBucketBytesLimiterErrStr) } return nil diff --git a/pkg/storegateway/limiter_test.go b/pkg/storegateway/limiter_test.go index c02d4e081e..2244cbc6cb 100644 --- a/pkg/storegateway/limiter_test.go +++ b/pkg/storegateway/limiter_test.go @@ -1,6 +1,7 @@ package storegateway import ( + "strconv" "testing" "github.com/prometheus/client_golang/prometheus" @@ -35,51 +36,69 @@ func TestCompositeLimiter(t *testing.T) { func TestNewTokenBucketBytesLimiter(t *testing.T) { tests := map[string]struct { - tokenToRetrieve uint64 + tokensToRetrieve []uint64 requestTokenBucketSize int64 userTokenBucketSize int64 instanceTokenBucketSize int64 expectedRequestTokenRemaining int64 expectedUserTokenRemaining int64 expectedInstanceTokenRemaining int64 - expectedErrStr string getTokensToRetrieve func(tokens uint64, dataType store.StoreDataType) int64 + errCode int dryRun bool }{ "should retrieve buckets from all buckets": { - tokenToRetrieve: 1, + tokensToRetrieve: []uint64{1}, requestTokenBucketSize: 1, userTokenBucketSize: 1, instanceTokenBucketSize: 1, }, "should succeed if there is enough request token, regardless of user or instance bucket": { - tokenToRetrieve: 1, + tokensToRetrieve: []uint64{1}, requestTokenBucketSize: 1, userTokenBucketSize: 0, instanceTokenBucketSize: 0, expectedUserTokenRemaining: -1, expectedInstanceTokenRemaining: -1, }, - "should fail if not enough user tokens remaining": { - tokenToRetrieve: 2, + "should throw 429 if not enough user tokens remaining": { + tokensToRetrieve: []uint64{1, 1}, requestTokenBucketSize: 1, userTokenBucketSize: 1, instanceTokenBucketSize: 2, - expectedErrStr: "not enough tokens in user token bucket", + errCode: 429, expectedRequestTokenRemaining: -1, expectedUserTokenRemaining: -1, }, - "should fail if not enough instance tokens remaining": { - tokenToRetrieve: 2, + "should throw 422 if request is greater than user token bucket size": { + tokensToRetrieve: []uint64{2}, + requestTokenBucketSize: 1, + userTokenBucketSize: 1, + instanceTokenBucketSize: 2, + errCode: 422, + expectedRequestTokenRemaining: -1, + expectedUserTokenRemaining: -1, + }, + "should throw 429 if not enough instance tokesn remaining": { + tokensToRetrieve: []uint64{1, 1}, requestTokenBucketSize: 1, userTokenBucketSize: 2, instanceTokenBucketSize: 1, - expectedErrStr: "not enough tokens in instance token bucket", + errCode: 429, + expectedRequestTokenRemaining: -1, + expectedInstanceTokenRemaining: -1, + }, + "should throw 422 if request is greater than instance token bucket size": { + tokensToRetrieve: []uint64{2}, + requestTokenBucketSize: 1, + userTokenBucketSize: 2, + instanceTokenBucketSize: 1, + errCode: 422, expectedRequestTokenRemaining: -1, expectedInstanceTokenRemaining: -1, }, "should use getTokensToRetrieve to calculate tokens": { - tokenToRetrieve: 1, + tokensToRetrieve: []uint64{1}, getTokensToRetrieve: func(tokens uint64, dataType store.StoreDataType) int64 { if dataType == store.PostingsFetched { return 0 @@ -88,7 +107,7 @@ func TestNewTokenBucketBytesLimiter(t *testing.T) { }, }, "should not fail if dryRun": { - tokenToRetrieve: 1, + tokensToRetrieve: []uint64{1}, expectedRequestTokenRemaining: -1, expectedUserTokenRemaining: -1, expectedInstanceTokenRemaining: -1, @@ -110,14 +129,17 @@ func TestNewTokenBucketBytesLimiter(t *testing.T) { } l := newTokenBucketBytesLimiter(requestTokenBucket, userTokenBucket, instanceTokenBucket, testData.dryRun, prometheus.NewCounter(prometheus.CounterOpts{}), getTokensToRetrieve) - err := l.ReserveWithType(testData.tokenToRetrieve, store.PostingsFetched) + var err error + for _, token := range testData.tokensToRetrieve { + err = l.ReserveWithType(token, store.PostingsFetched) + } assert.Equal(t, testData.expectedRequestTokenRemaining, requestTokenBucket.Retrieve(0)) assert.Equal(t, testData.expectedUserTokenRemaining, userTokenBucket.Retrieve(0)) assert.Equal(t, testData.expectedInstanceTokenRemaining, instanceTokenBucket.Retrieve(0)) - if testData.expectedErrStr != "" { - assert.ErrorContains(t, err, testData.expectedErrStr) + if testData.errCode > 0 { + assert.ErrorContains(t, err, strconv.Itoa(testData.errCode)) } else { assert.NoError(t, err) } diff --git a/pkg/util/token_bucket.go b/pkg/util/token_bucket.go index 3c3fdcbb62..6e33c7aaf8 100644 --- a/pkg/util/token_bucket.go +++ b/pkg/util/token_bucket.go @@ -46,6 +46,10 @@ func (t *TokenBucket) Retrieve(amount int64) int64 { return t.remainingTokens } +func (t *TokenBucket) MaxCapacity() int64 { + return t.maxCapacity +} + func (t *TokenBucket) updateTokens() { now := time.Now() refilledTokens := int64(now.Sub(t.lastRefill).Seconds() * float64(t.refillRate)) diff --git a/pkg/util/token_bucket_test.go b/pkg/util/token_bucket_test.go index a5a0bd932d..155e36ce75 100644 --- a/pkg/util/token_bucket_test.go +++ b/pkg/util/token_bucket_test.go @@ -15,3 +15,9 @@ func TestTokenBucket_Retrieve(t *testing.T) { time.Sleep(time.Second) assert.Positive(t, bucket.Retrieve(5)) } + +func TestTokenBucket_MaxCapacity(t *testing.T) { + bucket := NewTokenBucket(10, nil) + + assert.Equal(t, int64(10), bucket.MaxCapacity()) +} From 26dc6f4fd18ac02ff8bdba40e6301863caab8beb Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Thu, 11 Jul 2024 21:16:08 -0700 Subject: [PATCH 22/26] Hide token factors from doc Signed-off-by: Justin Jung --- docs/blocks-storage/querier.md | 24 --------------------- docs/blocks-storage/store-gateway.md | 24 --------------------- docs/configuration/config-file-reference.md | 24 --------------------- pkg/storage/tsdb/config.go | 12 +++++------ 4 files changed, 6 insertions(+), 78 deletions(-) diff --git a/docs/blocks-storage/querier.md b/docs/blocks-storage/querier.md index f98f402447..cbfbfcbb69 100644 --- a/docs/blocks-storage/querier.md +++ b/docs/blocks-storage/querier.md @@ -1362,30 +1362,6 @@ blocks_storage: # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.request-token-bucket-size [request_token_bucket_size: | default = 4194304] - # Multiplication factor used for fetched postings token - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-postings-token-factor - [fetched_postings_token_factor: | default = 0] - - # Multiplication factor used for touched postings token - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-postings-token-factor - [touched_postings_token_factor: | default = 5] - - # Multiplication factor used for fetched series token - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-series-token-factor - [fetched_series_token_factor: | default = 0] - - # Multiplication factor used for touched series token - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-series-token-factor - [touched_series_token_factor: | default = 25] - - # Multiplication factor used for fetched chunks token - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-chunks-token-factor - [fetched_chunks_token_factor: | default = 0] - - # Multiplication factor used for touched chunks token - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-chunks-token-factor - [touched_chunks_token_factor: | default = 1] - tsdb: # Local directory to store TSDBs in the ingesters. # CLI flag: -blocks-storage.tsdb.dir diff --git a/docs/blocks-storage/store-gateway.md b/docs/blocks-storage/store-gateway.md index 24232bf966..cf3406a22e 100644 --- a/docs/blocks-storage/store-gateway.md +++ b/docs/blocks-storage/store-gateway.md @@ -1487,30 +1487,6 @@ blocks_storage: # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.request-token-bucket-size [request_token_bucket_size: | default = 4194304] - # Multiplication factor used for fetched postings token - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-postings-token-factor - [fetched_postings_token_factor: | default = 0] - - # Multiplication factor used for touched postings token - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-postings-token-factor - [touched_postings_token_factor: | default = 5] - - # Multiplication factor used for fetched series token - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-series-token-factor - [fetched_series_token_factor: | default = 0] - - # Multiplication factor used for touched series token - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-series-token-factor - [touched_series_token_factor: | default = 25] - - # Multiplication factor used for fetched chunks token - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-chunks-token-factor - [fetched_chunks_token_factor: | default = 0] - - # Multiplication factor used for touched chunks token - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-chunks-token-factor - [touched_chunks_token_factor: | default = 1] - tsdb: # Local directory to store TSDBs in the ingesters. # CLI flag: -blocks-storage.tsdb.dir diff --git a/docs/configuration/config-file-reference.md b/docs/configuration/config-file-reference.md index bc47adb96e..9ae350b2cb 100644 --- a/docs/configuration/config-file-reference.md +++ b/docs/configuration/config-file-reference.md @@ -1920,30 +1920,6 @@ bucket_store: # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.request-token-bucket-size [request_token_bucket_size: | default = 4194304] - # Multiplication factor used for fetched postings token - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-postings-token-factor - [fetched_postings_token_factor: | default = 0] - - # Multiplication factor used for touched postings token - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-postings-token-factor - [touched_postings_token_factor: | default = 5] - - # Multiplication factor used for fetched series token - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-series-token-factor - [fetched_series_token_factor: | default = 0] - - # Multiplication factor used for touched series token - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-series-token-factor - [touched_series_token_factor: | default = 25] - - # Multiplication factor used for fetched chunks token - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.fetched-chunks-token-factor - [fetched_chunks_token_factor: | default = 0] - - # Multiplication factor used for touched chunks token - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.touched-chunks-token-factor - [touched_chunks_token_factor: | default = 1] - tsdb: # Local directory to store TSDBs in the ingesters. # CLI flag: -blocks-storage.tsdb.dir diff --git a/pkg/storage/tsdb/config.go b/pkg/storage/tsdb/config.go index 6a7720fae1..4e69f17f2b 100644 --- a/pkg/storage/tsdb/config.go +++ b/pkg/storage/tsdb/config.go @@ -303,12 +303,12 @@ type TokenBucketBytesLimiterConfig struct { InstanceTokenBucketSize int64 `yaml:"instance_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"` + FetchedPostingsTokenFactor float64 `yaml:"fetched_postings_token_factor" doc:"hidden"` + TouchedPostingsTokenFactor float64 `yaml:"touched_postings_token_factor" doc:"hidden"` + FetchedSeriesTokenFactor float64 `yaml:"fetched_series_token_factor" doc:"hidden"` + TouchedSeriesTokenFactor float64 `yaml:"touched_series_token_factor" doc:"hidden"` + FetchedChunksTokenFactor float64 `yaml:"fetched_chunks_token_factor" doc:"hidden"` + TouchedChunksTokenFactor float64 `yaml:"touched_chunks_token_factor" doc:"hidden"` } // RegisterFlags registers the BucketStore flags From 78a856eced7af70628b56ece9d586e91e88478c2 Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Thu, 11 Jul 2024 22:39:18 -0700 Subject: [PATCH 23/26] Simplified config by combining dryrun and enabled Signed-off-by: Justin Jung --- docs/blocks-storage/querier.md | 11 ++++----- docs/blocks-storage/store-gateway.md | 11 ++++----- docs/configuration/config-file-reference.md | 11 ++++----- pkg/storage/tsdb/config.go | 25 +++++++++++++++++---- pkg/storegateway/bucket_stores.go | 6 ++--- pkg/storegateway/bucket_stores_test.go | 4 ++-- pkg/storegateway/limiter.go | 5 +++-- 7 files changed, 41 insertions(+), 32 deletions(-) diff --git a/docs/blocks-storage/querier.md b/docs/blocks-storage/querier.md index cbfbfcbb69..a9206c82ba 100644 --- a/docs/blocks-storage/querier.md +++ b/docs/blocks-storage/querier.md @@ -1342,13 +1342,10 @@ blocks_storage: [series_batch_size: | default = 10000] token_bucket_bytes_limiter: - # Whether token bucket limiter is enabled - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.enabled - [enabled: | default = false] - - # Whether the token bucket limiter is in dry run mode - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.dry-run - [dry_run: | default = false] + # Token bucket bytes limiter mode. Supported values are: disabled, dryrun, + # enabled + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.mode + [mode: | default = "disabled"] # Instance token bucket size # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.instance-token-bucket-size diff --git a/docs/blocks-storage/store-gateway.md b/docs/blocks-storage/store-gateway.md index cf3406a22e..6c31046b1f 100644 --- a/docs/blocks-storage/store-gateway.md +++ b/docs/blocks-storage/store-gateway.md @@ -1467,13 +1467,10 @@ blocks_storage: [series_batch_size: | default = 10000] token_bucket_bytes_limiter: - # Whether token bucket limiter is enabled - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.enabled - [enabled: | default = false] - - # Whether the token bucket limiter is in dry run mode - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.dry-run - [dry_run: | default = false] + # Token bucket bytes limiter mode. Supported values are: disabled, dryrun, + # enabled + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.mode + [mode: | default = "disabled"] # Instance token bucket size # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.instance-token-bucket-size diff --git a/docs/configuration/config-file-reference.md b/docs/configuration/config-file-reference.md index 9ae350b2cb..167b50592e 100644 --- a/docs/configuration/config-file-reference.md +++ b/docs/configuration/config-file-reference.md @@ -1900,13 +1900,10 @@ bucket_store: [series_batch_size: | default = 10000] token_bucket_bytes_limiter: - # Whether token bucket limiter is enabled - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.enabled - [enabled: | default = false] - - # Whether the token bucket limiter is in dry run mode - # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.dry-run - [dry_run: | default = false] + # Token bucket bytes limiter mode. Supported values are: disabled, dryrun, + # enabled + # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.mode + [mode: | default = "disabled"] # Instance token bucket size # CLI flag: -blocks-storage.bucket-store.token-bucket-bytes-limiter.instance-token-bucket-size diff --git a/pkg/storage/tsdb/config.go b/pkg/storage/tsdb/config.go index 4e69f17f2b..bd3099dba9 100644 --- a/pkg/storage/tsdb/config.go +++ b/pkg/storage/tsdb/config.go @@ -2,6 +2,7 @@ package tsdb import ( "flag" + "fmt" "path/filepath" "strings" "time" @@ -52,6 +53,7 @@ var ( ErrInvalidBucketIndexBlockDiscoveryStrategy = errors.New("bucket index block discovery strategy can only be enabled when bucket index is enabled") ErrBlockDiscoveryStrategy = errors.New("invalid block discovery strategy") + ErrInvalidTokenBucketBytesLimiterMode = errors.New("invalid token bucket bytes limiter mode") ) // BlocksStorageConfig holds the config information for the blocks storage. @@ -298,8 +300,7 @@ type BucketStoreConfig struct { } type TokenBucketBytesLimiterConfig struct { - Enabled bool `yaml:"enabled"` - DryRun bool `yaml:"dry_run"` + Mode string `yaml:"mode"` InstanceTokenBucketSize int64 `yaml:"instance_token_bucket_size"` UserTokenBucketSize int64 `yaml:"user_token_bucket_size"` RequestTokenBucketSize int64 `yaml:"request_token_bucket_size"` @@ -342,8 +343,7 @@ 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.TokenBucketBytesLimiter.Enabled, "blocks-storage.bucket-store.token-bucket-bytes-limiter.enabled", false, "Whether token bucket limiter is enabled") - f.BoolVar(&cfg.TokenBucketBytesLimiter.DryRun, "blocks-storage.bucket-store.token-bucket-bytes-limiter.dry-run", false, "Whether the token bucket limiter is in dry run mode") + f.StringVar(&cfg.TokenBucketBytesLimiter.Mode, "blocks-storage.bucket-store.token-bucket-bytes-limiter.mode", string(TokenBucketBytesLimiterDisabled), fmt.Sprintf("Token bucket bytes limiter mode. Supported values are: %s", strings.Join(supportedTokenBucketBytesLimiterModes, ", "))) f.Int64Var(&cfg.TokenBucketBytesLimiter.InstanceTokenBucketSize, "blocks-storage.bucket-store.token-bucket-bytes-limiter.instance-token-bucket-size", int64(820*units.Mebibyte), "Instance token bucket size") f.Int64Var(&cfg.TokenBucketBytesLimiter.UserTokenBucketSize, "blocks-storage.bucket-store.token-bucket-bytes-limiter.user-token-bucket-size", int64(615*units.Mebibyte), "User token bucket size") f.Int64Var(&cfg.TokenBucketBytesLimiter.RequestTokenBucketSize, "blocks-storage.bucket-store.token-bucket-bytes-limiter.request-token-bucket-size", int64(4*units.Mebibyte), "Request token bucket size") @@ -372,6 +372,9 @@ func (cfg *BucketStoreConfig) Validate() error { if !util.StringsContain(supportedBlockDiscoveryStrategies, cfg.BlockDiscoveryStrategy) { return ErrInvalidBucketIndexBlockDiscoveryStrategy } + if !util.StringsContain(supportedTokenBucketBytesLimiterModes, cfg.TokenBucketBytesLimiter.Mode) { + return ErrInvalidTokenBucketBytesLimiterMode + } return nil } @@ -403,3 +406,17 @@ var supportedBlockDiscoveryStrategies = []string{ string(RecursiveDiscovery), string(BucketIndexDiscovery), } + +type TokenBucketBytesLimiterMode string + +const ( + TokenBucketBytesLimiterDisabled TokenBucketBytesLimiterMode = "disabled" + TokenBucketBytesLimiterDryRun TokenBucketBytesLimiterMode = "dryrun" + TokenBucketBytesLimiterEnabled TokenBucketBytesLimiterMode = "enabled" +) + +var supportedTokenBucketBytesLimiterModes = []string{ + string(TokenBucketBytesLimiterDisabled), + string(TokenBucketBytesLimiterDryRun), + string(TokenBucketBytesLimiterEnabled), +} diff --git a/pkg/storegateway/bucket_stores.go b/pkg/storegateway/bucket_stores.go index ebc409cc29..50333a19ce 100644 --- a/pkg/storegateway/bucket_stores.go +++ b/pkg/storegateway/bucket_stores.go @@ -150,7 +150,7 @@ func NewBucketStores(cfg tsdb.BlocksStorageConfig, shardingStrategy ShardingStra return nil, errors.Wrap(err, "create chunks bytes pool") } - if u.cfg.BucketStore.TokenBucketBytesLimiter.Enabled { + if u.cfg.BucketStore.TokenBucketBytesLimiter.Mode == string(tsdb.TokenBucketBytesLimiterDryRun) { u.instanceTokenBucket = util.NewTokenBucket(cfg.BucketStore.TokenBucketBytesLimiter.InstanceTokenBucketSize, promauto.With(reg).NewGauge(prometheus.GaugeOpts{ Name: "cortex_bucket_stores_instance_token_bucket_remaining", Help: "Number of tokens left in instance token bucket.", @@ -488,7 +488,7 @@ func (u *BucketStores) closeEmptyBucketStore(userID string) error { unlockInDefer = false u.storesMu.Unlock() - if u.cfg.BucketStore.TokenBucketBytesLimiter.Enabled { + if u.cfg.BucketStore.TokenBucketBytesLimiter.Mode == string(tsdb.TokenBucketBytesLimiterEnabled) { u.userTokenBucketsMu.Lock() delete(u.userTokenBuckets, userID) u.userTokenBucketsMu.Unlock() @@ -631,7 +631,7 @@ func (u *BucketStores) getOrCreateStore(userID string) (*store.BucketStore, erro bucketStoreOpts = append(bucketStoreOpts, store.WithDebugLogging()) } - if u.cfg.BucketStore.TokenBucketBytesLimiter.Enabled { + if u.cfg.BucketStore.TokenBucketBytesLimiter.Mode == string(tsdb.TokenBucketBytesLimiterEnabled) { u.userTokenBucketsMu.Lock() u.userTokenBuckets[userID] = util.NewTokenBucket(u.cfg.BucketStore.TokenBucketBytesLimiter.UserTokenBucketSize, nil) u.userTokenBucketsMu.Unlock() diff --git a/pkg/storegateway/bucket_stores_test.go b/pkg/storegateway/bucket_stores_test.go index e24581c06e..cba50c106a 100644 --- a/pkg/storegateway/bucket_stores_test.go +++ b/pkg/storegateway/bucket_stores_test.go @@ -771,7 +771,7 @@ func TestBucketStores_tokenBuckets(t *testing.T) { ctx := context.Background() cfg := prepareStorageConfig(t) - cfg.BucketStore.TokenBucketBytesLimiter.Enabled = true + cfg.BucketStore.TokenBucketBytesLimiter.Mode = string(cortex_tsdb.TokenBucketBytesLimiterEnabled) storageDir := t.TempDir() userToMetric := map[string]string{ @@ -807,7 +807,7 @@ func TestBucketStores_tokenBuckets(t *testing.T) { assert.Nil(t, stores.getUserTokenBucket("user-1")) assert.Nil(t, stores.getUserTokenBucket("user-2")) - cfg.BucketStore.TokenBucketBytesLimiter.Enabled = false + cfg.BucketStore.TokenBucketBytesLimiter.Mode = string(cortex_tsdb.TokenBucketBytesLimiterDisabled) sharding.users = []string{user1, user2} reg = prometheus.NewPedanticRegistry() stores, err = NewBucketStores(cfg, &sharding, objstore.WithNoopInstr(bucket), defaultLimitsOverrides(t), mockLoggingLevel(), log.NewNopLogger(), reg) diff --git a/pkg/storegateway/limiter.go b/pkg/storegateway/limiter.go index 37bc795d82..17dd4365ca 100644 --- a/pkg/storegateway/limiter.go +++ b/pkg/storegateway/limiter.go @@ -137,9 +137,10 @@ func newBytesLimiterFactory(limits *validation.Overrides, userID string, userTok // each time a new limiter is instantiated. limiters = append(limiters, store.NewLimiter(uint64(limits.MaxDownloadedBytesPerRequest(userID)), failedCounter)) - if tokenBucketBytesLimiterCfg.Enabled { + if tokenBucketBytesLimiterCfg.Mode == string(tsdb.TokenBucketBytesLimiterEnabled) { requestTokenBucket := util.NewTokenBucket(tokenBucketBytesLimiterCfg.RequestTokenBucketSize, nil) - limiters = append(limiters, newTokenBucketBytesLimiter(requestTokenBucket, userTokenBucket, instanceTokenBucket, tokenBucketBytesLimiterCfg.DryRun, failedCounter, getTokensToRetrieve)) + dryRun := tokenBucketBytesLimiterCfg.Mode == string(tsdb.TokenBucketBytesLimiterDryRun) + limiters = append(limiters, newTokenBucketBytesLimiter(requestTokenBucket, userTokenBucket, instanceTokenBucket, dryRun, failedCounter, getTokensToRetrieve)) } return &compositeLimiter{ From c7f61915085f9707f872c231e2ae798dfa810a75 Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Thu, 11 Jul 2024 22:50:23 -0700 Subject: [PATCH 24/26] Remove test log Signed-off-by: Justin Jung --- pkg/storegateway/limiter.go | 2 -- 1 file changed, 2 deletions(-) diff --git a/pkg/storegateway/limiter.go b/pkg/storegateway/limiter.go index 17dd4365ca..44c558a579 100644 --- a/pkg/storegateway/limiter.go +++ b/pkg/storegateway/limiter.go @@ -1,7 +1,6 @@ package storegateway import ( - "fmt" "net/http" "sync" @@ -78,7 +77,6 @@ func (t *tokenBucketBytesLimiter) ReserveWithType(num uint64, dataType store.Sto errCode := 0 - fmt.Printf("tokensToRetrieve: %d, maxCapacity: %d", tokensToRetrieve, t.userTokenBucket.MaxCapacity()) if tokensToRetrieve > t.userTokenBucket.MaxCapacity() || tokensToRetrieve > t.instanceTokenBucket.MaxCapacity() { errCode = http.StatusUnprocessableEntity } else if userTokenRemaining < 0 || instanceTokenRemaining < 0 { From faba9b89a27861cc3b1aece08347b1d2804c8476 Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Thu, 11 Jul 2024 23:33:37 -0700 Subject: [PATCH 25/26] Fix tests Signed-off-by: Justin Jung --- pkg/storegateway/bucket_stores.go | 2 +- pkg/storegateway/limiter.go | 8 ++++---- pkg/storegateway/limiter_test.go | 10 +++++----- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/pkg/storegateway/bucket_stores.go b/pkg/storegateway/bucket_stores.go index 50333a19ce..7ce3c2b7ec 100644 --- a/pkg/storegateway/bucket_stores.go +++ b/pkg/storegateway/bucket_stores.go @@ -150,7 +150,7 @@ func NewBucketStores(cfg tsdb.BlocksStorageConfig, shardingStrategy ShardingStra return nil, errors.Wrap(err, "create chunks bytes pool") } - if u.cfg.BucketStore.TokenBucketBytesLimiter.Mode == string(tsdb.TokenBucketBytesLimiterDryRun) { + if u.cfg.BucketStore.TokenBucketBytesLimiter.Mode != string(tsdb.TokenBucketBytesLimiterDisabled) { u.instanceTokenBucket = util.NewTokenBucket(cfg.BucketStore.TokenBucketBytesLimiter.InstanceTokenBucketSize, promauto.With(reg).NewGauge(prometheus.GaugeOpts{ Name: "cortex_bucket_stores_instance_token_bucket_remaining", Help: "Number of tokens left in instance token bucket.", diff --git a/pkg/storegateway/limiter.go b/pkg/storegateway/limiter.go index 44c558a579..a34911d5aa 100644 --- a/pkg/storegateway/limiter.go +++ b/pkg/storegateway/limiter.go @@ -34,11 +34,11 @@ func (c *limiter) ReserveWithType(num uint64, _ store.StoreDataType) error { return nil } -type compositeLimiter struct { +type compositeBytesLimiter struct { limiters []store.BytesLimiter } -func (c *compositeLimiter) ReserveWithType(num uint64, dataType store.StoreDataType) error { +func (c *compositeBytesLimiter) ReserveWithType(num uint64, dataType store.StoreDataType) error { for _, l := range c.limiters { if err := l.ReserveWithType(num, dataType); err != nil { return err // nested limiters are expected to return httpgrpc error @@ -133,7 +133,7 @@ func newBytesLimiterFactory(limits *validation.Overrides, userID string, userTok limiters := []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. - limiters = append(limiters, store.NewLimiter(uint64(limits.MaxDownloadedBytesPerRequest(userID)), failedCounter)) + limiters = append(limiters, &limiter{limiter: store.NewLimiter(uint64(limits.MaxDownloadedBytesPerRequest(userID)), failedCounter)}) if tokenBucketBytesLimiterCfg.Mode == string(tsdb.TokenBucketBytesLimiterEnabled) { requestTokenBucket := util.NewTokenBucket(tokenBucketBytesLimiterCfg.RequestTokenBucketSize, nil) @@ -141,7 +141,7 @@ func newBytesLimiterFactory(limits *validation.Overrides, userID string, userTok limiters = append(limiters, newTokenBucketBytesLimiter(requestTokenBucket, userTokenBucket, instanceTokenBucket, dryRun, failedCounter, getTokensToRetrieve)) } - return &compositeLimiter{ + return &compositeBytesLimiter{ limiters: limiters, } } diff --git a/pkg/storegateway/limiter_test.go b/pkg/storegateway/limiter_test.go index 2244cbc6cb..c2a61de249 100644 --- a/pkg/storegateway/limiter_test.go +++ b/pkg/storegateway/limiter_test.go @@ -1,7 +1,7 @@ package storegateway import ( - "strconv" + "fmt" "testing" "github.com/prometheus/client_golang/prometheus" @@ -23,10 +23,10 @@ func TestLimiter(t *testing.T) { } func TestCompositeLimiter(t *testing.T) { - l := &compositeLimiter{ + l := &compositeBytesLimiter{ limiters: []store.BytesLimiter{ - store.NewLimiter(2, prometheus.NewCounter(prometheus.CounterOpts{})), - store.NewLimiter(1, prometheus.NewCounter(prometheus.CounterOpts{})), + &limiter{limiter: store.NewLimiter(2, prometheus.NewCounter(prometheus.CounterOpts{}))}, + &limiter{limiter: store.NewLimiter(1, prometheus.NewCounter(prometheus.CounterOpts{}))}, }, } @@ -139,7 +139,7 @@ func TestNewTokenBucketBytesLimiter(t *testing.T) { assert.Equal(t, testData.expectedInstanceTokenRemaining, instanceTokenBucket.Retrieve(0)) if testData.errCode > 0 { - assert.ErrorContains(t, err, strconv.Itoa(testData.errCode)) + assert.ErrorContains(t, err, fmt.Sprintf("(%d)", testData.errCode)) } else { assert.NoError(t, err) } From 172a62d79126bbec2c4c4347cab8197bfa188c61 Mon Sep 17 00:00:00 2001 From: Justin Jung Date: Fri, 12 Jul 2024 00:08:14 -0700 Subject: [PATCH 26/26] Fix Signed-off-by: Justin Jung --- pkg/storegateway/bucket_stores.go | 4 ++-- pkg/storegateway/bucket_stores_test.go | 11 +++++++++++ pkg/storegateway/limiter.go | 2 +- 3 files changed, 14 insertions(+), 3 deletions(-) diff --git a/pkg/storegateway/bucket_stores.go b/pkg/storegateway/bucket_stores.go index 7ce3c2b7ec..e95e7dd6bc 100644 --- a/pkg/storegateway/bucket_stores.go +++ b/pkg/storegateway/bucket_stores.go @@ -488,7 +488,7 @@ func (u *BucketStores) closeEmptyBucketStore(userID string) error { unlockInDefer = false u.storesMu.Unlock() - if u.cfg.BucketStore.TokenBucketBytesLimiter.Mode == string(tsdb.TokenBucketBytesLimiterEnabled) { + if u.cfg.BucketStore.TokenBucketBytesLimiter.Mode != string(tsdb.TokenBucketBytesLimiterDisabled) { u.userTokenBucketsMu.Lock() delete(u.userTokenBuckets, userID) u.userTokenBucketsMu.Unlock() @@ -631,7 +631,7 @@ func (u *BucketStores) getOrCreateStore(userID string) (*store.BucketStore, erro bucketStoreOpts = append(bucketStoreOpts, store.WithDebugLogging()) } - if u.cfg.BucketStore.TokenBucketBytesLimiter.Mode == string(tsdb.TokenBucketBytesLimiterEnabled) { + if u.cfg.BucketStore.TokenBucketBytesLimiter.Mode != string(tsdb.TokenBucketBytesLimiterDisabled) { u.userTokenBucketsMu.Lock() u.userTokenBuckets[userID] = util.NewTokenBucket(u.cfg.BucketStore.TokenBucketBytesLimiter.UserTokenBucketSize, nil) u.userTokenBucketsMu.Unlock() diff --git a/pkg/storegateway/bucket_stores_test.go b/pkg/storegateway/bucket_stores_test.go index cba50c106a..d3efcfc112 100644 --- a/pkg/storegateway/bucket_stores_test.go +++ b/pkg/storegateway/bucket_stores_test.go @@ -807,6 +807,17 @@ func TestBucketStores_tokenBuckets(t *testing.T) { assert.Nil(t, stores.getUserTokenBucket("user-1")) assert.Nil(t, stores.getUserTokenBucket("user-2")) + cfg.BucketStore.TokenBucketBytesLimiter.Mode = string(cortex_tsdb.TokenBucketBytesLimiterDryRun) + sharding.users = []string{user1, user2} + reg = prometheus.NewPedanticRegistry() + stores, err = NewBucketStores(cfg, &sharding, objstore.WithNoopInstr(bucket), defaultLimitsOverrides(t), mockLoggingLevel(), log.NewNopLogger(), reg) + assert.NoError(t, err) + assert.NotNil(t, stores.instanceTokenBucket) + + assert.NoError(t, stores.InitialSync(ctx)) + assert.NotNil(t, stores.getUserTokenBucket("user-1")) + assert.NotNil(t, stores.getUserTokenBucket("user-2")) + cfg.BucketStore.TokenBucketBytesLimiter.Mode = string(cortex_tsdb.TokenBucketBytesLimiterDisabled) sharding.users = []string{user1, user2} reg = prometheus.NewPedanticRegistry() diff --git a/pkg/storegateway/limiter.go b/pkg/storegateway/limiter.go index a34911d5aa..7a633e2422 100644 --- a/pkg/storegateway/limiter.go +++ b/pkg/storegateway/limiter.go @@ -135,7 +135,7 @@ func newBytesLimiterFactory(limits *validation.Overrides, userID string, userTok // each time a new limiter is instantiated. limiters = append(limiters, &limiter{limiter: store.NewLimiter(uint64(limits.MaxDownloadedBytesPerRequest(userID)), failedCounter)}) - if tokenBucketBytesLimiterCfg.Mode == string(tsdb.TokenBucketBytesLimiterEnabled) { + if tokenBucketBytesLimiterCfg.Mode != string(tsdb.TokenBucketBytesLimiterDisabled) { requestTokenBucket := util.NewTokenBucket(tokenBucketBytesLimiterCfg.RequestTokenBucketSize, nil) dryRun := tokenBucketBytesLimiterCfg.Mode == string(tsdb.TokenBucketBytesLimiterDryRun) limiters = append(limiters, newTokenBucketBytesLimiter(requestTokenBucket, userTokenBucket, instanceTokenBucket, dryRun, failedCounter, getTokensToRetrieve))