130 lines
3.2 KiB
Go
130 lines
3.2 KiB
Go
package tb
|
|
|
|
import (
|
|
"math"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
// Bucket defines a generic lock-free implementation of a Token Bucket.
|
|
type Bucket struct {
|
|
inc int64
|
|
tokens int64
|
|
capacity int64
|
|
freq time.Duration
|
|
closing chan struct{}
|
|
}
|
|
|
|
// NewBucket returns a full Bucket with c capacity and starts a filling
|
|
// go-routine which ticks every freq. The number of tokens added on each tick
|
|
// is computed dynamically to be even across the duration of a second.
|
|
//
|
|
// If freq == -1 then the filling go-routine won't be started. Otherwise,
|
|
// If freq < 1/c seconds, then it will be adjusted to 1/c seconds.
|
|
func NewBucket(c int64, freq time.Duration) *Bucket {
|
|
b := &Bucket{tokens: c, capacity: c, closing: make(chan struct{})}
|
|
|
|
if freq == -1 {
|
|
return b
|
|
} else if evenFreq := time.Duration(1e9 / c); freq < evenFreq {
|
|
freq = evenFreq
|
|
}
|
|
|
|
b.freq = freq
|
|
b.inc = int64(math.Floor(.5 + (float64(c) * freq.Seconds())))
|
|
|
|
go b.fill()
|
|
|
|
return b
|
|
}
|
|
|
|
// Take attempts to take n tokens out of the bucket.
|
|
// If tokens == 0, nothing will be taken.
|
|
// If n <= tokens, n tokens will be taken.
|
|
// If n > tokens, all tokens will be taken.
|
|
//
|
|
// This method is thread-safe.
|
|
func (b *Bucket) Take(n int64) (taken int64) {
|
|
for {
|
|
if tokens := atomic.LoadInt64(&b.tokens); tokens == 0 {
|
|
return 0
|
|
} else if n <= tokens {
|
|
if !atomic.CompareAndSwapInt64(&b.tokens, tokens, tokens-n) {
|
|
continue
|
|
}
|
|
return n
|
|
} else if atomic.CompareAndSwapInt64(&b.tokens, tokens, 0) { // Spill
|
|
return tokens
|
|
}
|
|
}
|
|
}
|
|
|
|
// Put attempts to add n tokens to the bucket.
|
|
// If tokens == capacity, nothing will be added.
|
|
// If n <= capacity - tokens, n tokens will be added.
|
|
// If n > capacity - tokens, capacity - tokens will be added.
|
|
//
|
|
// This method is thread-safe.
|
|
func (b *Bucket) Put(n int64) (added int64) {
|
|
for {
|
|
if tokens := atomic.LoadInt64(&b.tokens); tokens == b.capacity {
|
|
return 0
|
|
} else if left := b.capacity - tokens; n <= left {
|
|
if !atomic.CompareAndSwapInt64(&b.tokens, tokens, tokens+n) {
|
|
continue
|
|
}
|
|
return n
|
|
} else if atomic.CompareAndSwapInt64(&b.tokens, tokens, b.capacity) {
|
|
return left
|
|
}
|
|
}
|
|
}
|
|
|
|
// Wait waits for n amount of tokens to be available.
|
|
// If n tokens are immediatelly available it doesn't sleep.
|
|
// Otherwise, it sleeps the minimum amount of time required for the remaining
|
|
// tokens to be available. It returns the wait duration.
|
|
//
|
|
// This method is thread-safe.
|
|
func (b *Bucket) Wait(n int64) time.Duration {
|
|
var rem int64
|
|
if rem = n - b.Take(n); rem == 0 {
|
|
return 0
|
|
}
|
|
|
|
var wait time.Duration
|
|
for rem > 0 {
|
|
sleep := b.wait(rem)
|
|
wait += sleep
|
|
time.Sleep(sleep)
|
|
rem -= b.Take(rem)
|
|
}
|
|
return wait
|
|
}
|
|
|
|
// Close stops the filling go-routine given it was started.
|
|
func (b *Bucket) Close() error {
|
|
close(b.closing)
|
|
return nil
|
|
}
|
|
|
|
// wait returns the minimum amount of time required for n tokens to be available.
|
|
// if n > capacity, n will be adjusted to capacity
|
|
func (b *Bucket) wait(n int64) time.Duration {
|
|
return time.Duration(int64(math.Ceil(math.Min(float64(n), float64(b.capacity))/float64(b.inc))) *
|
|
b.freq.Nanoseconds())
|
|
}
|
|
|
|
func (b *Bucket) fill() {
|
|
ticker := time.NewTicker(b.freq)
|
|
defer ticker.Stop()
|
|
|
|
for _ = range ticker.C {
|
|
select {
|
|
case <-b.closing:
|
|
return
|
|
default:
|
|
b.Put(b.inc)
|
|
}
|
|
}
|
|
}
|