token_bucket.go

package equalizer

import (
	"context"
	"fmt"
	"runtime"
	"time"

	"github.com/reugn/equalizer/internal/async"
)

// A TokenBucket represents a rate limiter based on a custom implementation of the
// token bucket algorithm with a refill interval. Implements the equalizer.Limiter
// interface.
//
// A TokenBucket is safe for use by multiple goroutines simultaneously.
//
// The underlying tokenBucket instance ensures that the background goroutine does
// not keep the main TokenBucket object from being garbage collected. When it is
// garbage collected, the finalizer stops the background goroutine, after
// which the underlying tokenBucket can be collected.
type TokenBucket struct {
	t *tokenBucket
}

var _ Limiter = (*TokenBucket)(nil)

type tokenBucket struct {
	capacity       int
	refillInterval time.Duration
	permits        chan token
	tokenSupplier  *async.Task
}

// NewTokenBucket allocates and returns a new TokenBucket rate limiter, where
// capacity is the capacity of the token bucket and refillInterval is the token
// bucket refill interval.
func NewTokenBucket(capacity int, refillInterval time.Duration) (*TokenBucket, error) {
	if capacity < 1 {
		return nil, fmt.Errorf("nonpositive capacity: %d", capacity)
	}
	underlying := &tokenBucket{
		capacity:       capacity,
		refillInterval: refillInterval,
		permits:        make(chan token, capacity),
		tokenSupplier:  async.NewTask(refillInterval),
	}
	underlying.fillTokenBucket(capacity)
	// initialize a goroutine responsible for periodically refilling the
	// token bucket
	underlying.tokenSupplier.Run(underlying.refill)
	tokenBucket := &TokenBucket{
		t: underlying,
	}
	// the finalizer may run as soon as the tokenBucket becomes unreachable
	runtime.SetFinalizer(tokenBucket, stopTokenSupplier)

	return tokenBucket, nil
}

// fillTokenBucket fills up the permits channel with the capacity number of tokens.
func (tb *tokenBucket) fillTokenBucket(capacity int) {
	for i := 0; i < capacity; i++ {
		tb.permits <- token{}
	}
}

// refill refills the token bucket.
func (tb *tokenBucket) refill() {
	issued := tb.capacity - len(tb.permits)
	tb.fillTokenBucket(issued)
}

// Acquire blocks the calling goroutine until a token is acquired, the Context
// is canceled, or the wait time exceeds the Context's Deadline.
func (tb *TokenBucket) Acquire(ctx context.Context) error {
	select {
	case <-tb.t.permits:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// TryAcquire attempts to acquire a token without blocking.
// Returns true if a token was acquired, false if no tokens are available.
func (tb *TokenBucket) TryAcquire() bool {
	select {
	case <-tb.t.permits:
		return true
	default:
		return false
	}
}

// stopTokenSupplier is the callback function used to terminate the
// goroutine responsible for periodically refilling the token bucket.
func stopTokenSupplier(t *TokenBucket) {
	t.t.tokenSupplier.Stop()
}