Go语言分布式限流实现:令牌桶集群同步与Redis滑动窗口方案

单机限流的瓶颈

Go语言中单机限流常用golang.org/x/time/rate包的令牌桶算法。但当服务部署为多实例集群时,单机限流无法控制全局总速率。假设API限流1000 QPS,部署5个实例,每实例配200 QPS,在负载不均时某些实例可能先触达阈值而其他实例仍有余量,全局QPS上限形同虚设。

Redis滑动窗口限流

分布式限流最直接的方式是将计数器放到Redis。滑动窗口比固定窗口更平滑,避免窗口边界的突发流量。基于Redis Sorted Set的实现:

package ratelimit

import (
	"context"
	"fmt"
	"time"

	"github.com/redis/go-redis/v9"
)

type SlidingWindowLimiter struct {
	client *redis.Client
	limit  int64
	window time.Duration
}

func NewSlidingWindowLimiter(client *redis.Client, limit int64, window time.Duration) *SlidingWindowLimiter {
	return &SlidingWindowLimiter{client: client, limit: limit, window: window}
}

func (l *SlidingWindowLimiter) Allow(ctx context.Context, key string) (bool, error) {
	now := time.Now().UnixMilli()
	windowStart := now - l.window.Milliseconds()
	pipe := l.client.Pipeline()

	// 移除窗口外的记录
	pipe.ZRemRangeByScore(ctx, key, "0", fmt.Sprintf("%d", windowStart))
	// 添加当前请求
	pipe.ZAdd(ctx, key, redis.Z{Score: float64(now), Member: now})
	// 计数
	countCmd := pipe.ZCard(ctx, key)
	// 设置过期
	pipe.Expire(ctx, key, l.window+time.Second)

	if _, err := pipe.Exec(ctx); err != nil {
		return false, err
	}
	return countCmd.Val() <= l.limit, nil
}

ZSet的Score存储请求时间戳,Member用唯一ID防重复。每次请求执行三步:清理过期记录、写入新记录、检查窗口内总数。Pipeline将多命令打包一次RTT完成。

令牌桶集群同步

滑动窗口方案简洁但每请求一次Redis调用,高并发场景下Redis成为瓶颈。令牌桶方案的思路是每个实例从Redis定期领取一批令牌,本地消耗,减少Redis交互频率:

type ClusterTokenBucket struct {
	client    *redis.Client
	key       string
	rate      float64 // 每秒填充令牌数
	burst     int64   // 桶容量
	local     *rate.Limiter
	lastSync  time.Time
	syncMutex sync.Mutex
}

func (b *ClusterTokenBucket) Allow(ctx context.Context) bool {
	if b.local.Allow() {
		return true
	}
	// 本地令牌不足,从Redis领取
	b.syncMutex.Lock()
	defer b.syncMutex.Unlock()

	if time.Since(b.lastSync) < 100*time.Millisecond {
		return false
	}

	claimed := b.claimFromRedis(ctx)
	if claimed > 0 {
		b.local = rate.NewLimiter(rate.Limit(b.rate), int(claimed))
		b.lastSync = time.Now()
		return b.local.Allow()
	}
	return false
}

func (b *ClusterTokenBucket) claimFromRedis(ctx context.Context) int64 {
	// Lua脚本保证原子性
	script := redis.NewScript(`
	local current = tonumber(redis.call("GET", KEYS[1]) or "0")
	local claim = math.min(ARGV[1], current + math.floor(ARGV[2] * (now - last_sync)))
	redis.call("SET", KEYS[1], current + refill - claim)
	return claim
	`)
	result, err := script.Run(ctx, b.client, []string{b.key}, b.burst/4, b.rate).Int64()
	if err != nil {
		return 0
	}
	return result
}

核心是用Lua脚本在Redis端原子地执行”计算补充令牌+扣减领取令牌”,避免竞态条件。每次领取桶容量的1/4,本地消耗完再领取下一批。

Lua脚本保证原子性

分布式限流中,多客户端并发扣减令牌需要原子操作。Redis的单线程模型保证了Lua脚本的原子执行。完整的限流Lua脚本:

local key = KEYS[1]
local limit = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local member = ARGV[4]

redis.call("ZREMRANGEBYSCORE", key, "0", tostring(now - window))
local count = redis.call("ZCARD", key)
if count < limit then
    redis.call("ZADD", key, now, member)
    redis.call("EXPIRE", key, math.ceil(window / 1000) + 1)
    return 1
else
    return 0
end

两种方案的性能对比

在10000 QPS压力测试中,滑动窗口方案Redis QPS约为请求QPS的3倍(Pipeline中3条命令),令牌桶方案Redis QPS约为请求QPS的0.1倍(每10秒同步一次)。滑动窗口内存开销与窗口内请求数成正比,令牌桶方案Redis只存一个计数器。

选择建议:API网关、计费系统等需要精确计数的场景用滑动窗口;微服务内部限流、容忍轻微超发的场景用令牌桶集群同步。

原创文章,作者:小编,如若转载,请注明出处:https://www.yunthe.com/go-yu-yan-fen-bu-shi-xian-liu-shi-xian-ling-pai-tong-ji-qun/

(0)
小编小编
上一篇 1天前
下一篇 1天前

相关推荐