SpringGateway - Redis限流组件之Lua脚本&Java实现

在Spring Cloud Gateway中，限流作为网关最基本的功能，Spring Cloud Gateway官方就提供了RequestRateLimiterGatewayFilterFactory这个类，适用Redis和lua脚本实现了令牌桶的方式。

具体实现逻辑在RequestRateLimiterGatewayFilterFactory类中,默认调用参数如下：

List<String> keys = getKeys(id);

// The arguments to the LUA script. time() returns unixtime in seconds.
List<String> scriptArgs = Arrays.asList(replenishRate + "", burstCapacity + "",
					Instant.now().getEpochSecond() + "", "1");

// allowed, tokens_left = redis.eval(SCRIPT, keys, args)
Flux<List<Long>> flux = this.redisTemplate.execute(this.script, keys, scriptArgs);

static List<String> getKeys(String id) {
		// use `{}` around keys to use Redis Key hash tags
		// this allows for using redis cluster

		// Make a unique key per user.
		String prefix = "request_rate_limiter.{" + id;

		// You need two Redis keys for Token Bucket.
		String tokenKey = prefix + "}.tokens";
		String timestampKey = prefix + "}.timestamp";
		return Arrays.asList(tokenKey, timestampKey);
}

lua脚本在Scripts文件夹下：request_rate_limitter.lua

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 }

Java 版本实现

package com.liuwei.springboot.algorithm.ratelimit;

import java.util.concurrent.ConcurrentHashMap;

public class RedisLimitter {
	
	public static void main(String[] args) {
		RedisLimitter limitter = new RedisLimitter();
		RateDataStore dataStore  = new RateLimitterLocalStore();
		boolean checkResult = limitter.rateLimit("uuid",1d,10d,50d,dataStore);
		System.out.println(checkResult);
	}
	
	/**
	 * 
	 * @param requestkey 请求唯一标识
	 * @param requested	 请求量
	 * @param rate	令牌桶填充平均速率，单位：秒
	 * @param capacity	 令牌桶上限
	 * @return
	 */
	public boolean rateLimit(String requestkey, double requested,double rate,double capacity,RateDataStore dataStore) {
		// https://blog.csdn.net/weixin_42073629/article/details/106934827
		String tokens_key = String.format("request_rate_limiter.%s.tokens", requestkey);  		// 令牌桶剩余令牌数
		String timestamp_key = String.format("request_rate_limiter.%s.timestamp", requestkey);	// 令牌桶最后填充令牌时间，单位：秒
		
		// double rate = 10;  		// 令牌桶填充平均速率，单位：秒
		// double capacity = 50;  	// 令牌桶上限
		double now = System.currentTimeMillis();	// 当前时间戳
		// double requested = 1;	// 请求量
		
		// 计算令牌桶填充满令牌需要多久时间，单位：秒
		// 如果是Redis * 2 保证时间充足,  如果设置永不过期也不影响功能
		double fill_time = capacity/rate;			
		double ttl = Math.floor(fill_time*2);  		
		
		// 获得令牌桶剩余令牌数( last_tokens ) 
		Double last_tokens = dataStore.getRateData(tokens_key);	
		if(last_tokens == null) last_tokens = capacity;
		
		// 令牌桶最后填充令牌时间(last_refreshed) 
		Double last_refreshed = dataStore.getRateData(timestamp_key);
		if(last_refreshed == null) last_refreshed = 0d;
		
		// 填充令牌，计算新的令牌桶剩余令牌数( filled_tokens )。填充不超过令牌桶令牌上限
		double delta = Math.max(0, now-last_refreshed);
		double filled_tokens = Math.min(capacity, last_tokens+(delta*rate));
		
		boolean allowed = filled_tokens >= requested;
		double new_tokens = filled_tokens;
		double allowed_num = 0;
		if(allowed) {
			new_tokens = filled_tokens - requested;
			allowed_num = requested;
		}
		// redis.call("setex", tokens_key, ttl, new_tokens)
		// redis.call("setex", timestamp_key, ttl, now)
		// return { allowed_num, new_tokens }
		
		dataStore.setRateData(tokens_key, new_tokens);
		dataStore.setRateData(timestamp_key, now);
		System.out.println(String.format("allowed_num:%s, new_tokens:%s ",allowed_num, new_tokens));
		return allowed;
	}
	
	/**
	 *  限流工具全局存储, 可基于数据库或Redis实现
	 * @author LIUWEI122
	 *
	 */
	public static interface RateDataStore{
		public Double getRateData(String key);
		public void setRateData(String key,Double value);
		public void setRateData(String key,double ttl,Double value);
	}
	
	public static class RateLimitterLocalStore implements RateDataStore{
		private static ConcurrentHashMap<String,Double> store = new ConcurrentHashMap<>();
		@Override
		public Double getRateData(String key) {
			return store.get(key);
		}

		@Override
		public void setRateData(String key, Double value) {
			store.put(key,value);
		}

		@Override
		public void setRateData(String key, double ttl, Double value) {
			store.put(key,value);
		}
	}

}