slider.go

package equalizer

import (
	"context"
	"fmt"
	"runtime"
	"sync/atomic"
	"time"

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

// A Slider represents a rate limiter which is based on a sliding window
// with a specified quota capacity. Implements the equalizer.Limiter interface.
//
// A Slider is safe for use by multiple goroutines simultaneously.
//
// The underlying slider instance ensures that the background goroutine does
// not keep the main Slider object from being garbage collected. When it is
// garbage collected, the finalizer stops the background goroutine, after
// which the underlying slider can be collected.
type Slider struct {
	slider *slider
}

var _ Limiter = (*Slider)(nil)

type slider struct {
	window          time.Duration
	slidingInterval time.Duration
	capacity        int
	permits         chan token
	issued          chan int64
	windowStart     atomic.Int64
	windowShifter   *async.Task
}

// NewSlider allocates and returns a new Slider rate limiter, where
// window is the fixed duration of the sliding window,
// slidingInterval controls how frequently a new sliding window is started and
// capacity is the quota limit for the window.
func NewSlider(window, slidingInterval time.Duration, capacity int) (*Slider, error) {
	if capacity < 1 {
		return nil, fmt.Errorf("nonpositive capacity: %d", capacity)
	}
	if slidingInterval > window {
		return nil, fmt.Errorf("slidingInterval must not exceed window")
	}
	underlying := &slider{
		window:          window,
		slidingInterval: slidingInterval,
		capacity:        capacity,
		permits:         make(chan token, capacity),
		issued:          make(chan int64, capacity),
		windowShifter:   async.NewTask(slidingInterval),
	}
	underlying.init()
	// initialize a goroutine responsible for periodically moving the
	// window forward
	underlying.windowShifter.Run(underlying.slide)
	slider := &Slider{
		slider: underlying,
	}
	// the finalizer may run as soon as the slider becomes unreachable
	runtime.SetFinalizer(slider, stopWindowShifter)

	return slider, nil
}

// init initializes the permits channel by populating it with the
// specified number of tokens.
func (s *slider) init() {
	for i := 0; i < s.capacity; i++ {
		s.permits <- token{}
	}
}

// slide moves the sliding window forward.
func (s *slider) slide() {
	now := time.Now().UnixNano()
	s.windowStart.Store(now)
	for issuedTime := range s.issued {
		s.permits <- token{}
		if issuedTime >= now {
			break
		}
	}
}

// Acquire blocks the calling goroutine until a token is acquired, the Context
// is canceled, or the wait time exceeds the Context's Deadline.
func (s *Slider) Acquire(ctx context.Context) error {
	select {
	case <-s.slider.permits:
		s.slider.issued <- s.slider.windowStart.Load()
		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 (s *Slider) TryAcquire() bool {
	select {
	case <-s.slider.permits:
		s.slider.issued <- s.slider.windowStart.Load()
		return true
	default:
		return false
	}
}

// stopWindowShifter is the callback function used to terminate the goroutine
// responsible for periodically moving the sliding window.
func stopWindowShifter(s *Slider) {
	s.slider.windowShifter.Stop()
}