Go语言实现请求频率限制的方法实践

type CounterLimiter struct {
	mu       sync.Mutex
	count    int
	interval time.Duration
	maxReq   int
	lastReset time.Time
}
func NewCounterLimiter(interval time.Duration, maxReq int) *CounterLimiter {
	return &CounterLimiter{
		interval:  interval,
		maxReq:    maxReq,
		lastReset: time.Now(),
	}
}
func (c *CounterLimiter) Allow() bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	// 检查是否需要重置计数器
	if time.Since(c.lastReset) > c.interval {
		c.count = 0
		c.lastReset = time.Now()
	}
	// 检查是否超过限制
	if c.count >= c.maxReq {
		return false
	}
	c.count++
	return true
}
// HTTP中间件示例
func RateLimitMiddleware(limiter *CounterLimiter) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if !limiter.Allow() {
				w.Header().Set("Retry-After", "60")
				http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
				return
			}
			next.ServeHTTP(w, r)
		})
	}
}

const luaScript = `
local key = KEYS[1]
local now = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local max = tonumber(ARGV[3])
redis.call('ZREMRANGEBYSCORE', key, 0, now - window)
local current = redis.call('ZCARD', key)
if current >= max then
    return 0
end
redis.call('ZADD', key, now, now)
redis.call('EXPIRE', key, math.ceil(window/1000))
return 1
`
type RedisLimiter struct {
	client   *redis.Client
	script   *redis.Script
	maxReq   int
	window   time.Duration
}
func NewRedisLimiter(client *redis.Client, window time.Duration, maxReq int) *RedisLimiter {
	return &RedisLimiter{
		client: client,
		script: redis.NewScript(luaScript),
		maxReq: maxReq,
		window: window,
	}
}
func (r *RedisLimiter) Allow(userID string) bool {
	ctx := context.Background()
	key := fmt.Sprintf("rate_limit:%s", userID)
	now := time.Now().UnixMilli()
	result, err := r.script.Run(ctx, r.client, 
		[]string{key}, now, r.window.Milliseconds(), r.maxReq).Int()
	return err == nil && result == 1
}

type TokenBucket struct {
	capacity    int           // 桶容量
	tokens      int           // 当前令牌数
	fillRate    time.Duration // 添加令牌间隔
	lastRefill  time.Time     // 上次添加时间
	mu          sync.Mutex
}
func NewTokenBucket(capacity int, rate time.Duration) *TokenBucket {
	return &TokenBucket{
		capacity:   capacity,
		tokens:     capacity,
		fillRate:   rate,
		lastRefill: time.Now(),
	}
}
func (b *TokenBucket) Allow() bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	// 补充令牌
	now := time.Now()
	elapsed := now.Sub(b.lastRefill)
	newTokens := int(elapsed / b.fillRate)
	if newTokens > 0 {
		b.tokens += newTokens
		if b.tokens > b.capacity {
			b.tokens = b.capacity
		}
		b.lastRefill = now
	}
	// 检查令牌是否足够
	if b.tokens <= 0 {
		return false
	}
	b.tokens--
	return true
}

func main() {
	r := chi.NewRouter()
	
	// 创建限流器：每分钟1000次
	limiter := tollbooth.NewLimiter(1000/60.0, nil)
	limiter.SetIPLookups([]string{"X-Real-IP", "RemoteAddr", "X-Forwarded-For"})
	
	// 应用中间件
	r.Use(tollbooth_chi.LimitHandler(limiter))
	
	r.Get("/api/protected", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("Protected content"))
	})
	
	http.ListenAndServe(":8080", r)
}