golang 实现基于redis的并行流量控制（计数锁）

在业务开发中，有时需要对某个操作在整个集群中限制并发度，例如限制大模型对话的并行数。基于redis zset实现计数锁，做个笔记。

关键词：并行流量控制、计数锁

package redisutilimport ("context""fmt""math""time""github.com/go-redis/redis/v9"
)// AcquireZSetLock 借助redis zset数据结构实现分布式计数锁。可用于计数任务运行数，防止超限。返回值：zset大小、释放锁的函数、错误信息
func AcquireZSetLock(ctx context.Context, c redis.Client, key string, element string, zsetMaxSize int,expiresIn time.Duration, syncWait time.Duration) (int, func() error, error) {ctx, cancel := context.WithTimeout(ctx, syncWait)defer cancel()for i := 0; ; i++ {select {case <-ctx.Done(): // 接到取消信号，按插入失败处理return -1, func() error { return nil }, ctx.Err()default:}size, err := insertElementToZsetLock(ctx, c, key, element, zsetMaxSize, expiresIn)if err != nil {second := 0.4 + 0.6*math.Exp(-0.17*float64(i)) // f(i=0) = 1.0; f(i=10) = 0.5096，即第10次就会衰减到0.5096秒second = max(second, 0.5)                      // 最小间隔0.5秒，防止过于频繁的请求time.Sleep(time.Duration(second*1000) * time.Millisecond)}releaseFunc := func() error {result, err := c.ZRem(context.Background(), key, element).Result()if err != nil {return fmt.Errorf("redis zrem error: %v. return=%d", err, result)}return nil}return size, releaseFunc, nil}
}// insertElementToZsetLock 插入元素到zset，并删除已过期的元素
func insertElementToZsetLock(ctx context.Context, c redis.Client, key string, element string, zsetMaxSize int, expiresIn time.Duration) (int, error) {luaScript := `local zsetName = KEYS[1]local memberName = ARGV[1]local currentTime = tonumber(ARGV[2])local deadTime = tonumber(ARGV[3])local sizeLimit = tonumber(ARGV[4])-- 删除已过期的元素redis.call("ZREMRANGEBYSCORE", zsetName, "-inf", currentTime)-- 获取集合的大小local setSize = redis.call('ZCard', zsetName)-- 如果集合大小小于限制值，则添加元素，并返回集合大小if setSize < sizeLimit thenredis.call('ZAdd', zsetName, deadTime, memberName)local expireTime = deadTime - currentTimeif expireTime > 0 thenredis.call('EXPIRE', zsetName, expireTime)endreturn setSize+1endreturn -1`currentTime := time.Now().Unix()deadTime := time.Now().Add(expiresIn).Unix() // 过期时间 Unix秒ret, err := c.Do(ctx, "EVAL", luaScript, 1, key, element, currentTime, deadTime, zsetMaxSize).Result()if err != nil {return -1, err}if ret.(int64) < 0 {return zsetMaxSize, fmt.Errorf("zset size reach max size: %d", zsetMaxSize)}return int(ret.(int64)), nil
}

使用示例：

size, release, err := AcquireZSetLock(ctx, client, key, element, 10, 10*time.Second, 3*time.Second)
defer release()
if err != nil {fmt.Println(err)
}