SpringGateway-Lua script & Java implementation of Redis current limiting component

In Spring Cloud Gateway, current limiting is the most basic function of the gateway. Spring Cloud Gateway officially provides the RequestRateLimiterGatewayFilterFactory class, which uses Redis and lua scripts to implement token buckets.

The specific implementation logic is in the RequestRateLimiterGatewayFilterFactory class, and the default call parameters are as follows:

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);
}

The lua script is in the Scripts folder: 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 version implementation

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);
		}
	}

}