diff --git a/ratelimiter_test.go b/ratelimiter_test.go index d3f0ef6..aff4854 100644 --- a/ratelimiter_test.go +++ b/ratelimiter_test.go @@ -1,101 +1,191 @@ package ratelimiter import ( + "errors" "sync" "testing" + "testing/synctest" "time" ) -func TestRateLimiter_Limit_Errors(t *testing.T) { - type TestCase struct { - Name string - Conf Config - N int - ErrCount int - DT time.Duration +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}}, } - - cases := []TestCase{ - { - Name: "no burst, no wait", - Conf: Config{ - BurstLimit: 0, - FillPeriod: 100 * time.Millisecond, - MaxWaitCount: 0, - }, - N: 32, - ErrCount: 32, - DT: 0, - }, { - Name: "no wait", - Conf: Config{ - BurstLimit: 10, - FillPeriod: 100 * time.Millisecond, - MaxWaitCount: 0, - }, - N: 32, - ErrCount: 22, - DT: 0, - }, { - Name: "no burst", - Conf: Config{ - BurstLimit: 0, - FillPeriod: 10 * time.Millisecond, - MaxWaitCount: 10, - }, - N: 32, - ErrCount: 22, - DT: 100 * time.Millisecond, - }, { - Name: "burst and wait", - Conf: Config{ - BurstLimit: 10, - FillPeriod: 10 * time.Millisecond, - MaxWaitCount: 10, - }, - N: 32, - ErrCount: 12, - DT: 100 * time.Millisecond, - }, - } - for _, tc := range cases { - wg := sync.WaitGroup{} - l := New(tc.Conf) - errs := make([]error, tc.N) - - t0 := time.Now() - - for i := 0; i < tc.N; i++ { - wg.Add(1) - go func(i int) { - errs[i] = l.Limit() - wg.Done() - }(i) - } - - wg.Wait() - - dt := time.Since(t0) - - errCount := 0 - for _, err := range errs { - if err != nil { - errCount++ - } - } - - if errCount != tc.ErrCount { - t.Fatalf("%s: Expected %d errors but got %d.", - tc.Name, tc.ErrCount, errCount) - } - - if dt < tc.DT { - t.Fatal(tc.Name, dt, tc.DT) - } - - if dt > tc.DT+10*time.Millisecond { - t.Fatal(tc.Name, dt, tc.DT) - } + 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) + } + }) +} + +// TestLimitMultiple: multiple tokens are consumed per call. +// +// Config: BurstLimit=4, FillPeriod=1s, MaxWaitCount=2 → minWaitTime=-4s, maxWaitTime=2s. +// +// Call 1: LimitMultiple(3) → waitTime = -4s+3s = -1s → immediate +// Call 2: LimitMultiple(3) → waitTime = -1s+3s = 2s → sleeps 2s +// Call 3: LimitMultiple(3) → waitTime = 2s+3s = 5s > 2s → ErrBackoff +func TestLimitMultiple(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + t0 := time.Now() + l := New(Config{BurstLimit: 4, FillPeriod: time.Second, MaxWaitCount: 2}) + + if err := l.LimitMultiple(3); err != nil { + t.Fatalf("call 1: %v", err) + } + if err := l.LimitMultiple(3); err != nil { + t.Fatalf("call 2: %v", err) + } + if err := l.LimitMultiple(3); !errors.Is(err, ErrBackoff) { + t.Fatalf("call 3: want ErrBackoff, got %v", err) + } + if elapsed := time.Since(t0); elapsed != 2*time.Second { + t.Errorf("elapsed: want 2s, got %v", elapsed) + } + }) +}