《后端架构设计·高性能与高可用》6.1 限流算法与分布式限流

// file: internal/platform/ratelimit/ratelimit.go
// tier: L1
// verify: go test ./internal/platform/ratelimit/
package ratelimit

import (
	"math"
	"sync"
	"time"
)

// Clock is injected so every limiter can be tested without sleeping.
type Clock interface{ Now() time.Time }

type ManualClock struct{ now time.Time }

func NewManualClock(t time.Time) *ManualClock  { return &ManualClock{now: t} }
func (c *ManualClock) Now() time.Time          { return c.now }
func (c *ManualClock) Advance(d time.Duration) { c.now = c.now.Add(d) }

type Decision struct {
	Allowed    bool
	RetryAfter time.Duration
}

type Limiter interface{ Allow() Decision }

type TokenBucket struct {
	mu                     sync.Mutex
	clock                  Clock
	capacity, rate, tokens float64
	last                   time.Time
}

func NewTokenBucket(clock Clock, capacity int, ratePerSecond float64) *TokenBucket {
	return &TokenBucket{clock: clock, capacity: float64(capacity), rate: ratePerSecond,
		tokens: float64(capacity), last: clock.Now()}
}

func (b *TokenBucket) Allow() Decision {
	b.mu.Lock()
	defer b.mu.Unlock()
	now := b.clock.Now()
	elapsed := now.Sub(b.last).Seconds()
	if elapsed < 0 {
		b.last = now
		elapsed = 0
	}
	if elapsed > 0 {
		b.tokens = math.Min(b.capacity, b.tokens+elapsed*b.rate)
		b.last = now
	}
	if b.tokens >= 1 {
		b.tokens--
		return Decision{Allowed: true}
	}
	need := (1 - b.tokens) / b.rate
	return Decision{RetryAfter: time.Duration(math.Ceil(need * float64(time.Second)))}
}

// SlidingWindow keeps admission timestamps and rejects while the window is full.
type SlidingWindow struct {
	mu     sync.Mutex
	clock  Clock
	limit  int
	window time.Duration
	times  []time.Time
}

func NewSlidingWindow(clock Clock, limit int, window time.Duration) *SlidingWindow {
	return &SlidingWindow{clock: clock, limit: limit, window: window}
}

func (w *SlidingWindow) Allow() Decision {
	w.mu.Lock()
	defer w.mu.Unlock()
	now := w.clock.Now()
	kept := w.times[:0]
	for _, t := range w.times {
		if t.After(now.Add(-w.window)) {
			kept = append(kept, t)
		}
	}
	w.times = kept
	if len(w.times) < w.limit {
		w.times = append(w.times, now)
		return Decision{Allowed: true}
	}
	retry := w.times[0].Add(w.window).Sub(now)
	if retry < 0 {
		retry = 0
	}
	return Decision{RetryAfter: retry}
}

bucket   allowed=30 denied=5 max_at_once=10
window   allowed=21 denied=14 max_at_once=10
bucket20 allowed=35 denied=0 max_at_once=15

// file: internal/platform/ratelimit/lease.go
// tier: L1
// verify: go test ./internal/platform/ratelimit/
package ratelimit

import (
	"math"
	"sync"
	"time"
)

// Coordinator is an in-process stand-in for a shared quota service: it models
// only the arithmetic and the two clocks, not Redis.
type Coordinator struct {
	mu          sync.Mutex
	clock       Clock
	capacity    float64
	rate        float64
	tokens      float64
	last        time.Time
	leaseTokens int
	ttl         time.Duration
	granted     int
}

func NewCoordinator(clock Clock, rate float64, burst, leaseTokens int, ttl time.Duration) *Coordinator {
	return &Coordinator{clock: clock, capacity: float64(burst), rate: rate, tokens: float64(burst),
		last: clock.Now(), leaseTokens: leaseTokens, ttl: ttl}
}

func (c *Coordinator) RequestLease() (int, time.Time, bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	now := c.clock.Now()
	if elapsed := now.Sub(c.last).Seconds(); elapsed > 0 {
		c.tokens = math.Min(c.capacity, c.tokens+elapsed*c.rate)
		c.last = now
	}
	n := int(math.Floor(c.tokens))
	if n > c.leaseTokens {
		n = c.leaseTokens
	}
	if n < 1 {
		return 0, time.Time{}, false
	}
	c.tokens -= float64(n)
	c.granted += n
	return n, now.Add(c.ttl), true
}

type InstanceLeaseLimiter struct {
	mu        sync.Mutex
	clock     Clock
	skew      time.Duration
	coord     *Coordinator
	tokens    float64
	expiresAt time.Time
	allowed   int
	denied    int
	wasted    int
}

func NewInstanceLeaseLimiter(clock Clock, skew time.Duration, coord *Coordinator) *InstanceLeaseLimiter {
	return &InstanceLeaseLimiter{clock: clock, skew: skew, coord: coord}
}

func (i *InstanceLeaseLimiter) Allow() Decision {
	i.mu.Lock()
	defer i.mu.Unlock()
	now := i.clock.Now().Add(i.skew)
	if !i.expiresAt.IsZero() && !now.Before(i.expiresAt) {
		// Expired: unused tokens are lost; the coordinator never reissues them.
		if i.tokens >= 1 {
			i.wasted += int(math.Floor(i.tokens))
		}
		i.tokens = 0
		i.expiresAt = time.Time{}
	}
	if i.tokens < 1 {
		if n, expires, ok := i.coord.RequestLease(); ok {
			i.tokens = math.Min(float64(i.coord.leaseTokens), i.tokens+float64(n))
			i.expiresAt = expires
		}
	}
	if i.tokens >= 1 {
		i.tokens--
		i.allowed++
		return Decision{Allowed: true}
	}
	i.denied++
	return Decision{}
}

// file: internal/platform/ratelimit/ratelimit_test.go
// tier: L1
// verify: go test ./internal/platform/ratelimit/
package ratelimit

import (
	"testing"
	"time"
)

var base = time.Date(2026, 10, 18, 8, 0, 0, 0, time.UTC)

type arrival struct {
	at time.Duration
	n  int
}

type trace struct {
	allowed, denied, burst int
}

func replay(l Limiter, clock *ManualClock, as []arrival) trace {
	tr := trace{}
	var cursor time.Duration
	for _, a := range as {
		clock.Advance(a.at - cursor)
		cursor = a.at
		group := 0
		for n := 0; n < a.n; n++ {
			if l.Allow().Allowed {
				tr.allowed++
				group++
			} else {
				tr.denied++
			}
		}
		if group > tr.burst {
			tr.burst = group
		}
	}
	return tr
}

func TestTokenBucket(t *testing.T) {
	clock := NewManualClock(base)
	b := NewTokenBucket(clock, 10, 10)
	for i := 0; i < 10; i++ {
		if !b.Allow().Allowed {
			t.Fatalf("burst %d rejected", i+1)
		}
	}
	if d := b.Allow(); d.Allowed || d.RetryAfter != 100*time.Millisecond {
		t.Fatalf("11th=%+v", d)
	}
	clock.Advance(100 * time.Millisecond)
	if !b.Allow().Allowed {
		t.Fatal("100ms must refill one token")
	}
	clock.Advance(-time.Hour)
	if b.Allow().Allowed {
		t.Fatal("backwards clock must not mint tokens")
	}
}

func TestReplayComparison(t *testing.T) {
	as := []arrival{{0, 15}}
	for i := 1; i <= 20; i++ {
		as = append(as, arrival{time.Duration(i) * 100 * time.Millisecond, 1})
	}
	clock := NewManualClock(base)
	bucket := replay(NewTokenBucket(clock, 10, 10), clock, as)
	clock = NewManualClock(base)
	window := replay(NewSlidingWindow(clock, 10, time.Second), clock, as)
	clock = NewManualClock(base)
	big := replay(NewTokenBucket(clock, 20, 10), clock, as)
	t.Logf("bucket   allowed=%d denied=%d max_at_once=%d", bucket.allowed, bucket.denied, bucket.burst)
	t.Logf("window   allowed=%d denied=%d max_at_once=%d", window.allowed, window.denied, window.burst)
	t.Logf("bucket20 allowed=%d denied=%d max_at_once=%d", big.allowed, big.denied, big.burst)
	if bucket.allowed != 30 || window.allowed != 21 || bucket.burst != 10 || window.burst != 10 || big.burst != 15 {
		t.Fatalf("comparison mismatch: %+v %+v %+v", bucket, window, big)
	}
}

func TestLeaseSkew(t *testing.T) {
	every := []time.Duration{10 * time.Millisecond, 33 * time.Millisecond, 50 * time.Millisecond}
	for _, skew := range []time.Duration{0, 150 * time.Millisecond, -150 * time.Millisecond} {
		clock := NewManualClock(base)
		coord := NewCoordinator(clock, 150, 60, 20, 400*time.Millisecond)
		insts := make([]*InstanceLeaseLimiter, len(every))
		for k := range every {
			insts[k] = NewInstanceLeaseLimiter(clock, skew, coord)
		}
		allowed, denied, wasted := 0, 0, 0
		for ms := 0; ms <= 10000; ms++ {
			clock.Advance(time.Millisecond)
			for k, e := range every {
				if ms%int(e/time.Millisecond) == 0 && insts[k].Allow().Allowed {
					allowed++
				}
			}
		}
		shares := make([]int, len(insts))
		for k, inst := range insts {
			denied += inst.denied
			wasted += inst.wasted
			shares[k] = inst.allowed
		}
		t.Logf("skew=%-6s demand=%d allowed=%d denied=%d wasted=%d granted=%d shares=%v",
			skew, allowed+denied, allowed, denied, wasted, coord.granted, shares)
		if allowed > coord.granted {
			t.Fatalf("skew=%s allowed=%d exceeds granted=%d", skew, allowed, coord.granted)
		}
	}
}