src

Go monorepo.
git clone git://code.dwrz.net/src
Log | Files | Refs

token_rate_limit.go (2076B)

      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 type canceledError struct {
     34 	Err error
     35 }
     36 
     37 func (c canceledError) CanceledError() bool { return true }
     38 func (c canceledError) Unwrap() error       { return c.Err }
     39 func (c canceledError) Error() string {
     40 	return fmt.Sprintf("canceled, %v", c.Err)
     41 }
     42 
     43 // GetToken may cause a available pool of retry quota to be
     44 // decremented. Will return an error if the decremented value can not be
     45 // reduced from the retry quota.
     46 func (l *TokenRateLimit) GetToken(ctx context.Context, cost uint) (func() error, error) {
     47 	select {
     48 	case <-ctx.Done():
     49 		return nil, canceledError{Err: ctx.Err()}
     50 	default:
     51 	}
     52 	if avail, ok := l.bucket.Retrieve(cost); !ok {
     53 		return nil, QuotaExceededError{Available: avail, Requested: cost}
     54 	}
     55 
     56 	return rateToken{
     57 		tokenCost: cost,
     58 		bucket:    l.bucket,
     59 	}.release, nil
     60 }
     61 
     62 // AddTokens increments the token bucket by a fixed amount.
     63 func (l *TokenRateLimit) AddTokens(v uint) error {
     64 	l.bucket.Refund(v)
     65 	return nil
     66 }
     67 
     68 // Remaining returns the number of remaining tokens in the bucket.
     69 func (l *TokenRateLimit) Remaining() uint {
     70 	return l.bucket.Remaining()
     71 }
     72 
     73 // QuotaExceededError provides the SDK error when the retries for a given
     74 // token bucket have been exhausted.
     75 type QuotaExceededError struct {
     76 	Available uint
     77 	Requested uint
     78 }
     79 
     80 func (e QuotaExceededError) Error() string {
     81 	return fmt.Sprintf("retry quota exceeded, %d available, %d requested",
     82 		e.Available, e.Requested)
     83 }