-
Notifications
You must be signed in to change notification settings - Fork 186
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Signed-off-by: Cody Littley <[email protected]>
- Loading branch information
1 parent
fbe6d3b
commit afd5894
Showing
12 changed files
with
1,224 additions
and
54 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,102 @@ | ||
| package limiter | ||
|
|
||
| import ( | ||
| "fmt" | ||
| "golang.org/x/time/rate" | ||
| "sync" | ||
| "time" | ||
| ) | ||
|
|
||
| // BlobRateLimiter enforces rate limits on GetBlob operations. | ||
| type BlobRateLimiter struct { | ||
|
|
||
| // config is the rate limit configuration. | ||
| config *Config | ||
|
|
||
| // opLimiter enforces rate limits on the maximum rate of GetBlob operations | ||
| opLimiter *rate.Limiter | ||
|
|
||
| // bandwidthLimiter enforces rate limits on the maximum bandwidth consumed by GetBlob operations. Only the size | ||
| // of the blob data is considered, not the size of the entire response. | ||
| bandwidthLimiter *rate.Limiter | ||
|
|
||
| // operationsInFlight is the number of GetBlob operations currently in flight. | ||
| operationsInFlight int | ||
|
|
||
| // this lock is used to provide thread safety | ||
| lock sync.Mutex | ||
| } | ||
|
|
||
| // NewBlobRateLimiter creates a new BlobRateLimiter. | ||
| func NewBlobRateLimiter(config *Config) *BlobRateLimiter { | ||
| globalGetBlobOpLimiter := rate.NewLimiter( | ||
| rate.Limit(config.MaxGetBlobOpsPerSecond), | ||
| config.GetBlobOpsBurstiness) | ||
|
|
||
| globalGetBlobBandwidthLimiter := rate.NewLimiter( | ||
| rate.Limit(config.MaxGetBlobBytesPerSecond), | ||
| config.GetBlobBytesBurstiness) | ||
|
|
||
| return &BlobRateLimiter{ | ||
| config: config, | ||
| opLimiter: globalGetBlobOpLimiter, | ||
| bandwidthLimiter: globalGetBlobBandwidthLimiter, | ||
| } | ||
| } | ||
|
|
||
| // BeginGetBlobOperation should be called when a GetBlob operation is about to begin. If it returns an error, | ||
| // the operation should not be performed. If it does not return an error, FinishGetBlobOperation should be | ||
| // called when the operation completes. | ||
| func (l *BlobRateLimiter) BeginGetBlobOperation(now time.Time) error { | ||
| if l == nil { | ||
| // If the rate limiter is nil, do not enforce rate limits. | ||
| return nil | ||
| } | ||
|
|
||
| l.lock.Lock() | ||
| defer l.lock.Unlock() | ||
|
|
||
| if l.operationsInFlight >= l.config.MaxConcurrentGetBlobOps { | ||
| return fmt.Errorf("global concurrent request limit exceeded for getBlob operations, try again later") | ||
| } | ||
| if l.opLimiter.TokensAt(now) < 1 { | ||
| return fmt.Errorf("global rate limit exceeded for getBlob operations, try again later") | ||
| } | ||
|
|
||
| l.operationsInFlight++ | ||
| l.opLimiter.AllowN(now, 1) | ||
|
|
||
| return nil | ||
| } | ||
|
|
||
| // FinishGetBlobOperation should be called exactly once for each time BeginGetBlobOperation is called and | ||
| // returns nil. | ||
| func (l *BlobRateLimiter) FinishGetBlobOperation() { | ||
| if l == nil { | ||
| // If the rate limiter is nil, do not enforce rate limits. | ||
| return | ||
| } | ||
|
|
||
| l.lock.Lock() | ||
| defer l.lock.Unlock() | ||
|
|
||
| l.operationsInFlight-- | ||
| } | ||
|
|
||
| // RequestGetBlobBandwidth should be called when a GetBlob is about to start downloading blob data | ||
| // from S3. It returns an error if there is insufficient bandwidth available. If it returns nil, the | ||
| // operation should proceed. | ||
| func (l *BlobRateLimiter) RequestGetBlobBandwidth(now time.Time, bytes uint32) error { | ||
| if l == nil { | ||
| // If the rate limiter is nil, do not enforce rate limits. | ||
| return nil | ||
| } | ||
|
|
||
| // no locking needed, the only thing we touch here is the bandwidthLimiter, which is inherently thread-safe | ||
|
|
||
| allowed := l.bandwidthLimiter.AllowN(now, int(bytes)) | ||
| if !allowed { | ||
| return fmt.Errorf("global rate limit exceeded for getBlob bandwidth, try again later") | ||
| } | ||
| return nil | ||
| } |
Oops, something went wrong.