益竹月 发表于 2025-11-19 23:05:00

新来的外包,在大群分享了它的限流算法的实现


1. 令牌桶按用户维度限流

前文golang/x/time/rate演示了基于整体请求速率的令牌桶限流;
那基于用户id、ip、apikey请求速率的限流(更贴近生产的需求), 阁下又该如何应对?
那这个问题就从全局速率变成了按照用户维度(group by userid)来做限流,那么

[*]早先的全局的rateLimiter就要变成 userid:rateLimiter的键值对, select count( * ) from table ---> select userid, count(*) fromtable group by userid
[*]使用缓存组件来存储维度键值对: 缓存的剔除机制来清理不再访问的键值对 (30min过期,10min周期清理内存)。
var userLimiters = cache.New(time.Minute*30, 10) // 10 items per minute
func limiterForUser(userID string) *rate.Limiter {
        if v, found := userLimiters.Get(userID); found {
                return v.(*rate.Limiter)
        }

        l := rate.NewLimiter(rate.Every(time.Minute/60), 10)
        userLimiters.Set(userID, l, cache.DefaultExpiration)
        return l
}

// 更细化的限流: 针对同一用户的请求次数限速, 增加了细粒度的用户维度,需要维护 用户与对应限速器的映射关系
func userRatelimitMiddleware(c *gin.Context) {
        userID := c.GetString("userID")//从每个请求context的key中取得信息, 这个key对于req context是排他性的
        if userID == "" {
                c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized"})
                return
        }
        if userID == "" {
                userID = c.GetString("x-api-key")
        }

        if userID == "" {
                userID = c.ClientIP()
        }
        limiter := limiterForUser(userID) // 通过userid维度找到对应的限速器
        if !limiter.Allow() {
                c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{"error": "Too many requests"})
                return
        }
        c.Next()
}2. redis 作为限流器第三方存储

这个思路也是极其常见的行为: redis可以成为用户令牌桶的全局中心存储: 当多个负载层需要读写用户限流器时,与redis交互。
本次通过golang的实战,深入理解基于redis的令牌桶限流器的算法实现。
①请求到达负载层,被负载层识别为userid=junio
②负载层请求redis获取该用户的token bucket的当前状态:
hget userbucket:junio tokens last_time
③基于当前时间now和last_time,计算流逝的时间,再根据rate计算这一阶段下发了多少tokens:delta=(now-last_time) * r/1000,加上redis原始记录的token,就是本次请求时bucket中能用的tokens, 注意:令牌数量最多不能超过cap
④ 如果tokens>=1, 表示桶中有令牌,可放行请求,tokens数量减1
⑤ 最后将本次处理完后的 tokens和last_time=now写入原用户令牌桶
hsetuserbucket:juniotokens20 last_time 990

使用redis 中的hashmap存储用户的tokenbucket状态,应用存在读取redis- 计算- 回写redis过程,使用redis lua的脚本执行三个动作,以保证线程安全。
为什么lua脚本能保证线程安全呢?
主要得益于 Redis 的单线程架构和原子性执行机制: 加载并执行lua脚本时所有的redis操作作为一个整体完成; 整个脚本执行期间没有其他命令可以插入。
// 读取- 计算 - 重新赋值都在一个 lua 脚本里面
var redisScript = `
        local key = KEYS
        local capacity = tonumber(ARGV)
        local rate = tonumber(ARGV)
        local now = tonumber(ARGV)
        local tokens =tonumber(redis.call('hget', key, 'tokens') or '-1')
        local last_time = tonumber(redis.call('hget', key, 'last_time') or'-1')

        if tokens== -1 or last_time == -1 then
                tokens = capacity
                last_time = now
        else
                local elapsed = now - last_time
      if elapsed < 0
                        then elapsed = 0
                end
                local delta= elapsed * rate / 1000
                tokens = tokens + delta
                if tokens > capacity then
                        tokens = capacity
                end
      last_time = now
        end
        local allow = 0
        if tokens >= 1 then
                allow = 1
                tokens= tokens - 1
        else       
                allow = 0
        end

        redis.call('hset', key, 'tokens', tokens)
        redis.call('hset', key, 'last_time', last_time)
    redis.call('PEXPIRE', key,math.max(1000, 2 * math.ceil((capacity / rate) * 5000)))
        return allow
`注意

[*]上面还使用的redis expire机制: redis expire不是滑动过期,但是每次被请求触发执行的时候就重新设置TTL, 表现为“滑动过期”。
[*]除了hset/hget ,还有hmget可用,另外这些操作还有配套的TTL指令,eg:hset key EXAT1740470400 FIELDS 2 field1 "Hello" field2 "World"。
golang应用层的写法如下:
func (r *RedisLimiter) Allow(c *gin.Context, userid string) bool {
        key := r.keyprefix + userid // 定位这个用户的token bucket
        now := time.Now().UnixMilli()
        // Check if the key exists in Redis
        rCmd := r.redis.Eval(redisScript, []string{key}, r.cap, r.rate, now)
        res, err := rCmd.Result()
        if err != nil {
                log.Printf("get from redis failure. ", err)
                return false
        }
        if allow, ok := res.(int64); ok { // 注意:lua返回的0,1 值对应golang的int64
                log.Printf("%v %v \n", allow, res)
                return allow == 1
        } else {
                log.Printf("get from redis failure. ", err)
                return false
        }
}
至此限流第二弹结束了,本文紧接掘金爆文
来源:程序园用户自行投稿发布,如果侵权,请联系站长删除
免责声明:如果侵犯了您的权益,请联系站长,我们会及时删除侵权内容,谢谢合作!

焦听云 发表于 2025-11-25 13:59:38

喜欢鼓捣这些软件,现在用得少,谢谢分享!

甦忻愉 发表于 2025-11-26 17:28:01

鼓励转贴优秀软件安全工具和文档!

邹语彤 发表于 2025-11-27 02:35:11

很好很强大我过来先占个楼 待编辑
页: [1]
查看完整版本: 新来的外包,在大群分享了它的限流算法的实现