diff --git a/router/token_bucket.go b/router/token_bucket.go new file mode 100644 index 00000000..254de6eb --- /dev/null +++ b/router/token_bucket.go @@ -0,0 +1,42 @@ +package router + +import ( + "time" +) + +type TokenBucket struct { + capacity int64 // Maximum capacity of bucket + tokenInterval time.Duration // Token replenishment rate + refillDuration time.Duration // Time to refill from empty + earliestUnspentToken time.Time +} + +func NewTokenBucket(capacity int64, tokenInterval time.Duration) *TokenBucket { + tb := TokenBucket{ + capacity: capacity, + tokenInterval: tokenInterval, + refillDuration: tokenInterval * time.Duration(capacity)} + + tb.earliestUnspentToken = tb.capacityToken() + + return &tb +} + +func (tb *TokenBucket) Wait() { + // If earliest unspent token is in the future, sleep until then + time.Sleep(tb.earliestUnspentToken.Sub(time.Now())) + + // Alternatively, enforce bucket capacity if necessary + capacityToken := tb.capacityToken() + if tb.earliestUnspentToken.Before(capacityToken) { + tb.earliestUnspentToken = capacityToken + } + + // 'Remove' a token from the bucket + tb.earliestUnspentToken = tb.earliestUnspentToken.Add(tb.tokenInterval) +} + +// Determine the historic token timestamp representing a full bucket +func (tb *TokenBucket) capacityToken() time.Time { + return time.Now().Add(-tb.refillDuration).Truncate(tb.tokenInterval) +}