code.dwrz.net

Go monorepo.
Log | Files | Refs

adaptive_token_bucket.go (2383B)

      1 package retry
      2 
      3 import (
      4 	"math"
      5 	"sync"
      6 )
      7 
      8 // adaptiveTokenBucket provides a concurrency safe utility for adding and
      9 // removing tokens from the available token bucket.
     10 type adaptiveTokenBucket struct {
     11 	remainingTokens float64
     12 	maxCapacity     float64
     13 	minCapacity     float64
     14 	mu              sync.Mutex
     15 }
     16 
     17 // newAdaptiveTokenBucket returns an initialized adaptiveTokenBucket with the
     18 // capacity specified.
     19 func newAdaptiveTokenBucket(i float64) *adaptiveTokenBucket {
     20 	return &adaptiveTokenBucket{
     21 		remainingTokens: i,
     22 		maxCapacity:     i,
     23 		minCapacity:     1,
     24 	}
     25 }
     26 
     27 // Retrieve attempts to reduce the available tokens by the amount requested. If
     28 // there are tokens available true will be returned along with the number of
     29 // available tokens remaining. If amount requested is larger than the available
     30 // capacity, false will be returned along with the available capacity. If the
     31 // amount is less than the available capacity, the capacity will be reduced by
     32 // that amount, and the remaining capacity and true will be returned.
     33 func (t *adaptiveTokenBucket) Retrieve(amount float64) (available float64, retrieved bool) {
     34 	t.mu.Lock()
     35 	defer t.mu.Unlock()
     36 
     37 	if amount > t.remainingTokens {
     38 		return t.remainingTokens, false
     39 	}
     40 
     41 	t.remainingTokens -= amount
     42 	return t.remainingTokens, true
     43 }
     44 
     45 // Refund returns the amount of tokens back to the available token bucket, up
     46 // to the initial capacity.
     47 func (t *adaptiveTokenBucket) Refund(amount float64) {
     48 	t.mu.Lock()
     49 	defer t.mu.Unlock()
     50 
     51 	// Capacity cannot exceed max capacity.
     52 	t.remainingTokens = math.Min(t.remainingTokens+amount, t.maxCapacity)
     53 }
     54 
     55 // Capacity returns the maximum capacity of tokens that the bucket could
     56 // contain.
     57 func (t *adaptiveTokenBucket) Capacity() float64 {
     58 	t.mu.Lock()
     59 	defer t.mu.Unlock()
     60 
     61 	return t.maxCapacity
     62 }
     63 
     64 // Remaining returns the number of tokens that remaining in the bucket.
     65 func (t *adaptiveTokenBucket) Remaining() float64 {
     66 	t.mu.Lock()
     67 	defer t.mu.Unlock()
     68 
     69 	return t.remainingTokens
     70 }
     71 
     72 // Resize adjusts the size of the token bucket. Returns the capacity remaining.
     73 func (t *adaptiveTokenBucket) Resize(size float64) float64 {
     74 	t.mu.Lock()
     75 	defer t.mu.Unlock()
     76 
     77 	t.maxCapacity = math.Max(size, t.minCapacity)
     78 
     79 	// Capacity needs to be capped at max capacity, if max size reduced.
     80 	t.remainingTokens = math.Min(t.remainingTokens, t.maxCapacity)
     81 
     82 	return t.remainingTokens
     83 }