redis令牌桶算法拦截率远低于预期

38 阅读2分钟

最近在写一个接口限流的springboot 的 stater,用Augment生成了基于redis的令牌桶限流脚本


private static final String TOKEN_BUCKET_SCRIPT =
    "local key = KEYS[1]\n" +
    "local capacity = tonumber(ARGV[1])\n" +
    "local refill_rate = tonumber(ARGV[2])\n" +
    "local time_window = tonumber(ARGV[3])\n" +
    "local now = tonumber(ARGV[4])\n" +
    "\n" +
    "local bucket = redis.call('HMGET', key, 'tokens', 'last_refill')\n" +
    "local tokens = tonumber(bucket[1]) or capacity\n" +
    "local last_refill = tonumber(bucket[2]) or now\n" +
    "\n" +
    "-- 计算需要补充的令牌数\n" +
    "local elapsed = math.max(0, now - last_refill)\n" +
    "local tokens_to_add = math.floor(elapsed / 1000 * refill_rate)\n" +
    "tokens = math.min(capacity, tokens + tokens_to_add)\n" +
    "\n" +
    "local allowed = 0\n" +
    "if tokens > 0 then\n" +
    "    tokens = tokens - 1\n" +
    "    allowed = 1\n" +
    "end\n" +
    "\n" +
    "-- 更新令牌桶状态\n" +
    "redis.call('HMSET', key, 'tokens', tokens, 'last_refill', now)\n" +
    "redis.call('EXPIRE', key, time_window * 2)\n" +
    "return {allowed, tokens}";

看着没啥问题,当我把令牌桶容量和填充速率设为10的时候,用Jmeter在10秒内跑了1000个请求,发现被拦击的数量只有不到100个,理论上通过的请求大概是10*10+10=110个左右,拦截的应该有900个左右。在加了日志之后发现了其实是由于last_refill导致的,如果后一个请求比前一个请求的时间还早,会导致last_refill被反复更新,导致释放了更多的令牌。既然是填充时间,那么就只有填充了才更新 修改如下

private static final String TOKEN_BUCKET_SCRIPT =
    "local key = KEYS[1]\n" +
    "local capacity = tonumber(ARGV[1])\n" +
    "local refill_rate = tonumber(ARGV[2])\n" +
    "local time_window = tonumber(ARGV[3])\n" +
    "local now = tonumber(ARGV[4])\n" +
    "\n" +
    "local bucket = redis.call('HMGET', key, 'tokens', 'last_refill')\n" +
    "local tokens = tonumber(bucket[1]) or capacity\n" +
    "local last_refill = tonumber(bucket[2]) or now\n" +
    "\n" +
    "-- 计算需要补充的令牌数(修复:确保时间间隔计算正确)\n" +
    "local elapsed = math.max(0, now - last_refill)\n" +
    "-- refill_rate是每秒补充的令牌数,elapsed是秒数\n" +
    "local tokens_to_add = math.floor(elapsed * refill_rate)\n" +
    "\n" +
    "-- 只有当需要添加令牌时才更新last_refill,避免重复填充\n" +
    "local new_last_refill = last_refill\n" +
    "if tokens_to_add > 0 then\n" +
    "    tokens = math.min(capacity, tokens + tokens_to_add)\n" +
    "    new_last_refill = now\n" +
    "end\n" +
    "\n" +
    "local allowed = 0\n" +
    "if tokens > 0 then\n" +
    "    tokens = tokens - 1\n" +
    "    allowed = 1\n" +
    "end\n" +
    "\n" +
    "-- 更新令牌桶状态\n" +
    "redis.call('HMSET', key, 'tokens', tokens, 'last_refill', new_last_refill)\n" +
    "redis.call('EXPIRE', key, time_window * 2)\n" +
    "return {allowed, tokens}";