《互联网产品的服务端架构·交易、增长与平台化》6.2 秒杀链路与限流

// file: internal/activity/seckill.go
// tier: L1
// verify: go test ./internal/activity/

package activity

import (
	"errors"
	"fmt"
)

var (
	// ErrInvalidLimiter 表示令牌桶参数不合法。
	ErrInvalidLimiter = errors.New("activity: invalid limiter")
	// ErrLimiterClockBackwards 表示传入的 now 早于桶的上次结算时间。
	ErrLimiterClockBackwards = errors.New("activity: limiter clock backwards")
)

// Bucket 是令牌桶的当前状态。Tokens 是可用令牌数，允许小数；
// Rate 是每秒补充的令牌数；Burst 是容量上限；LastAt 是上次结算时间，
// 用 Unix 秒表示，允许小数。
type Bucket struct {
	Tokens float64
	Rate   float64
	Burst  int
	LastAt float64
}

// Allow 判断在 now 时刻是否放行一次请求。它会先按流逝时间补充令牌，
// 补充后不超过 Burst，然后尝试扣掉一个令牌。
//
// 无论放行还是拒绝，LastAt 都会推进到 now，剩余令牌也会写回，
// 这样被拒绝的请求不会“攒”出额外的令牌。now 早于 LastAt 说明时钟
// 或调用顺序异常，直接报错且不修改桶状态。
func Allow(b *Bucket, now float64) (bool, error) {
	if b == nil {
		return false, fmt.Errorf("%w: nil bucket", ErrInvalidLimiter)
	}
	if b.Rate <= 0 || b.Burst <= 0 {
		return false, fmt.Errorf("%w: rate=%v burst=%d", ErrInvalidLimiter, b.Rate, b.Burst)
	}
	if b.Tokens < 0 {
		return false, fmt.Errorf("%w: negative tokens %v", ErrInvalidLimiter, b.Tokens)
	}
	if now < b.LastAt {
		return false, fmt.Errorf("%w: now=%v last=%v", ErrLimiterClockBackwards, now, b.LastAt)
	}
	tokens := b.Tokens + (now-b.LastAt)*b.Rate
	if limit := float64(b.Burst); tokens > limit {
		tokens = limit
	}
	b.LastAt = now
	if tokens >= 1 {
		b.Tokens = tokens - 1
		return true, nil
	}
	b.Tokens = tokens
	return false, nil
}

// file: internal/activity/seckill_test.go
// tier: L1
// verify: go test ./internal/activity/

package activity

import (
	"errors"
	"testing"
)

func TestAllowRefillsOverFractionalTime(t *testing.T) {
	b := &Bucket{Tokens: 0, Rate: 4, Burst: 2, LastAt: 0}
	if ok, err := Allow(b, 0.25); err != nil || !ok {
		t.Fatalf("Allow(0.25) = %v, %v; want true, nil", ok, err)
	}
	if b.Tokens != 0 {
		t.Fatalf("tokens = %v, want 0", b.Tokens)
	}
	if ok, err := Allow(b, 0.5); err != nil || !ok {
		t.Fatalf("Allow(0.5) = %v, %v; want true, nil", ok, err)
	}
	if b.LastAt != 0.5 {
		t.Fatalf("LastAt = %v, want 0.5", b.LastAt)
	}
}

func TestAllowRejectKeepsRemainderAndAdvancesClock(t *testing.T) {
	b := &Bucket{Tokens: 0, Rate: 0.5, Burst: 2, LastAt: 10}
	ok, err := Allow(b, 11)
	if err != nil || ok {
		t.Fatalf("Allow(11) = %v, %v; want false, nil", ok, err)
	}
	if b.LastAt != 11 || b.Tokens != 0.5 {
		t.Fatalf("after reject: LastAt=%v Tokens=%v, want 11 and 0.5", b.LastAt, b.Tokens)
	}
	ok, err = Allow(b, 11.5)
	if err != nil || ok {
		t.Fatalf("Allow(11.5) = %v, %v; want false, nil", ok, err)
	}
	if b.LastAt != 11.5 || b.Tokens != 0.75 {
		t.Fatalf("after second reject: LastAt=%v Tokens=%v, want 11.5 and 0.75", b.LastAt, b.Tokens)
	}
	ok, err = Allow(b, 12)
	if err != nil || !ok {
		t.Fatalf("Allow(12) = %v, %v; want true, nil", ok, err)
	}
	if b.Tokens != 0 {
		t.Fatalf("tokens = %v, want 0", b.Tokens)
	}
}

func TestAllowCapsAtBurst(t *testing.T) {
	b := &Bucket{Tokens: 0, Rate: 10, Burst: 3, LastAt: 0}
	ok, err := Allow(b, 100)
	if err != nil || !ok {
		t.Fatalf("Allow(100) = %v, %v; want true, nil", ok, err)
	}
	if b.Tokens != 2 {
		t.Fatalf("tokens = %v, want 2 after cap at burst 3", b.Tokens)
	}
}

func TestAllowRejectsInvalidState(t *testing.T) {
	if _, err := Allow(nil, 0); !errors.Is(err, ErrInvalidLimiter) {
		t.Fatalf("nil bucket = %v, want ErrInvalidLimiter", err)
	}
	if _, err := Allow(&Bucket{Rate: 0, Burst: 1}, 0); !errors.Is(err, ErrInvalidLimiter) {
		t.Fatalf("zero rate = %v, want ErrInvalidLimiter", err)
	}
	if _, err := Allow(&Bucket{Rate: 1, Burst: 0}, 0); !errors.Is(err, ErrInvalidLimiter) {
		t.Fatalf("zero burst = %v, want ErrInvalidLimiter", err)
	}
	if _, err := Allow(&Bucket{Tokens: -1, Rate: 1, Burst: 1}, 0); !errors.Is(err, ErrInvalidLimiter) {
		t.Fatalf("negative tokens = %v, want ErrInvalidLimiter", err)
	}
	b := &Bucket{Tokens: 0.5, Rate: 1, Burst: 1, LastAt: 10}
	if _, err := Allow(b, 9.5); !errors.Is(err, ErrLimiterClockBackwards) {
		t.Fatalf("backwards clock = %v, want ErrLimiterClockBackwards", err)
	}
	if b.LastAt != 10 || b.Tokens != 0.5 {
		t.Fatalf("error mutated bucket: LastAt=%v Tokens=%v", b.LastAt, b.Tokens)
	}
}

cd .local/checks/webapp-lab && go test ./internal/activity/ -run 'TestAllow' -count=1 -v

=== RUN   TestAllowRefillsOverFractionalTime
--- PASS: TestAllowRefillsOverFractionalTime (0.00s)
=== RUN   TestAllowRejectKeepsRemainderAndAdvancesClock
--- PASS: TestAllowRejectKeepsRemainderAndAdvancesClock (0.00s)
=== RUN   TestAllowCapsAtBurst
--- PASS: TestAllowCapsAtBurst (0.00s)
=== RUN   TestAllowRejectsInvalidState
--- PASS: TestAllowRejectsInvalidState (0.00s)
PASS
ok  	webapp/internal/activity	0.264s