【Springboot starter 组件开发】限流组件 RateLimiter

一、摘要

1. 基于guava的RateLimiter，实现限流
2. 基于redis + lua脚本(推荐，准确性高)，实现限流
3. 掌握springboot starter的开发流程
4. 源码地址：ratelimiter-spring-boot-starter

二、基于guava实现

2.1 核心依赖

<dependency>
    <groupId>com.google.guava</groupId>
    <artifactId>guava</artifactId>
    <version>23.5-jre</version>
</dependency>

2.2 核心逻辑

@Slf4j
public class GuavaLimiter implements LimiterManager {

    private final Map<String, RateLimiter> limiterMap = Maps.newConcurrentMap();

    @Override
    public boolean tryAccess(LimiterEntity entity) {
        if (StringUtils.isBlank(entity.getKey())) {
            throw new LimiterException("Guava limiter key cannot be empty");
        }
        RateLimiter rateLimiter = getRateLimiter(entity);
        if (rateLimiter == null) {
            return false;
        }
        boolean result = rateLimiter.tryAcquire(entity.getLimit(), 200, TimeUnit.MILLISECONDS);
        log.info("Guava limiter tryAccess, key={}, result={}", entity.getKey(), result);
        return result;
    }

    private RateLimiter getRateLimiter(LimiterEntity entity) {
        String key = entity.getKey();
        // 先看缓存中是否存在
        if (!limiterMap.containsKey(key)) {
            // 缓存中不存在，则创建令牌桶，预热时间设置为1s
            RateLimiter rateLimiter = RateLimiter.create(entity.getLimit(), 1, TimeUnit.SECONDS);
            limiterMap.put(key, rateLimiter);
            log.info("Guava limiter new bucket, key={}, permits={}", key, entity.getLimit());
            return rateLimiter;
        }
        return limiterMap.get(key);
    }
}

三、基于Redis + lua脚本实现

3.1 核心依赖

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-data-redis</artifactId>
    <scope>provided</scope>
</dependency>

3.2 核心逻辑

@Slf4j
public class RedisLimiter implements LimiterManager {

    private final StringRedisTemplate stringRedisTemplate;

    public RedisLimiter(StringRedisTemplate stringRedisTemplate) {
        this.stringRedisTemplate = stringRedisTemplate;
    }

    private final static DefaultRedisScript<Long> redisScript;

    static {
        log.info("Redis Limiter lua is already loaded");
        redisScript = new DefaultRedisScript<>();
        redisScript.setResultType(Long.class);
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("lua\\limit.lua")));
    }

    @Override
    public boolean tryAccess(LimiterEntity entity) {
        if (StringUtils.isBlank(entity.getKey())) {
            throw new LimiterException("Redis limiter key cannot be empty");
        }
        List<String> keys = Collections.singletonList(entity.getKey());

        Long count = stringRedisTemplate.execute(redisScript, keys,
                "" + entity.getLimit(), "" + entity.getExpire());

        log.info("Redis limiter tryAccess, key={}, count={} ", entity.getKey(), count);

        return count != null && count != 0;
    }
}

lua脚本:

-- 简单的令牌桶算法的变体
local key = KEYS[1]
local limit = tonumber(ARGV[1])
local expire = tonumber(ARGV[2])

local current = tonumber(redis.call('get', key) or "0")

if current == 0 then
    -- 当前计数为0，说明这是第一次请求或计数已重置
    redis.call('SET', key, '1') -- 设置计数器为1
    redis.call('EXPIRE', key, expire) -- 设置过期时间
    return 1
elseif current + 1 > limit then
    -- 达到限流阈值
    return 0
else
    -- 计数器递增
    redis.call('INCRBY', key, 1)
    return current + 1
end