Redis+Lua实现分布式限流

限流接口

public interface EasyRateLimiter {
 
    void acquire(int permits);
}

调用示例

@Resource(name = "redisEasyRateLimiterImpl")
private EasyRateLimiter easyRateLimiter;
 
@Test
public void test() throws InterruptedException {
    LOGGER.info("限流前");
    easyRateLimiter.acquire(1);
    LOGGER.info("限流后");
}

限流实现

import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
 
import java.util.Collections;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import java.util.logging.Logger;
 

public class RedisEasyRateLimiterImpl implements EasyRateLimiter {
 
    private static final Logger LOGGER = Logger.getLogger(RedisEasyRateLimiterImpl.class.getName());
 
    private String redisKey;
 
    private Integer qps;// 限流器速率，自行注入
 
    private RedisTemplate<String, Long> redisTemplate;// spring redis template，需要注意，由于需要与lua脚本交互，Key、Value需要使用String序列化
 
    private RedisScript<Long> redisScript;// spring redis script
 
    public void acquire(int permits) {
        Long currentTime = System.currentTimeMillis();
        //@TODO 调用TIME命令获取统一时间
        //redisTemplate.getConnectionFactory().getConnection().time();//如此获取会造成连接泄漏
        LOGGER.info("current time: " + currentTime);
        Long holdTime = redisTemplate.execute(redisScript, Collections.singletonList(redisKey),
                String.valueOf(currentTime), String.valueOf(qps), String.valueOf(permits), UUID.randomUUID().toString().replace("-", ""));
        LOGGER.info("holdTime ms: " + holdTime);
        if (holdTime > 0) {
            try {
                TimeUnit.MILLISECONDS.sleep(holdTime);// 当前线程睡眠等待
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
    }
 
    // setter and getter...
}

lua脚本

local nextPermissionTime = redis.call('GET', KEYS[1])
local currentTime = tonumber(ARGV[1])
local qps = tonumber(ARGV[2])
local permits = tonumber(ARGV[3])
local blockingTime = 0
 
if nextPermissionTime then
    nextPermissionTime = tonumber(nextPermissionTime)
    if nextPermissionTime > currentTime then
        blockingTime = nextPermissionTime - currentTime
    else
        nextPermissionTime = currentTime
    end
else
    nextPermissionTime = currentTime
end
 
nextPermissionTime = nextPermissionTime + 1 / qps * permits * 1000
redis.call('SET', KEYS[1], nextPermissionTime)
return blockingTime