ApiBoot RateLimiter redis方式Lua脚本

Author: hengboy

api-boot-project/api-boot-plugins/api-boot-plugin-rate-limiter/src/main/resources/META-INF/scripts/qps-rate-limiter.lua

@@ -0,0 +1,50 @@
+local tokens_key = KEYS[1]
+local timestamp_key = KEYS[2]
+--redis.log(redis.LOG_WARNING, "tokens_key " .. tokens_key)
+
+local rate = tonumber(ARGV[1])
+local capacity = tonumber(ARGV[2])
+local now = tonumber(ARGV[3])
+local requested = tonumber(ARGV[4])
+
+local fill_time = capacity/rate
+local ttl = math.floor(fill_time*2)
+
+--redis.log(redis.LOG_WARNING, "rate " .. ARGV[1])
+--redis.log(redis.LOG_WARNING, "capacity " .. ARGV[2])
+--redis.log(redis.LOG_WARNING, "now " .. ARGV[3])
+--redis.log(redis.LOG_WARNING, "requested " .. ARGV[4])
+--redis.log(redis.LOG_WARNING, "filltime " .. fill_time)
+--redis.log(redis.LOG_WARNING, "ttl " .. ttl)
+
+local last_tokens = tonumber(redis.call("get", tokens_key))
+if last_tokens == nil then
+  last_tokens = capacity
+end
+--redis.log(redis.LOG_WARNING, "last_tokens " .. last_tokens)
+
+local last_refreshed = tonumber(redis.call("get", timestamp_key))
+if last_refreshed == nil then
+  last_refreshed = 0
+end
+--redis.log(redis.LOG_WARNING, "last_refreshed " .. last_refreshed)
+
+local delta = math.max(0, now-last_refreshed)
+local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
+local allowed = filled_tokens >= requested
+local new_tokens = filled_tokens
+local allowed_num = 0
+if allowed then
+  new_tokens = filled_tokens - requested
+  allowed_num = 1
+end
+
+--redis.log(redis.LOG_WARNING, "delta " .. delta)
+--redis.log(redis.LOG_WARNING, "filled_tokens " .. filled_tokens)
+--redis.log(redis.LOG_WARNING, "allowed_num " .. allowed_num)
+--redis.log(redis.LOG_WARNING, "new_tokens " .. new_tokens)
+
+redis.call("setex", tokens_key, ttl, new_tokens)
+redis.call("setex", timestamp_key, ttl, now)
+
+return { allowed_num, new_tokens }