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 }