Go语言限流器实战：令牌桶与滑动窗口算法实现及中间件集成

小编 • • 后端开发

Go语言在高并发服务场景中对流量控制有强需求。限流器（Rate Limiter）保护后端服务不被突发流量压垮，是微服务治理的基础组件。Go标准库golang.org/x/time/rate提供了生产级的令牌桶实现，但实际业务中往往需要更灵活的多维度限流策略。本文从算法原理到工程实现，覆盖令牌桶、滑动窗口两种限流算法的Go语言实战。

令牌桶算法原理与标准库使用

令牌桶（Token Bucket）的核心机制：系统以固定速率向桶中放入令牌，桶有最大容量上限。每个请求到来时从桶中取走一个令牌，取到则放行，取不到则拒绝或排队。令牌桶的优势在于允许一定程度的突发流量——当桶中积攒了令牌时，可以实现短时高于平均速率的吞吐。

Go标准库的rate.Limiter使用：

package main

import (
    "context"
    "fmt"
    "time"
    "golang.org/x/time/rate"
)

func main() {
    // 创建限流器：每秒10个令牌，桶容量20
    // rate.Limit(10) = 每秒10个请求
    // 20 = 突发容量，允许一次最多20个请求通过
    limiter := rate.NewLimiter(rate.Limit(10), 20)

    for i := 0; i < 50; i++ {
        // Wait方法：取不到令牌时阻塞等待
        err := limiter.Wait(context.Background())
        if err != nil {
            fmt.Printf("请求 %d 被拒绝: %v\n", i, err)
            continue
        }
        fmt.Printf("请求 %d 通过: %s\n", i, time.Now().Format("15:04:05.000"))
    }
}

Allow方法适用于非阻塞场景，直接返回是否允许通过：

func handleRequest(limiter *rate.Limiter) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        // Allow消耗一个令牌，返回false说明限流
        if !limiter.Allow() {
            w.WriteHeader(http.StatusTooManyRequests)
            w.Write([]byte(`{"code":429,"message":"请求过于频繁"}`))
            return
        }
        w.WriteHeader(http.StatusOK)
        w.Write([]byte(`{"code":0,"message":"ok"}`))
    }
}

基于Redis的分布式令牌桶实现

单机限流器在多实例部署时无法共享状态。基于Redis的分布式限流器通过Lua脚本保证原子性，所有实例共享同一计数桶：

package ratelimit

import (
    "context"
    "fmt"
    "time"
    "github.com/redis/go-redis/v9"
)

type RedisTokenBucket struct {
    client   *redis.Client
    rate     int           // 每秒令牌数
    capacity int           // 桶容量
    key      string        // Redis key前缀
}

func NewRedisTokenBucket(client *redis.Client, rate, capacity int, key string) *RedisTokenBucket {
    return &RedisTokenBucket{
        client:   client,
        rate:     rate,
        capacity: capacity,
        key:      fmt.Sprintf("ratelimit:%s", key),
    }
}

// Lua脚本保证获取令牌的原子性
const tokenBucketScript = `
local key = KEYS[1]
local rate = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local requested = tonumber(ARGV[4])

local data = redis.call('HMGET', key, 'tokens', 'last_time')
local tokens = tonumber(data[1])
local last_time = tonumber(data[2])

-- 初始化
if tokens == nil then
    tokens = capacity
    last_time = now
end

-- 计算自上次请求以来新增的令牌数
local elapsed = now - last_time
local new_tokens = elapsed * rate / 1000
tokens = math.min(capacity, tokens + new_tokens)

-- 判断是否放行
local allowed = 0
if tokens >= requested then
    tokens = tokens - requested
    allowed = 1
end

-- 更新状态
redis.call('HMSET', key, 'tokens', tokens, 'last_time', now)
redis.call('EXPIRE', key, 60)

return allowed
`

func (r *RedisTokenBucket) Allow(ctx context.Context) (bool, error) {
    now := time.Now().UnixMilli()
    result, err := r.client.Eval(ctx, tokenBucketScript, []string{r.key},
        r.rate, r.capacity, now, 1).Int()
    if err != nil {
        return false, err
    }
    return result == 1, nil
}

Lua脚本在Redis中以原子方式执行，避免了”先读取再写入”的竞态条件。EXPIRE设置60秒过期，防止无效key长期占用内存。参数中的now使用毫秒时间戳，保证多实例间时钟基本同步。对于精度要求更高的场景，可以引入NTP时间同步。

滑动窗口限流算法实现

滑动窗口（Sliding Window）相比固定窗口避免了临界点突发问题。核心思路是维护一个时间窗口内的请求计数，窗口随时间滑动。Redis的Sorted Set天然适合实现这个方案：

package ratelimit

import (
    "context"
    "fmt"
    "time"
    "github.com/redis/go-redis/v9"
)

type SlidingWindowLimiter struct {
    client *redis.Client
    key    string
    limit  int64         // 窗口内最大请求数
    window time.Duration // 窗口大小
}

func NewSlidingWindowLimiter(client *redis.Client, key string, limit int64, window time.Duration) *SlidingWindowLimiter {
    return &SlidingWindowLimiter{
        client: client,
        key:    fmt.Sprintf("sliding:%s", key),
        limit:  limit,
        window: window,
    }
}

const slidingWindowScript = `
local key = KEYS[1]
local now = tonumber(ARGV[1])
local window_ms = tonumber(ARGV[2])
local limit = tonumber(ARGV[3])
local member = ARGV[4]

-- 计算窗口起始时间
local window_start = now - window_ms

-- 移除窗口外的旧请求记录
redis.call('ZREMRANGEBYSCORE', key, '-inf', window_start)

-- 获取当前窗口内的请求数
local count = redis.call('ZCARD', key)

if count < limit then
    -- 添加当前请求，score为当前时间戳
    redis.call('ZADD', key, now, member)
    redis.call('EXPIRE', key, math.ceil(window_ms / 1000))
    return 1
else
    return 0
end
`

func (s *SlidingWindowLimiter) Allow(ctx context.Context) (bool, error) {
    now := time.Now().UnixMilli()
    // 使用纳秒时间戳保证member唯一性
    member := fmt.Sprintf("%d-%d", now, time.Now().UnixNano())
    
    result, err := s.client.Eval(ctx, slidingWindowScript, []string{s.key},
        now, s.window.Milliseconds(), s.limit, member).Int()
    if err != nil {
        return false, err
    }
    return result == 1, nil
}

ZREMRANGEBYSCORE清理窗口外的历史请求，ZCARD统计当前窗口内的请求总数。每次请求将当前时间戳作为score写入Sorted Set，member使用纳秒级唯一标识避免覆盖。

HTTP中间件集成与多维度限流

实际业务中，限流维度通常不止一个——全局总流量、单IP频率、单用户频率需要组合控制。以下中间件支持多维度限流：

package middleware

import (
    "context"
    "net/http"
    "strconv"
    "sync"
    "github.com/redis/go-redis/v9"
)

type RateLimitConfig struct {
    GlobalRate int // 全局每秒请求数
    GlobalCap  int
    IPRate     int // 单IP每秒请求数
    IPCap      int
    UserRate   int // 单用户每秒请求数
    UserCap    int
}

type MultiDimensionLimiter struct {
    client   *redis.Client
    config   RateLimitConfig
    limiters sync.Map // 本地缓存，减少Redis访问
}

func NewMultiDimensionLimiter(client *redis.Client, config RateLimitConfig) *MultiDimensionLimiter {
    return &MultiDimensionLimiter{client: client, config: config}
}

func (m *MultiDimensionLimiter) Middleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        ctx := r.Context()
        
        // 全局限流
        globalLimiter := NewRedisTokenBucket(m.client, m.config.GlobalRate, m.config.GlobalCap, "global")
        allowed, err := globalLimiter.Allow(ctx)
        if err != nil || !allowed {
            respondRateLimited(w, "全局限流")
            return
        }

        // IP限流
        ip := getClientIP(r)
        ipLimiter := NewRedisTokenBucket(m.client, m.config.IPRate, m.config.IPCap, "ip:"+ip)
        allowed, err = ipLimiter.Allow(ctx)
        if err != nil || !allowed {
            respondRateLimited(w, "IP限流")
            return
        }

        // 用户限流（已认证用户）
        userID := getUserIDFromContext(ctx)
        if userID != "" {
            userLimiter := NewRedisTokenBucket(m.client, m.config.UserRate, m.config.UserCap, "user:"+userID)
            allowed, err = userLimiter.Allow(ctx)
            if err != nil || !allowed {
                respondRateLimited(w, "用户限流")
                return
            }
        }

        next.ServeHTTP(w, r)
    })
}

func respondRateLimited(w http.ResponseWriter, dimension string) {
    w.Header().Set("Content-Type", "application/json")
    w.Header().Set("Retry-After", "1")
    w.WriteHeader(http.StatusTooManyRequests)
    w.Write([]byte(`{"code":429,"dimension":"` + dimension + `","message":"请求过于频繁，请稍后重试"}`))
}

func getClientIP(r *http.Request) string {
    // 从X-Forwarded-For或X-Real-IP获取
    if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
        return strings.Split(xff, ",")[0]
    }
    if xri := r.Header.Get("X-Real-IP"); xri != "" {
        return xri
    }
    return r.RemoteAddr
}

多维度限流的执行顺序应该从粗到细：先全局限流快速拦截攻击，再IP限流防止单点滥用，最后用户限流保障公平性。每个维度独立计数，互不影响。429响应中携带Retry-After头部，指导客户端退避重试。维度信息写入响应体，便于客户端日志排查限流原因。

对于Go 1.22+项目，可以直接用标准库的ServeMux配合中间件使用：

func main() {
    rdb := redis.NewClient(&redis.Options{Addr: "localhost:6379"})
    limiter := NewMultiDimensionLimiter(rdb, RateLimitConfig{
        GlobalRate: 10000, GlobalCap: 20000,
        IPRate: 100, IPCap: 200,
        UserRate: 50, UserCap: 100,
    })

    mux := http.NewServeMux()
    mux.HandleFunc("POST /api/orders", handleCreateOrder)
    
    // 限流中间件包裹所有路由
    server := &http.Server{
        Addr:    ":8080",
        Handler: limiter.Middleware(mux),
    }
    server.ListenAndServe()
}

原创文章，作者：小编，如若转载，请注明出处：https://www.yunthe.com/go-yu-yan-xian-liu-qi-shi-zhan-ling-pai-tong-yu-hua-dong/

赞 (0)
小编小编
0
上一篇 5小时前
下一篇 5小时前

相关推荐