【Springboot starter 组件开发】限流组件 RateLimiter
- 一、摘要
- 二、基于guava实现
-
- [2.1 核心依赖](#2.1 核心依赖)
- [2.2 核心逻辑](#2.2 核心逻辑)
- [三、基于Redis + lua脚本实现](#三、基于Redis + lua脚本实现)
-
- [3.1 核心依赖](#3.1 核心依赖)
- [3.2 核心逻辑](#3.2 核心逻辑)
一、摘要
- 基于guava的RateLimiter,实现限流
- 基于redis + lua脚本(推荐,准确性高),实现限流
- 掌握springboot starter的开发流程
- 源码地址:ratelimiter-spring-boot-starter
二、基于guava实现
2.1 核心依赖
xml
<dependency>
<groupId>com.google.guava</groupId>
<artifactId>guava</artifactId>
<version>23.5-jre</version>
</dependency>
2.2 核心逻辑
java
@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.getPermitsPerSecond(), entity.getTimeout(), TimeUnit.SECONDS);
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.getPermitsPerSecond(), 1, TimeUnit.SECONDS);
limiterMap.put(key, rateLimiter);
log.info("Guava limiter new bucket, key={}, permits={}", key, entity.getPermitsPerSecond());
return rateLimiter;
}
return limiterMap.get(key);
}
}
三、基于Redis + lua脚本实现
3.1 核心依赖
xml
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-data-redis</artifactId>
<scope>provided</scope>
</dependency>
3.2 核心逻辑
java
@Slf4j
public class RedisLimiter implements LimiterManager {
private final StringRedisTemplate stringRedisTemplate;
public RedisLimiter(StringRedisTemplate stringRedisTemplate) {
this.stringRedisTemplate = stringRedisTemplate;
}
@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());
double permitsPerSecond = entity.getPermitsPerSecond();
long timeout = entity.getTimeout();
RedisScript<Long> redisScript = new DefaultRedisScript<>(buildLuaScript(), Long.class);
Long count = stringRedisTemplate.execute(redisScript, keys, "" + permitsPerSecond, "" + timeout);
log.info("Redis limiter tryAccess, key={}, count={} ", entity.getKey(), count);
return count != null && count != 0;
}
private String buildLuaScript() {
return "--获取KEY\n" +
"local key = KEYS[1]\n" +
"\n" +
"local limit = tonumber(ARGV[1])\n" +
"\n" +
"local curentLimit = tonumber(redis.call('get', key) or \"0\")\n" +
"\n" +
"if curentLimit + 1 > limit\n" +
" then return 0\n" +
"else\n" +
" -- 自增长 1\n" +
" redis.call('INCRBY', key, 1)\n" +
" -- 设置过期时间\n" +
" redis.call('EXPIRE', key, ARGV[2])\n" +
" return curentLimit + 1\n" +
"end";
}
}