code.dwrz.net

Go monorepo.
Log | Files | Refs

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 }