equalizer.go

package equalizer

import (
	"fmt"
	"math/big"
	"strings"
	"sync"
)

// An Equalizer represents an adaptive rate limiter based on a bit array.
//
// The Equalizer uses a round-robin bit array with a moving head to manage
// quotas.
// The quota management algorithm is simple and works in the following way.
// To request a permit in a non-blocking manner use the TryAcquire method.
// The Equalizer will locate the appropriate position on the tape using the
// offset manager and return the value, denoting whether the request is allowed
// or not. To update the tape state, a notification method (Success or Failure)
// should be invoked based on the operation status.
// The Reset and Purge methods allow for immediate transition of the limiter
// state to whether permissive or restrictive.
//
// An Equalizer is safe for use by multiple goroutines simultaneously.
type Equalizer struct {
	sync.RWMutex
	// tape is the underlying bit array
	tape *big.Int
	// seed is the initial state of the bit array tape
	seed *big.Int
	// mask is the positive bitmask
	mask *big.Int
	// offset is the next index offset manager
	offset Offset
	// notifyHead is the moving pointer for notifications
	notifyHead int
	// adjustable is the number of unmasked bits
	adjustable int
}

// NewEqualizer instantiates and returns a new Equalizer rate limiter, where
// size is the length of the bit array, reserved is the number of reserved
// positive bits and offset is an instance of the equalizer.Offset strategy.
func NewEqualizer(size, reserved int, offset Offset) (*Equalizer, error) {
	if offset == nil {
		return nil, fmt.Errorf("offset is nil")
	}
	if size < 1 {
		return nil, fmt.Errorf("nonpositive size: %d", size)
	}
	if reserved < 1 {
		return nil, fmt.Errorf("nonpositive reserved: %d", reserved)
	}
	if reserved > size {
		return nil, fmt.Errorf("reserved must not exceed size")
	}
	// init the seed bit array
	var seed big.Int
	seed.SetString(strings.Repeat("1", size), 2)

	// init the positive bitmask
	var mask big.Int
	mask.SetString(strings.Repeat("1", reserved), 2)
	mask.Lsh(&mask, uint(size-reserved))

	// init the operational bit array tape
	var tape big.Int
	tape.Set(&seed)

	return &Equalizer{
		tape:       &tape,
		seed:       &seed,
		mask:       &mask,
		offset:     offset,
		adjustable: size - reserved,
	}, nil
}

// TryAcquire moves the tape head to the next index and returns the value.
func (eq *Equalizer) TryAcquire() bool {
	eq.RLock()
	defer eq.RUnlock()

	head := eq.offset.NextIndex()
	return eq.tape.Bit(head) > 0
}

// Success notifies the equalizer with n successful operations.
func (eq *Equalizer) Success(n int) {
	eq.notify(n, 1)
}

// Failure notifies the equalizer with n failed operations.
func (eq *Equalizer) Failure(n int) {
	eq.notify(n, 0)
}

// notify advances the notification head by n bits, assigning the given value.
func (eq *Equalizer) notify(n int, value uint) {
	eq.Lock()
	defer eq.Unlock()

	for i := 0; i < n; i++ {
		pos := eq.notifyHead % eq.adjustable
		eq.tape.SetBit(eq.tape, pos, value)
		eq.notifyHead++
	}
}

// Reset resets the tape to its initial state.
func (eq *Equalizer) Reset() {
	eq.Lock()
	defer eq.Unlock()

	eq.tape.Set(eq.seed)
}

// Purge blanks out the tape to the positive bitmask state.
func (eq *Equalizer) Purge() {
	eq.Lock()
	defer eq.Unlock()

	eq.tape.Set(eq.mask)
}