前言 很多時候,由於種種不可描述的原因,我們需要針對單個介面實現介面限流,防止訪問次數過於頻繁。這裡就用 redis+aop 實現一個限流介面註解 @RedisLimit 代碼 點擊查看RedisLimit註解代碼 import java.lang.annotation.*; /** * 功能:分佈 ...
前言
- 很多時候,由於種種不可描述的原因,我們需要針對單個介面實現介面限流,防止訪問次數過於頻繁。這裡就用 redis+aop 實現一個限流介面註解
@RedisLimit 代碼
點擊查看RedisLimit註解代碼
import java.lang.annotation.*;
/**
* 功能:分散式介面限流註解
* @author love ice
* @create 2023-09-18 15:43
*/
@Target({ElementType.TYPE,ElementType.METHOD})
@Documented
@Retention(RetentionPolicy.RUNTIME)
public @interface RedisLimit {
/**
* redis中唯一key,一般用方法名字做區分
* 作用: 針對不同介面,做不同的限流控制
*/
String key() default "";
/**
* 限流時間內允許訪問次數 預設1
*/
long permitsPerSecond() default 1;
/**
* 限流時間,單位秒 預設60秒
*/
long expire() default 60;
/**
* 限流提示信息
*/
String msg() default "介面限流,請稍後重試";
}
AOP代碼
點擊查看aop代碼
import com.aliyuncs.utils.StringUtils;
import com.test.redis.Infrastructure.annotation.RedisLimit;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Component;
import javax.annotation.PostConstruct;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
/**
* redis限流切麵
*
* @author love ice
* @create 2023-09-18 15:44
*/
@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);
// 執行 lua 腳本
ResourceScriptSource resourceScriptSource = new ResourceScriptSource(new ClassPathResource("rateLimiter.lua"));
redisScript.setScriptSource(resourceScriptSource);
}
@Pointcut("@annotation(com.test.redis.Infrastructure.annotation.RedisLimit)")
private void check() {
}
@Before("check()")
private 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 + name;
log.info("限流的key:{}", limitKey);
if (StringUtils.isEmpty(key)) {
// 這裡是自定義異常,為了方便寫成了 RuntimeException
throw new RuntimeException("code:101,msg:介面中 key 參數不能為空");
}
long limit = redisLimit.permitsPerSecond();
long expire = redisLimit.expire();
// 把 key 放入 List 中
List<String> keys = new ArrayList<>(Collections.singletonList(key));
Long count = stringRedisTemplate.execute(redisScript, keys, String.valueOf(limit), String.valueOf(expire));
log.info("Access try count is {} for key ={}", count, keys);
if (count != null && count == 0) {
log.debug("令牌桶={}, 獲取令牌失效,介面觸發限流", key);
throw new RuntimeException("code:10X, redisLimit.msg()");
}
}
}
}
lua腳本代碼
註意:腳本代碼是放在 resources 文件下的,它的類型是 txt,名稱尾碼是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