2.2 太牛了！一个限流算法竟然能保护整个平台？

2.2 太牛了！一个限流算法竟然能保护整个平台？

在高并发的分布式系统中，限流是一个非常重要的保护机制。它能够防止系统被突发的大流量冲垮，保护系统的稳定性和可用性。本节将深入探讨几种常用的限流算法，并展示如何在通知平台中实现高效的限流保护。

限流算法概述

限流算法是用来控制系统请求处理速率的技术，主要目的是保护系统免受过多请求的冲击。常见的限流算法包括：

1. 计数器算法：最简单的限流算法
2. 滑动窗口算法：对计数器算法的改进
3. 漏桶算法：以固定速率处理请求
4. 令牌桶算法：允许一定程度的突发流量

计数器算法

计数器算法是最简单的限流算法，它在一定时间窗口内统计请求次数，当请求次数超过阈值时拒绝后续请求。

// CounterRateLimiter 计数器限流器
type CounterRateLimiter struct {
    // 限制数量
    limit int64
    
    // 时间窗口
    window time.Duration
    
    // 计数器映射
    counters map[string]*Counter
    
    // 互斥锁
    mutex sync.RWMutex
}

// Counter 计数器
type Counter struct {
    // 计数
    count int64
    
    // 过期时间
    expireAt time.Time
}

// NewCounterRateLimiter 创建计数器限流器
func NewCounterRateLimiter(limit int64, window time.Duration) *CounterRateLimiter {
    return &CounterRateLimiter{
        limit:    limit,
        window:   window,
        counters: make(map[string]*Counter),
    }
}

// Allow 是否允许通过
func (crl *CounterRateLimiter) Allow(key string) bool {
    crl.mutex.Lock()
    defer crl.mutex.Unlock()
    
    now := time.Now()
    
    // 获取计数器
    counter, exists := crl.counters[key]
    if !exists || counter.expireAt.Before(now) {
        // 创建新的计数器
        crl.counters[key] = &Counter{
            count:    1,
            expireAt: now.Add(crl.window),
        }
        return true
    }
    
    // 增加计数
    if counter.count < crl.limit {
        counter.count++
        return true
    }
    
    return false
}

// 清理过期计数器
func (crl *CounterRateLimiter) cleanup() {
    crl.mutex.Lock()
    defer crl.mutex.Unlock()
    
    now := time.Now()
    for key, counter := range crl.counters {
        if counter.expireAt.Before(now) {
            delete(crl.counters, key)
        }
    }
}

计数器算法的优点是实现简单，但缺点是存在临界问题。例如，如果限制1秒内最多100个请求，在第1秒的后半段和第2秒的前半段可能分别通过100个请求，导致短时间内通过200个请求。

滑动窗口算法

滑动窗口算法是对计数器算法的改进，它将时间窗口划分为多个小的时间段，更加精确地控制请求速率。

// SlidingWindowRateLimiter 滑动窗口限流器
type SlidingWindowRateLimiter struct {
    // 限制数量
    limit int64
    
    // 时间窗口
    window time.Duration
    
    // 窗口分段数
    segments int64
    
    // 滑动窗口映射
    windows map[string]*SlidingWindow
    
    // 互斥锁
    mutex sync.RWMutex
}

// SlidingWindow 滑动窗口
type SlidingWindow struct {
    // 请求记录
    requests []int64
    
    // 每个段的时间长度
    segmentDuration time.Duration
    
    // 最后更新时间
    lastUpdate time.Time
}

// NewSlidingWindowRateLimiter 创建滑动窗口限流器
func NewSlidingWindowRateLimiter(limit int64, window time.Duration, segments int64) *SlidingWindowRateLimiter {
    return &SlidingWindowRateLimiter{
        limit:    limit,
        window:   window,
        segments: segments,
        windows:  make(map[string]*SlidingWindow),
    }
}

// Allow 是否允许通过
func (swrl *SlidingWindowRateLimiter) Allow(key string) bool {
    swrl.mutex.Lock()
    defer swrl.mutex.Unlock()
    
    now := time.Now()
    
    // 获取滑动窗口
    window, exists := swrl.windows[key]
    if !exists {
        // 创建新的滑动窗口
        window = &SlidingWindow{
            requests:        make([]int64, swrl.segments),
            segmentDuration: swrl.window / time.Duration(swrl.segments),
            lastUpdate:      now,
        }
        swrl.windows[key] = window
    }
    
    // 计算当前段索引
    segmentIndex := (now.UnixNano() / int64(window.segmentDuration)) % swrl.segments
    
    // 清理过期段
    swrl.cleanupWindow(window, now)
    
    // 计算当前窗口内的请求数
    totalCount := int64(0)
    for _, count := range window.requests {
        totalCount += count
    }
    
    // 检查是否超过限制
    if totalCount >= swrl.limit {
        return false
    }
    
    // 增加当前段的计数
    window.requests[segmentIndex]++
    window.lastUpdate = now
    
    return true
}

// cleanupWindow 清理过期段
func (swrl *SlidingWindowRateLimiter) cleanupWindow(window *SlidingWindow, now time.Time) {
    // 计算过期时间
    expireTime := now.Add(-swrl.window)
    expireSegmentIndex := (expireTime.UnixNano() / int64(window.segmentDuration)) % swrl.segments
    
    // 清理过期段
    for i := int64(0); i < swrl.segments; i++ {
        segmentTime := time.Unix(0, (expireSegmentIndex+i)*int64(window.segmentDuration))
        if segmentTime.Before(expireTime) {
            window.requests[(expireSegmentIndex+i)%swrl.segments] = 0
        }
    }
}

滑动窗口算法比计数器算法更加精确，但实现复杂度也更高。

漏桶算法

漏桶算法是一种恒定速率处理请求的算法，它将请求看作是流入漏桶的水，而桶以恒定的速率漏水，当桶满时拒绝新请求。

// LeakyBucketRateLimiter 漏桶限流器
type LeakyBucketRateLimiter struct {
    // 桶容量
    capacity int64
    
    // 漏水速率（每秒）
    leakRate float64
    
    // 桶映射
    buckets map[string]*LeakyBucket
    
    // 互斥锁
    mutex sync.RWMutex
}

// LeakyBucket 漏桶
type LeakyBucket struct {
    // 当前水量
    water float64
    
    // 最大容量
    capacity int64
    
    // 漏水速率（每秒）
    leakRate float64
    
    // 上次更新时间
    lastUpdate time.Time
    
    // 互斥锁
    mutex sync.Mutex
}

// NewLeakyBucketRateLimiter 创建漏桶限流器
func NewLeakyBucketRateLimiter(capacity int64, leakRate float64) *LeakyBucketRateLimiter {
    return &LeakyBucketRateLimiter{
        capacity: capacity,
        leakRate: leakRate,
        buckets:  make(map[string]*LeakyBucket),
    }
}

// Allow 是否允许通过
func (lbrl *LeakyBucketRateLimiter) Allow(key string) bool {
    lbrl.mutex.Lock()
    defer lbrl.mutex.Unlock()
    
    // 获取漏桶
    bucket, exists := lbrl.buckets[key]
    if !exists {
        // 创建新的漏桶
        bucket = &LeakyBucket{
            water:    0,
            capacity: lbrl.capacity,
            leakRate: lbrl.leakRate,
            lastUpdate: time.Now(),
        }
        lbrl.buckets[key] = bucket
    }
    
    return bucket.Allow()
}

// Allow 是否允许通过
func (lb *LeakyBucket) Allow() bool {
    lb.mutex.Lock()
    defer lb.mutex.Unlock()
    
    now := time.Now()
    
    // 计算漏水量
    elapsed := now.Sub(lb.lastUpdate).Seconds()
    leaked := elapsed * lb.leakRate
    lb.lastUpdate = now
    
    // 更新水量
    lb.water -= leaked
    if lb.water < 0 {
        lb.water = 0
    }
    
    // 检查是否可以加入新水
    if lb.water+1 <= float64(lb.capacity) {
        lb.water++
        return true
    }
    
    return false
}

漏桶算法的优点是能够平滑处理请求，但缺点是不能应对突发流量。

令牌桶算法

令牌桶算法是一种更加灵活的限流算法，它允许一定程度的突发流量。系统以恒定速率生成令牌放入桶中，请求需要获取令牌才能被处理，当桶中没有令牌时拒绝请求。

// TokenBucketRateLimiter 令牌桶限流器
type TokenBucketRateLimiter struct {
    // 令牌生成速率（每秒）
    rate float64
    
    // 桶容量
    capacity int64
    
    // 桶映射
    buckets map[string]*TokenBucket
    
    // 互斥锁
    mutex sync.RWMutex
}

// TokenBucket 令牌桶
type TokenBucket struct {
    // 当前令牌数
    tokens float64
    
    // 最大令牌数
    capacity int64
    
    // 令牌生成速率（每秒）
    rate float64
    
    // 上次更新时间
    lastUpdate time.Time
    
    // 互斥锁
    mutex sync.Mutex
}

// NewTokenBucketRateLimiter 创建令牌桶限流器
func NewTokenBucketRateLimiter(rate float64, capacity int64) *TokenBucketRateLimiter {
    return &TokenBucketRateLimiter{
        rate:     rate,
        capacity: capacity,
        buckets:  make(map[string]*TokenBucket),
    }
}

// Allow 是否允许通过
func (tbrl *TokenBucketRateLimiter) Allow(key string) bool {
    return tbrl.AllowN(key, 1)
}

// AllowN 是否允许通过N个请求
func (tbrl *TokenBucketRateLimiter) AllowN(key string, n int64) bool {
    tbrl.mutex.Lock()
    defer tbrl.mutex.Unlock()
    
    // 获取令牌桶
    bucket, exists := tbrl.buckets[key]
    if !exists {
        // 创建新的令牌桶
        bucket = &TokenBucket{
            tokens:    float64(tbrl.capacity),
            capacity:  tbrl.capacity,
            rate:      tbrl.rate,
            lastUpdate: time.Now(),
        }
        tbrl.buckets[key] = bucket
    }
    
    return bucket.AllowN(n)
}

// WaitUntil 等待直到允许通过
func (tbrl *TokenBucketRateLimiter) WaitUntil(ctx context.Context, key string) error {
    for {
        if tbrl.Allow(key) {
            return nil
        }
        
        // 等待一小段时间再重试
        select {
        case <-time.After(10 * time.Millisecond):
            // 继续重试
        case <-ctx.Done():
            return ctx.Err()
        }
    }
}

// AllowN 是否允许通过N个请求
func (tb *TokenBucket) AllowN(n int64) bool {
    tb.mutex.Lock()
    defer tb.mutex.Unlock()
    
    now := time.Now()
    
    // 计算新增令牌数
    elapsed := now.Sub(tb.lastUpdate).Seconds()
    newTokens := elapsed * tb.rate
    tb.lastUpdate = now
    
    // 更新令牌数
    tb.tokens += newTokens
    if tb.tokens > float64(tb.capacity) {
        tb.tokens = float64(tb.capacity)
    }
    
    // 检查是否有足够的令牌
    if tb.tokens >= float64(n) {
        tb.tokens -= float64(n)
        return true
    }
    
    return false
}

令牌桶算法的优点是既能够控制平均速率，又能够应对一定程度的突发流量，是实际应用中最常用的限流算法之一。

分布式限流

在分布式系统中，单机限流无法满足需求，我们需要实现分布式限流。常见的分布式限流方案有：

1. 基于Redis的限流：使用Redis的原子操作实现限流
2. 基于令牌桶的分布式限流：将令牌桶存储在共享存储中
3. 基于漏桶的分布式限流：将漏桶状态存储在共享存储中

// RedisTokenBucketRateLimiter 基于Redis的令牌桶限流器
type RedisTokenBucketRateLimiter struct {
    // Redis客户端
    redisClient RedisClient
    
    // 令牌生成速率（每秒）
    rate float64
    
    // 桶容量
    capacity int64
    
    // 限流键前缀
    keyPrefix string
}

// NewRedisTokenBucketRateLimiter 创建基于Redis的令牌桶限流器
func NewRedisTokenBucketRateLimiter(
    redisClient RedisClient,
    rate float64,
    capacity int64,
    keyPrefix string,
) *RedisTokenBucketRateLimiter {
    return &RedisTokenBucketRateLimiter{
        redisClient: redisClient,
        rate:        rate,
        capacity:    capacity,
        keyPrefix:   keyPrefix,
    }
}

// Allow 是否允许通过
func (rtbrl *RedisTokenBucketRateLimiter) Allow(key string) bool {
    return rtbrl.AllowN(key, 1)
}

// AllowN 是否允许通过N个请求
func (rtbrl *RedisTokenBucketRateLimiter) AllowN(key string, n int64) bool {
    fullKey := rtbrl.keyPrefix + key
    
    // 使用Lua脚本保证原子性
    script := `
        local tokens_key = KEYS[1]
        local rate = tonumber(ARGV[1])
        local capacity = tonumber(ARGV[2])
        local requested = tonumber(ARGV[3])
        local now = tonumber(ARGV[4])
        
        local last_tokens = redis.call("GET", tokens_key .. ":tokens")
        if not last_tokens then
            last_tokens = capacity
        end
        
        local last_refreshed = redis.call("GET", tokens_key .. ":refreshed")
        if not last_refreshed then
            last_refreshed = 0
        end
        
        local delta = math.max(0, now - last_refreshed)
        local filled_tokens = math.min(capacity, last_tokens + (delta * rate))
        
        if filled_tokens >= requested then
            local new_tokens = filled_tokens - requested
            redis.call("SET", tokens_key .. ":tokens", new_tokens)
            redis.call("SET", tokens_key .. ":refreshed", now)
            return 1
        else
            redis.call("SET", tokens_key .. ":tokens", filled_tokens)
            redis.call("SET", tokens_key .. ":refreshed", now)
            return 0
        end
    `
    
    now := time.Now().Unix()
    result, err := rtbrl.redisClient.Eval(script, []string{fullKey}, 
        []interface{}{rtbrl.rate, rtbrl.capacity, n, now})
    if err != nil {
        // 出错时默认允许通过，避免因限流服务故障导致业务不可用
        return true
    }
    
    return result.(int64) == 1
}

限流策略管理

在实际应用中，我们需要根据不同业务场景配置不同的限流策略：

// RateLimitStrategy 限流策略
type RateLimitStrategy struct {
    // 策略ID
    ID string `json:"id" db:"id"`
    
    // 策略名称
    Name string `json:"name" db:"name"`
    
    // 策略类型
    Type RateLimitType `json:"type" db:"type"`
    
    // 配置参数
    Config RateLimitConfig `json:"config" db:"config"`
    
    // 适用范围
    Scope RateLimitScope `json:"scope" db:"scope"`
    
    // 创建时间
    CreatedAt time.Time `json:"created_at" db:"created_at"`
    
    // 更新时间
    UpdatedAt time.Time `json:"updated_at" db:"updated_at"`
}

// RateLimitType 限流类型
type RateLimitType string

const (
    // RateLimitTypeTokenBucket 令牌桶
    RateLimitTypeTokenBucket RateLimitType = "token_bucket"
    
    // RateLimitTypeLeakyBucket 漏桶
    RateLimitTypeLeakyBucket RateLimitType = "leaky_bucket"
    
    // RateLimitTypeSlidingWindow 滑动窗口
    RateLimitTypeSlidingWindow RateLimitType = "sliding_window"
)

// RateLimitConfig 限流配置
type RateLimitConfig struct {
    // 速率（每秒请求数）
    Rate float64 `json:"rate"`
    
    // 容量
    Capacity int64 `json:"capacity"`
    
    // 时间窗口
    Window time.Duration `json:"window"`
    
    // 窗口分段数
    Segments int64 `json:"segments"`
}

// RateLimitScope 限流范围
type RateLimitScope struct {
    // 业务方ID列表
    BizIDs []string `json:"biz_ids"`
    
    // API路径列表
    APIPaths []string `json:"api_paths"`
    
    // IP地址列表
    IPs []string `json:"ips"`
}

// RateLimitStrategyManager 限流策略管理器
type RateLimitStrategyManager struct {
    // 数据库连接
    db *sql.DB
    
    // 缓存
    cache Cache
    
    // 限流器工厂
    rateLimiterFactory RateLimiterFactory
}

// RateLimiterFactory 限流器工厂接口
type RateLimiterFactory interface {
    // CreateRateLimiter 创建限流器
    CreateRateLimiter(strategy *RateLimitStrategy) RateLimiter
}

// DefaultRateLimiterFactory 默认限流器工厂
type DefaultRateLimiterFactory struct {
    // Redis客户端（用于分布式限流）
    redisClient RedisClient
}

// CreateRateLimiter 创建限流器
func (drlf *DefaultRateLimiterFactory) CreateRateLimiter(strategy *RateLimitStrategy) RateLimiter {
    switch strategy.Type {
    case RateLimitTypeTokenBucket:
        // 如果配置了Redis客户端，则使用分布式令牌桶
        if drlf.redisClient != nil {
            return NewRedisTokenBucketRateLimiter(
                drlf.redisClient,
                strategy.Config.Rate,
                strategy.Config.Capacity,
                "rate_limit:",
            )
        }
        return NewTokenBucketRateLimiter(strategy.Config.Rate, strategy.Config.Capacity)
        
    case RateLimitTypeLeakyBucket:
        return NewLeakyBucketRateLimiter(strategy.Config.Capacity, strategy.Config.Rate)
        
    case RateLimitTypeSlidingWindow:
        return NewSlidingWindowRateLimiter(
            strategy.Config.Capacity,
            strategy.Config.Window,
            strategy.Config.Segments,
        )
        
    default:
        // 默认使用令牌桶
        return NewTokenBucketRateLimiter(strategy.Config.Rate, strategy.Config.Capacity)
    }
}

使用示例

// 初始化限流器
redisClient := NewRedisClient() // 假设的Redis客户端
rateLimiterFactory := &DefaultRateLimiterFactory{
    redisClient: redisClient,
}

// 创建限流策略管理器
strategyManager := NewRateLimitStrategyManager(db, cache, rateLimiterFactory)

// 创建令牌桶限流策略
strategy := &RateLimitStrategy{
    ID:   "default_strategy",
    Name: "默认限流策略",
    Type: RateLimitTypeTokenBucket,
    Config: RateLimitConfig{
        Rate:     100,  // 每秒100个请求
        Capacity: 200,  // 桶容量200
    },
    Scope: RateLimitScope{
        BizIDs: []string{"*"}, // 适用于所有业务方
    },
}

// 应用限流策略
rateLimiter := rateLimiterFactory.CreateRateLimiter(strategy)

// 在请求处理中使用限流器
func handleRequest(w http.ResponseWriter, r *http.Request) {
    // 获取业务方ID
    bizID := getBizIDFromRequest(r)
    
    // 检查是否允许通过
    if !rateLimiter.Allow(bizID) {
        http.Error(w, "请求过于频繁", http.StatusTooManyRequests)
        return
    }
    
    // 处理业务逻辑
    // ...
    
    w.WriteHeader(http.StatusOK)
    json.NewEncoder(w).Encode(map[string]interface{}{
        "code": 0,
        "msg":  "success",
    })
}

总结

通过合理的限流算法和策略，我们能够有效地保护系统免受流量冲击，确保系统的稳定性和可用性。不同的限流算法有各自的优缺点和适用场景：

1. 计数器算法：实现简单，但精度较低
2. 滑动窗口算法：精度较高，但实现复杂
3. 漏桶算法：能够平滑处理请求，但不能应对突发流量
4. 令牌桶算法：既能够控制平均速率，又能够应对一定程度的突发流量

在实际应用中，我们通常会结合使用多种限流算法，并通过分布式限流来应对大规模系统的限流需求。同时，还需要建立完善的限流策略管理体系，根据不同业务场景配置不同的限流策略。

在下一节中，我们将探讨如何设计API安全策略，进一步保护系统的安全性。