src

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

token_bucket.go (2387B)

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