ratelimiter/ratelimiter_test.go

package ratelimiter

import (
	"errors"
	"sync"
	"testing"
	"testing/synctest"
	"time"
)

func TestNew_Panics(t *testing.T) {
	cases := []struct {
		name string
		conf Config
	}{
		{"negative BurstLimit", Config{BurstLimit: -1, FillPeriod: time.Second, MaxWaitCount: 1}},
		{"zero FillPeriod", Config{BurstLimit: 1, FillPeriod: 0, MaxWaitCount: 1}},
		{"negative FillPeriod", Config{BurstLimit: 1, FillPeriod: -time.Second, MaxWaitCount: 1}},
		{"negative MaxWaitCount", Config{BurstLimit: 1, FillPeriod: time.Second, MaxWaitCount: -1}},
		{"zero BurstLimit and MaxWaitCount", Config{BurstLimit: 0, FillPeriod: time.Second, MaxWaitCount: 0}},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			defer func() {
				if recover() == nil {
					t.Error("expected panic, got none")
				}
			}()
			New(tc.conf)
		})
	}
}

// limitN runs n goroutines concurrently, each calling l.Limit(), and returns
// their errors in an unordered slice.
func limitN(l *Limiter, n int) []error {
	errs := make([]error, n)
	var wg sync.WaitGroup
	for i := range n {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			errs[i] = l.Limit()
		}(i)
	}
	wg.Wait()
	return errs
}

func backoffCount(errs []error) int {
	n := 0
	for _, err := range errs {
		if errors.Is(err, ErrBackoff) {
			n++
		}
	}
	return n
}

// TestLimit_Burst: burst capacity is consumed by the first BurstLimit requests;
// excess requests are rejected without any waiting.
func TestLimit_Burst(t *testing.T) {
	const burst, total = 5, 10
	synctest.Test(t, func(t *testing.T) {
		t0 := time.Now()
		l := New(Config{BurstLimit: burst, FillPeriod: time.Second, MaxWaitCount: 0})
		errs := limitN(l, total)
		if got := backoffCount(errs); got != total-burst {
			t.Errorf("backoffs: want %d, got %d", total-burst, got)
		}
		if elapsed := time.Since(t0); elapsed != 0 {
			t.Errorf("elapsed: want 0, got %v", elapsed)
		}
	})
}

// TestLimit_Queue: up to MaxWaitCount requests queue and sleep; excess are rejected.
func TestLimit_Queue(t *testing.T) {
	const maxWait, total = 5, 8
	const period = time.Second
	synctest.Test(t, func(t *testing.T) {
		t0 := time.Now()
		l := New(Config{BurstLimit: 0, FillPeriod: period, MaxWaitCount: maxWait})
		errs := limitN(l, total)
		if got := backoffCount(errs); got != total-maxWait {
			t.Errorf("backoffs: want %d, got %d", total-maxWait, got)
		}
		// The last queued request sleeps maxWait*period.
		if elapsed := time.Since(t0); elapsed != time.Duration(maxWait)*period {
			t.Errorf("elapsed: want %v, got %v", time.Duration(maxWait)*period, elapsed)
		}
	})
}

// TestLimit_BurstAndQueue: burst requests pass immediately; overflow queues;
// excess are rejected.
func TestLimit_BurstAndQueue(t *testing.T) {
	const burst, maxWait, total = 3, 5, 12
	const period = time.Second
	synctest.Test(t, func(t *testing.T) {
		t0 := time.Now()
		l := New(Config{BurstLimit: burst, FillPeriod: period, MaxWaitCount: maxWait})
		errs := limitN(l, total)
		if got := backoffCount(errs); got != total-(burst+maxWait) {
			t.Errorf("backoffs: want %d, got %d", total-(burst+maxWait), got)
		}
		// Burst passes at T=0; the last queued request sleeps maxWait*period.
		if elapsed := time.Since(t0); elapsed != time.Duration(maxWait)*period {
			t.Errorf("elapsed: want %v, got %v", time.Duration(maxWait)*period, elapsed)
		}
	})
}

// TestLimit_Refill: capacity refills as time passes, allowing further requests.
func TestLimit_Refill(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		l := New(Config{BurstLimit: 2, FillPeriod: time.Second, MaxWaitCount: 0})

		// Exhaust burst (waitTime: -2s → -1s → 0s).
		if err := l.Limit(); err != nil {
			t.Fatal("call 1:", err)
		}
		if err := l.Limit(); err != nil {
			t.Fatal("call 2:", err)
		}
		// Next would push waitTime to 1s > 0 = maxWaitTime.
		if err := l.Limit(); !errors.Is(err, ErrBackoff) {
			t.Fatalf("call 3: want ErrBackoff, got %v", err)
		}

		// After 2s, two tokens have refilled (waitTime: 0s - 2s = -2s, clamped to -2s).
		time.Sleep(2 * time.Second)

		// Burst is fully restored; same pattern repeats.
		if err := l.Limit(); err != nil {
			t.Fatal("call 4:", err)
		}
		if err := l.Limit(); err != nil {
			t.Fatal("call 5:", err)
		}
		if err := l.Limit(); !errors.Is(err, ErrBackoff) {
			t.Fatalf("call 6: want ErrBackoff, got %v", err)
		}
	})
}

// TestLimit_BurstCap: idle time does not accumulate credit beyond BurstLimit.
//
// After sleeping far longer than BurstLimit*FillPeriod, the limiter clamps
// waitTime back to minWaitTime, so only BurstLimit requests can burst through.
// Without the clamp, 10s of idle time at FillPeriod=1s would allow 10 requests.
func TestLimit_BurstCap(t *testing.T) {
	const burst = 2
	synctest.Test(t, func(t *testing.T) {
		l := New(Config{BurstLimit: burst, FillPeriod: time.Second, MaxWaitCount: 0})

		time.Sleep(10 * time.Second)

		errs := limitN(l, burst+3)
		if got := backoffCount(errs); got != 3 {
			t.Errorf("backoffs: want 3, got %d", got)
		}
	})
}