269 lines
8.8 KiB
Go
269 lines
8.8 KiB
Go
package fa
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
// TestRateLimiter_ShapesThroughput asserts that the rate limiter actually
|
|
// paces requests. Uses real time (small intervals) rather than synctest so
|
|
// it works on any Go 1.25+ host without GOEXPERIMENT.
|
|
func TestRateLimiter_ShapesThroughput(t *testing.T) {
|
|
lim := newRateLimiter(50*time.Millisecond, 1, false)
|
|
ctx := context.Background()
|
|
|
|
start := time.Now()
|
|
const N = 4
|
|
for i := 0; i < N; i++ {
|
|
if err := lim.wait(ctx); err != nil {
|
|
t.Fatalf("wait: %v", err)
|
|
}
|
|
}
|
|
elapsed := time.Since(start)
|
|
// With burst=1 and interval=50ms, N=4 should take at least 3*50ms = 150ms.
|
|
// Allow a small lower-bound floor; the first token is free with burst=1.
|
|
if elapsed < 100*time.Millisecond {
|
|
t.Fatalf("expected at least 100ms for %d requests; got %v", N, elapsed)
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_RespectsContextCancel(t *testing.T) {
|
|
lim := newRateLimiter(time.Hour, 1, false) // effectively never refills
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
|
|
// Drain the burst token.
|
|
if err := lim.wait(ctx); err != nil {
|
|
t.Fatalf("first wait: %v", err)
|
|
}
|
|
go func() {
|
|
time.Sleep(20 * time.Millisecond)
|
|
cancel()
|
|
}()
|
|
err := lim.wait(ctx)
|
|
if err == nil {
|
|
t.Fatal("expected error from cancelled context")
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_NilSafe(t *testing.T) {
|
|
var r *rateLimiter
|
|
if err := r.wait(context.Background()); err != nil {
|
|
t.Fatalf("nil-receiver wait should be no-op; got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestWithPriority_ContextMarker(t *testing.T) {
|
|
if got := priorityOf(context.Background()); got != PriorityNormal {
|
|
t.Errorf("undecorated context = %v; want PriorityNormal", got)
|
|
}
|
|
if got := priorityOf(nil); got != PriorityNormal {
|
|
t.Errorf("nil context = %v; want PriorityNormal", got)
|
|
}
|
|
if got := priorityOf(WithBackgroundPriority(context.Background())); got != PriorityBackground {
|
|
t.Errorf("WithBackgroundPriority = %v; want PriorityBackground", got)
|
|
}
|
|
if got := priorityOf(WithPriority(context.Background(), PriorityInteractive)); got != PriorityInteractive {
|
|
t.Errorf("WithPriority(Interactive) = %v; want PriorityInteractive", got)
|
|
}
|
|
// The marker survives further decoration.
|
|
ctx, cancel := context.WithCancel(WithPriority(context.Background(), PriorityLow))
|
|
defer cancel()
|
|
if got := priorityOf(ctx); got != PriorityLow {
|
|
t.Errorf("priority marker should survive a child context; got %v", got)
|
|
}
|
|
}
|
|
|
|
// orderRecorder collects the completion order of concurrent wait calls.
|
|
type orderRecorder struct {
|
|
mu sync.Mutex
|
|
seen []string
|
|
}
|
|
|
|
func (o *orderRecorder) wait(t *testing.T, wg *sync.WaitGroup, lim *rateLimiter, ctx context.Context, name string) {
|
|
t.Helper()
|
|
defer wg.Done()
|
|
if err := lim.wait(ctx); err != nil {
|
|
t.Errorf("%s wait: %v", name, err)
|
|
return
|
|
}
|
|
o.mu.Lock()
|
|
o.seen = append(o.seen, name)
|
|
o.mu.Unlock()
|
|
}
|
|
|
|
func (o *orderRecorder) joined() string {
|
|
o.mu.Lock()
|
|
defer o.mu.Unlock()
|
|
return strings.Join(o.seen, ",")
|
|
}
|
|
|
|
// TestRateLimiter_BackgroundDefersToForeground: with priority enabled, a
|
|
// background request (B) defers to BOTH foreground requests even C, which
|
|
// is issued *after* B so B completes last. A and C are the same priority
|
|
// and race, so their order relative to each other is unspecified.
|
|
func TestRateLimiter_BackgroundDefersToForeground(t *testing.T) {
|
|
lim := newRateLimiter(60*time.Millisecond, 1, true)
|
|
bg := context.Background()
|
|
if err := lim.wait(bg); err != nil { // drain the burst token
|
|
t.Fatalf("drain: %v", err)
|
|
}
|
|
|
|
var order orderRecorder
|
|
var wg sync.WaitGroup
|
|
wg.Add(3)
|
|
|
|
go order.wait(t, &wg, lim, bg, "A") // foreground, t0
|
|
time.Sleep(20 * time.Millisecond)
|
|
go order.wait(t, &wg, lim, WithBackgroundPriority(bg), "B") // background, t0+20ms
|
|
time.Sleep(20 * time.Millisecond)
|
|
go order.wait(t, &wg, lim, bg, "C") // foreground, t0+40ms
|
|
|
|
wg.Wait()
|
|
// A and C are both foreground (PriorityNormal); same-level requests race,
|
|
// so their relative order is unspecified. The guarantee is only that the
|
|
// background request B is served last.
|
|
if got := order.joined(); got != "A,C,B" && got != "C,A,B" {
|
|
t.Errorf("completion order = %q; want background (B) last, foreground A and C before it in either order", got)
|
|
}
|
|
}
|
|
|
|
// TestRateLimiter_PriorityDisabledServesFIFO: the same scenario with priority
|
|
// disabled the WithBackgroundPriority marker is inert, so requests are
|
|
// served strictly in the order they queued: A, B, C.
|
|
func TestRateLimiter_PriorityDisabledServesFIFO(t *testing.T) {
|
|
lim := newRateLimiter(60*time.Millisecond, 1, false)
|
|
bg := context.Background()
|
|
if err := lim.wait(bg); err != nil { // drain the burst token
|
|
t.Fatalf("drain: %v", err)
|
|
}
|
|
|
|
var order orderRecorder
|
|
var wg sync.WaitGroup
|
|
wg.Add(3)
|
|
|
|
go order.wait(t, &wg, lim, bg, "A")
|
|
time.Sleep(20 * time.Millisecond)
|
|
go order.wait(t, &wg, lim, WithBackgroundPriority(bg), "B") // marked, but ignored
|
|
time.Sleep(20 * time.Millisecond)
|
|
go order.wait(t, &wg, lim, bg, "C")
|
|
|
|
wg.Wait()
|
|
if got := order.joined(); got != "A,B,C" {
|
|
t.Errorf("completion order = %q; want \"A,B,C\" (priority off marker must be inert)", got)
|
|
}
|
|
}
|
|
|
|
// TestRateLimiter_BackgroundProceedsWithoutForeground confirms a background
|
|
// request does not hang when nothing foreground is competing.
|
|
func TestRateLimiter_BackgroundProceedsWithoutForeground(t *testing.T) {
|
|
lim := newRateLimiter(20*time.Millisecond, 1, true)
|
|
ctx := WithBackgroundPriority(context.Background())
|
|
for i := 0; i < 3; i++ {
|
|
c, cancel := context.WithTimeout(ctx, time.Second)
|
|
err := lim.wait(c)
|
|
cancel()
|
|
if err != nil {
|
|
t.Fatalf("background wait %d: %v", i, err)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestRateLimiter_BackgroundWaitRespectsCancel confirms a background request
|
|
// blocked behind a foreground request still honours context cancellation.
|
|
func TestRateLimiter_BackgroundWaitRespectsCancel(t *testing.T) {
|
|
lim := newRateLimiter(time.Hour, 1, true) // bucket effectively never refills
|
|
bg := context.Background()
|
|
if err := lim.wait(bg); err != nil {
|
|
t.Fatalf("drain: %v", err)
|
|
}
|
|
|
|
// A foreground request parks in wait forever, keeping the gate closed.
|
|
fgCtx, fgCancel := context.WithCancel(bg)
|
|
defer fgCancel()
|
|
fgStarted := make(chan struct{})
|
|
go func() {
|
|
close(fgStarted)
|
|
_ = lim.wait(fgCtx)
|
|
}()
|
|
<-fgStarted
|
|
time.Sleep(20 * time.Millisecond) // ensure the foreground counter is set
|
|
|
|
ctx, cancel := context.WithCancel(WithBackgroundPriority(bg))
|
|
go func() {
|
|
time.Sleep(20 * time.Millisecond)
|
|
cancel()
|
|
}()
|
|
if err := lim.wait(ctx); err == nil {
|
|
t.Fatal("expected cancellation error for background wait blocked behind foreground")
|
|
}
|
|
}
|
|
|
|
// TestNew_PrioritizedRateLimitingOption verifies the construction-time
|
|
// opt-in wires through to the limiter.
|
|
func TestNew_PrioritizedRateLimitingOption(t *testing.T) {
|
|
if New().limiter.priority {
|
|
t.Error("default client: priority scheduling should be off")
|
|
}
|
|
if !New(WithPrioritizedRateLimiting(true)).limiter.priority {
|
|
t.Error("WithPrioritizedRateLimiting(true): priority scheduling should be on")
|
|
}
|
|
if New(WithPrioritizedRateLimiting(false)).limiter.priority {
|
|
t.Error("WithPrioritizedRateLimiting(false): priority scheduling should be off")
|
|
}
|
|
}
|
|
|
|
func TestPriority_String(t *testing.T) {
|
|
cases := map[Priority]string{
|
|
PriorityInteractive: "interactive",
|
|
PriorityNormal: "normal",
|
|
PriorityLow: "low",
|
|
PriorityBackground: "background",
|
|
Priority(42): "unknown",
|
|
}
|
|
for p, want := range cases {
|
|
if got := p.String(); got != want {
|
|
t.Errorf("Priority(%d).String() = %q; want %q", int(p), got, want)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestWithPriority_ClampsOutOfRange(t *testing.T) {
|
|
if got := priorityOf(WithPriority(context.Background(), Priority(-5))); got != PriorityInteractive {
|
|
t.Errorf("below-range priority = %v; want PriorityInteractive", got)
|
|
}
|
|
if got := priorityOf(WithPriority(context.Background(), Priority(99))); got != PriorityBackground {
|
|
t.Errorf("above-range priority = %v; want PriorityBackground", got)
|
|
}
|
|
}
|
|
|
|
// TestRateLimiter_NLevelPriorityOrder: with priority enabled, four requests
|
|
// queued lowest-first must complete highest-priority-first, not in arrival
|
|
// order.
|
|
func TestRateLimiter_NLevelPriorityOrder(t *testing.T) {
|
|
lim := newRateLimiter(60*time.Millisecond, 1, true)
|
|
bg := context.Background()
|
|
if err := lim.wait(bg); err != nil { // drain the burst token
|
|
t.Fatalf("drain: %v", err)
|
|
}
|
|
|
|
var order orderRecorder
|
|
var wg sync.WaitGroup
|
|
wg.Add(4)
|
|
|
|
go order.wait(t, &wg, lim, WithPriority(bg, PriorityBackground), "background")
|
|
time.Sleep(15 * time.Millisecond)
|
|
go order.wait(t, &wg, lim, WithPriority(bg, PriorityLow), "low")
|
|
time.Sleep(15 * time.Millisecond)
|
|
go order.wait(t, &wg, lim, WithPriority(bg, PriorityNormal), "normal")
|
|
time.Sleep(15 * time.Millisecond)
|
|
go order.wait(t, &wg, lim, WithPriority(bg, PriorityInteractive), "interactive")
|
|
|
|
wg.Wait()
|
|
if got := order.joined(); got != "interactive,normal,low,background" {
|
|
t.Errorf("completion order = %q; want \"interactive,normal,low,background\"", got)
|
|
}
|
|
}
|