使用 Redis 实现四种限流算法

固定窗口

local key = KEYS[1]
local maxRequests = tonumber(ARGV[1])
local windowSize = tonumber(ARGV[2])

local count = tonumber(redis.call('get', key))

if not count then
    redis.call('SET', key, 1, 'EX', windowSize)
    return true
end

if count < maxRequests then
    redis.call('INCR', key)
    return true
end

return false

@Component
public class FixedWindowRateLimiter {

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    private static final String RATE_LIMIT_KEY_PREFIX = "rate_limit:fixed:";

    /**
     * 固定窗口限流
     *
     * @param key         限流 key
     * @param maxRequests 最大请求数
     * @param windowSize  窗口大小(秒)
     * @return 是否允许通过
     */
    public boolean tryAcquire(String key, int maxRequests, int windowSize) {
        key = RATE_LIMIT_KEY_PREFIX + key;
        DefaultRedisScript<Boolean> script = new DefaultRedisScript<>();
        script.setScriptSource(new ResourceScriptSource(new ClassPathResource("scripts/rateLimit.lua")));
        script.setResultType(Boolean.class);
        Boolean ok = redisTemplate.execute(script, List.of(key), maxRequests, windowSize);
        return Objects.equals(ok, Boolean.TRUE);
    }
}

滑动窗口

local key = KEYS[1]
local maxRequests = tonumber(ARGV[1])
local windowSize = tonumber(ARGV[2])
local currentTime = tonumber(ARGV[3])
local windowStart = tonumber(ARGV[4])

redis.call('ZREMRANGEBYSCORE', key, 0, windowStart)

local count = redis.call('ZCARD', key)

if count < maxRequests then
    redis.call('ZADD', key, currentTime, currentTime)
    redis.call('EXPIRE', key, windowSize)
    return true
end

return false

@Component
public class SlidingWindowRateLimiter {

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    private static final String RATE_LIMIT_KEY_PREFIX = "rate_limit:sliding:";

    /**
     * 滑动窗口限流
     *
     * @param key         限流 key
     * @param maxRequests 最大请求数
     * @param windowSize  窗口大小(秒)
     * @return 是否允许通过
     */
    public boolean tryAcquire(String key, int maxRequests, int windowSize) {
        key = RATE_LIMIT_KEY_PREFIX + key;
        long currentTime = System.currentTimeMillis();
        long windowStart = currentTime - windowSize * 1000L;
        DefaultRedisScript<Boolean> script = new DefaultRedisScript<>();
        script.setScriptSource(new ResourceScriptSource(new ClassPathResource("scripts/rateLimit.lua")));
        script.setResultType(Boolean.class);
        Boolean ok = redisTemplate.execute(script, List.of(key), maxRequests, windowSize, currentTime, windowStart);
        return Objects.equals(ok, Boolean.TRUE);
    }
}

漏桶

local key = KEYS[1]
local rate = tonumber(ARGV[1])
local currentTime = tonumber(ARGV[2])

local nextAvailableTime = tonumber(redis.call('GET', key))

if not nextAvailableTime or currentTime > nextAvailableTime then
    nextAvailableTime = currentTime
end

local intervalMs = math.ceil(1000 / rate)
local waitTime = nextAvailableTime - currentTime
nextAvailableTime = nextAvailableTime + intervalMs
redis.call('SET', key, nextAvailableTime, 'PX', (nextAvailableTime - currentTime))

return waitTime

@Component
public class LeakyBucketRateLimiter {

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    private static final String RATE_LIMIT_KEY_PREFIX = "rate_limit:leaky:";

    /**
     * 漏桶算法限流
     *
     * @param key  限流 key
     * @param rate 流出速率(请求/秒)
     * @return 是否允许通过
     */
    public boolean acquire(String key, int rate) throws InterruptedException {
        key = RATE_LIMIT_KEY_PREFIX + key;
        long currentTime = System.currentTimeMillis();
        DefaultRedisScript<Long> script = new DefaultRedisScript<>();
        script.setScriptSource(new ResourceScriptSource(new ClassPathResource("scripts/rateLimit.lua")));
        script.setResultType(Long.class);
        Long waitTime = redisTemplate.execute(script, List.of(key), rate, currentTime);
        assert waitTime != null;
        Thread.sleep(waitTime);
        return true;
    }
}

令牌桶

local key = KEYS[1]
local capacity = tonumber(ARGV[1])
local rate = tonumber(ARGV[2])
local currentTime = tonumber(ARGV[3])

local bucket = redis.call('HMGET', key, 'tokens', 'lastTime')
local tokens = capacity

if bucket[1] and bucket[2] then
    tokens = tonumber(bucket[1])
    local lastTime = tonumber(bucket[2])

    local newTokens = math.floor((currentTime - lastTime) * rate / 1000)
    tokens = math.min(capacity, tokens + newTokens)
end

if tokens >= 1 then
    tokens = tokens - 1
    redis.call('HSET', key, 'tokens', tokens, 'lastTime', currentTime)
    redis.call('EXPIRE', key, math.ceil(capacity / rate))
    return true
end

return false

@Component
public class TokenBucketRateLimiter {

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    private static final String RATE_LIMIT_KEY_PREFIX = "rate_limit:token:";

    /**
     * 令牌桶限流
     *
     * @param key      限流key
     * @param capacity 桶容量
     * @param rate     令牌生成速率(令牌/秒)
     * @return 是否允许通过
     */
    public boolean tryAcquire(String key, int capacity, int rate) {
        key = RATE_LIMIT_KEY_PREFIX + key;
        long currentTime = System.currentTimeMillis();
        DefaultRedisScript<Boolean> script = new DefaultRedisScript<>();
        script.setScriptSource(new ResourceScriptSource(new ClassPathResource("scripts/rateLimit.lua")));
        script.setResultType(Boolean.class);
        Boolean ok = redisTemplate.execute(script, List.of(key), capacity, rate, currentTime);
        return Objects.equals(ok, Boolean.TRUE);
    }
}