跳到主要内容

实现限流器

问题​

如何用 Go 实现常见的限流算法？

答案​

四种限流算法对比​

算法原理优点缺点
固定窗口时间窗口内计数简单窗口边界突发
滑动窗口滑动时间段计数平滑实现复杂
漏桶恒速流出流量均匀无法应对突发
令牌桶恒速放令牌允许突发参数需调优

1. 令牌桶（推荐）​

// 手写令牌桶
type TokenBucket struct {
    rate       float64   // 每秒产生的令牌数
    capacity   float64   // 桶容量
    tokens     float64   // 当前令牌数
    lastRefill time.Time
    mu         sync.Mutex
}

func NewTokenBucket(rate, capacity float64) *TokenBucket {
    return &TokenBucket{
        rate:       rate,
        capacity:   capacity,
        tokens:     capacity, // 初始满桶
        lastRefill: time.Now(),
    }
}

func (tb *TokenBucket) Allow() bool {
    tb.mu.Lock()
    defer tb.mu.Unlock()

    now := time.Now()
    // 补充令牌
    elapsed := now.Sub(tb.lastRefill).Seconds()
    tb.tokens += elapsed * tb.rate
    if tb.tokens > tb.capacity {
        tb.tokens = tb.capacity
    }
    tb.lastRefill = now

    // 消耗令牌
    if tb.tokens >= 1 {
        tb.tokens--
        return true
    }
    return false
}

标准库方案：

import "golang.org/x/time/rate"

// 每秒 100 个请求，允许突发 200
limiter := rate.NewLimiter(100, 200)

// 非阻塞判断
if limiter.Allow() {
    handleRequest()
}

// 阻塞等待令牌
limiter.Wait(ctx)

// 预约未来的令牌
reservation := limiter.Reserve()
time.Sleep(reservation.Delay())

2. 滑动窗口​

type SlidingWindow struct {
    mu       sync.Mutex
    window   time.Duration
    limit    int
    requests []time.Time
}

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

func (sw *SlidingWindow) Allow() bool {
    sw.mu.Lock()
    defer sw.mu.Unlock()

    now := time.Now()
    windowStart := now.Add(-sw.window)

    // 移除窗口外的记录
    idx := 0
    for idx < len(sw.requests) && sw.requests[idx].Before(windowStart) {
        idx++
    }
    sw.requests = sw.requests[idx:]

    // 判断是否超限
    if len(sw.requests) >= sw.limit {
        return false
    }

    sw.requests = append(sw.requests, now)
    return true
}

3. 漏桶​

type LeakyBucket struct {
    rate     float64   // 每秒漏出速率
    capacity float64   // 桶容量
    water    float64   // 当前水量
    lastLeak time.Time
    mu       sync.Mutex
}

func (lb *LeakyBucket) Allow() bool {
    lb.mu.Lock()
    defer lb.mu.Unlock()

    now := time.Now()
    // 漏水
    elapsed := now.Sub(lb.lastLeak).Seconds()
    lb.water -= elapsed * lb.rate
    if lb.water < 0 {
        lb.water = 0
    }
    lb.lastLeak = now

    // 加水
    if lb.water >= lb.capacity {
        return false
    }
    lb.water++
    return true
}

Gin 中间件​

func RateLimitMiddleware(rps float64, burst int) gin.HandlerFunc {
    limiter := rate.NewLimiter(rate.Limit(rps), burst)
    return func(c *gin.Context) {
        if !limiter.Allow() {
            c.JSON(429, gin.H{"error": "too many requests"})
            c.Abort()
            return
        }
        c.Next()
    }
}

// 按 IP 限流
func PerIPRateLimit(rps float64, burst int) gin.HandlerFunc {
    var mu sync.Mutex
    limiters := make(map[string]*rate.Limiter)

    return func(c *gin.Context) {
        ip := c.ClientIP()
        mu.Lock()
        lim, ok := limiters[ip]
        if !ok {
            lim = rate.NewLimiter(rate.Limit(rps), burst)
            limiters[ip] = lim
        }
        mu.Unlock()

        if !lim.Allow() {
            c.JSON(429, gin.H{"error": "rate limit exceeded"})
            c.Abort()
            return
        }
        c.Next()
    }
}

常见面试问题​

Q1: 令牌桶和漏桶的区别？​

答案：

  • 令牌桶：允许突发流量（桶满时可一次消耗多个令牌）
  • 漏桶：流出速率恒定，不允许突发
  • 大多数场景用令牌桶更合适

Q2: 分布式限流怎么做？​

答案：

  • Redis + Lua 脚本实现滑动窗口/令牌桶
  • 网关层限流（如 Sentinel、APISIX 插件）
  • 参考 系统设计 - 限流器

相关链接​