最近在写一个接口限流的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}";