token_rate_limit.go (2127B)
1 package ratelimit 2 3 import ( 4 "context" 5 "fmt" 6 ) 7 8 type rateToken struct { 9 tokenCost uint 10 bucket *TokenBucket 11 } 12 13 func (t rateToken) release() error { 14 t.bucket.Refund(t.tokenCost) 15 return nil 16 } 17 18 // TokenRateLimit provides a Token Bucket RateLimiter implementation 19 // that limits the overall number of retry attempts that can be made across 20 // operation invocations. 21 type TokenRateLimit struct { 22 bucket *TokenBucket 23 } 24 25 // NewTokenRateLimit returns an TokenRateLimit with default values. 26 // Functional options can configure the retry rate limiter. 27 func NewTokenRateLimit(tokens uint) *TokenRateLimit { 28 return &TokenRateLimit{ 29 bucket: NewTokenBucket(tokens), 30 } 31 } 32 33 func isTimeoutError(error) bool { 34 return false 35 } 36 37 type canceledError struct { 38 Err error 39 } 40 41 func (c canceledError) CanceledError() bool { return true } 42 func (c canceledError) Unwrap() error { return c.Err } 43 func (c canceledError) Error() string { 44 return fmt.Sprintf("canceled, %v", c.Err) 45 } 46 47 // GetToken may cause a available pool of retry quota to be 48 // decremented. Will return an error if the decremented value can not be 49 // reduced from the retry quota. 50 func (l *TokenRateLimit) GetToken(ctx context.Context, cost uint) (func() error, error) { 51 select { 52 case <-ctx.Done(): 53 return nil, canceledError{Err: ctx.Err()} 54 default: 55 } 56 if avail, ok := l.bucket.Retrieve(cost); !ok { 57 return nil, QuotaExceededError{Available: avail, Requested: cost} 58 } 59 60 return rateToken{ 61 tokenCost: cost, 62 bucket: l.bucket, 63 }.release, nil 64 } 65 66 // AddTokens increments the token bucket by a fixed amount. 67 func (l *TokenRateLimit) AddTokens(v uint) error { 68 l.bucket.Refund(v) 69 return nil 70 } 71 72 // Remaining returns the number of remaining tokens in the bucket. 73 func (l *TokenRateLimit) Remaining() uint { 74 return l.bucket.Remaining() 75 } 76 77 // QuotaExceededError provides the SDK error when the retries for a given 78 // token bucket have been exhausted. 79 type QuotaExceededError struct { 80 Available uint 81 Requested uint 82 } 83 84 func (e QuotaExceededError) Error() string { 85 return fmt.Sprintf("retry quota exceeded, %d available, %d requested", 86 e.Available, e.Requested) 87 }