大家好，我是飄渺。
在SpringBoot 如何進行限流，老鳥們都這麼玩的!一文中我們詳細介紹了爲什麼需要對接口進行限流，也介紹了常見的限流算法，最後還基於Guava工具類實現了接口限流。但是這種方式有個問題，無法實現分佈式限流。那今天我們來利用Redis + Lua 來實現分佈式限流。
Lua 腳本和 MySQL 數據庫的存儲過程比較相似，他們執行一組命令，所有命令的執行要麼全部成功或者失敗，以此達到原子性。也可以把 Lua 腳本理解爲，一段具有業務邏輯的代碼塊。
實現過程
第一步：引入Redis依賴包
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-data-redis</artifactId>
</dependency>
第二步：配置Redis
/**
* @author JAVA日知錄
* @date 2022/5/2 22:35
*/
@Configuration
public class RedisConfig {
@Bean
public RedisTemplate<String, Object> redisTemplate(RedisConnectionFactory factory) {
RedisTemplate<String, Object> template = new RedisTemplate<>();
template.setConnectionFactory(factory);
// 使用Jackson2JsonRedisSerializer來序列化和反序列化redis的value值（默認使用JDK的序列化方式）
Jackson2JsonRedisSerializer<Object> serializer = new Jackson2JsonRedisSerializer<>(Object.class);
ObjectMapper mapper = new ObjectMapper();
mapper.setVisibility(PropertyAccessor.ALL, JsonAutoDetect.Visibility.ANY);
mapper.activateDefaultTyping(LaissezFaireSubTypeValidator.instance,
ObjectMapper.DefaultTyping.NON_FINAL, JsonTypeInfo.As.PROPERTY);
serializer.setObjectMapper(mapper);
template.setValueSerializer(serializer);
// 使用StringRedisSerializer來序列化和反序列化redis的key值
template.setKeySerializer(new StringRedisSerializer());
template.afterPropertiesSet();
return template;
}
}
第二步：自定義限流注解
@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.METHOD})
@Documented
public @interface RedisLimit {
/**
* 資源的key,唯一
* 作用：不同的接口，不同的流量控制
*/
String key() default "";
/**
* 最多的訪問限制次數
*/
long permitsPerSecond() default 2;
/**
* 過期時間也可以理解爲單位時間，單位秒，默認60
*/
long expire() default 60;
/**
* 得不到令牌的提示語
*/
String msg() default "系統繁忙,請稍後再試.";
}
第三步：創建限流異常
/**
* @author JAVA日知錄
* Redis限流自定義異常
* @date 2022/5/2 21:43
*/
public class RedisLimitException extends RuntimeException{
public RedisLimitException(String msg) {
super( msg );
}
}
第四步：使用AOP切面攔截限流注解
/**
* Limit AOP
* @author JAVA日知錄
* @date 2021/9/24 3:07 下午
*/
@Slf4j
@Aspect
@Component
public class RedisLimitAop {
@Autowired
private StringRedisTemplate stringRedisTemplate;
@Pointcut("@annotation(com.jianzh5.blog.limit.redis.RedisLimit)")
private void check() {
}
@Before("check()")
public void before(JoinPoint joinPoint) {
MethodSignature signature = (MethodSignature) joinPoint.getSignature();
Method method = signature.getMethod();
//拿到RedisLimit註解，如果存在則說明需要限流
RedisLimit redisLimit = method.getAnnotation(RedisLimit.class);
if(redisLimit != null){
//獲取redis的key
String key = redisLimit.key();
String className = method.getDeclaringClass().getName();
String name = method.getName();
String limitKey = key + className + method.getName();
log.info(limitKey);
if(StringUtils.isEmpty(key)){
throw new RedisLimitException( "key cannot be null" );
}
long limit = redisLimit.permitsPerSecond();
long expire = redisLimit.expire();
List<String> keys = new ArrayList<>();
keys.add( key );
String luaScript = buildLuaScript();
RedisScript<Long> redisScript = new DefaultRedisScript<>( luaScript, Long.class );
Long count = stringRedisTemplate.execute( redisScript, keys, String.valueOf(limit), String.valueOf(expire) );
log.info( "Access try count is {} for key={}", count, key );
if (count != null && count == 0) {
log.debug("令牌桶={}，獲取令牌失敗",key);
throw new RedisLimitException(redisLimit.msg());
}
}
}
/**
* 構建redis lua腳本
* @return
*/
private String buildLuaScript() {
StringBuilder luaString = new StringBuilder();
luaString.append( "local key = KEYS[1]" );
//獲取ARGV內參數Limit
luaString.append( "\nlocal limit = tonumber(ARGV[1])" );
//獲取key的次數
luaString.append( "\nlocal curentLimit = tonumber(redis.call('get', key) or \"0\")" );
luaString.append( "\nif curentLimit + 1 > limit then" );
luaString.append( "\nreturn 0" );
luaString.append( "\nelse" );
//自增長 1
luaString.append( "\n redis.call(\"INCRBY\", key, 1)" );
//設置過期時間
luaString.append( "\nredis.call(\"EXPIRE\", key, ARGV[2])" );
luaString.append( "\nreturn curentLimit + 1" );
luaString.append( "\nend" );
return luaString.toString();
}
}
第五步：給需要限流的接口加上註解
/**
* 公衆號：JAVA日知錄
* 限流測試類基於Redis限流
*/
@Slf4j
@RestController
@RequestMapping("/limit/redis")
public class LimitRedisController {
/**
* 基於Redis AOP限流
*/
@GetMapping("/test")
@RedisLimit(key = "redis-limit:test", permitsPerSecond = 2, expire = 1, msg = "當前排隊人數較多，請稍後再試！")
public String test() {
log.info("限流成功。。。");
return "ok";
}
}
第六步：體驗效果
通過訪問測試地址： http://127.0.0.1:8080/limit/redis/test，反覆刷新並觀察輸出結果：
正常響應時：
{"status":100,"message":"操作成功","data":"ok","timestamp":1652343229643}
觸發限流時：
{"status":500,"message":"當前排隊人數較多，請稍後再試！","data":null,"timestamp":1652343239035}
通過觀察得之，基於自定義註解同樣實現了接口限流的效果。
優化
程序每次執行每次都需要通過buildLuaScript()方法構建lua執行腳本，比較 low，我們可以生成一個lua文件放在resources目錄下，利用@PostConstruct註解提前加載。
- 在resouces文件夾下創建lua文件 rateLimiter.lua
--獲取KEY
local key = KEYS[1]
local limit = tonumber(ARGV[1])
local curentLimit = tonumber(redis.call('get', key) or "0")
if curentLimit + 1 > limit
then return 0
else
-- 自增長 1
redis.call('INCRBY', key, 1)
-- 設置過期時間
redis.call('EXPIRE', key, ARGV[2])
return curentLimit + 1
end
- 修改RedisLimitAop，通過@PostConstruct注入DefaultRedisScript
@Slf4j
@Aspect
@Component
public class RedisLimitAop {
@Autowired
private StringRedisTemplate stringRedisTemplate;
private DefaultRedisScript<Long> redisScript;
@PostConstruct
public void init(){
redisScript = new DefaultRedisScript<>();
redisScript.setResultType(Long.class);
redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("rateLimiter.lua")));
}
...
}
小結
基於Redis + Lua 可以很方便地實現分佈式限流，算是SpringBoot老鳥系列限流文章的補充擴展。
那麼現在問題來了，我們現在有基於Guava實現的單機限流，又有基於Redis+Lua實現的分佈式限流，那能不能將兩種限流功能做成一個獨立的公共組件，讓使用方根據實際情況選擇對應的限流功能呢？
老鳥系列源碼已經上傳至GitHub，需要的在公號【JAVA日知錄】回覆關鍵字 0923 獲取源碼地址。