JAVA自定義注解實現(xiàn)接口/ip限流的示例代碼

更新時間：2023年07月28日 15:18:53 作者：二月二_

本文主要介紹了JAVA自定義注解實現(xiàn)接口/ip限流的示例代碼，文中通過示例代碼介紹的非常詳細，對大家的學習或者工作具有一定的參考學習價值，需要的朋友們下面隨著小編來一起學習學習吧

基本步驟：
1.創(chuàng)建枚舉類，定義限流類型
2.自定義注解@RateLimiter
3.創(chuàng)建lua腳本，保證redis中操作原子性（在resources/lua目錄下創(chuàng)建rateLimit.lua文件）
4.創(chuàng)建配置類，解決redis序列化與讀取lua腳本
5.自定義限流異常與全局異常
6.創(chuàng)建切面，進行限流操作
7.創(chuàng)建接口，進行驗證

基本步驟：

1.自定義注解 @RateLimiter，添加屬性：key（redis中key前綴）、time（redis中key過期時間，單位秒）、count（time中定義時間段內(nèi)可以訪問的次數(shù)），limitType（限流類型：1.限制接口；2.限制ip）

2.創(chuàng)建AOP，定義前置通知，獲取@RateLimiter注解上的各項值，客戶端每訪問一次接口，redis中存儲的current（當前時間段內(nèi)已訪問的次數(shù)）值+1，若current值大于count規(guī)定的值，拋出異常，提示到達允許訪問的次數(shù)；

具體代碼如下：

        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-data-redis</artifactId>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-web</artifactId>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-aop</artifactId>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-test</artifactId>
            <scope>test</scope>
        </dependency>
        <!-- redis json序列化 -->
        <dependency>
            <groupId>com.fasterxml.jackson.core</groupId>
            <artifactId>jackson-databind</artifactId>
        </dependency>
        <dependency>
            <groupId>com.fasterxml.jackson.datatype</groupId>
            <artifactId>jackson-datatype-jsr310</artifactId>
        </dependency>
        <dependency>
            <groupId>cn.hutool</groupId>
            <artifactId>hutool-all</artifactId>
            <version>5.8.9</version>
        </dependency>

1.創(chuàng)建枚舉類，定義限流類型

public enum LimitType {
    /**
     * 默認限流策略，針對某一個接口進行限流
     */
    DEFAULT,
    /**
     * 根據(jù)IP地址進行限流
     */
    IP;
}

2.自定義注解@RateLimiter

@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
public @interface RateLimiter {
    /**
     * 定義redis中限流key前綴 rate_limit:com.xxx.controller.HelloController-hello //HelloController中的hello方法
     */
    String key() default "rate_limit:";
    /**
     * 限流時間，單位秒
     * @return
     */
    int time() default 60;
    /**
     * 限流時間內(nèi)限流次數(shù)
     * @return
     */
    int count() default 100;
    /**
     * 限流類型：1.限制接口訪問次數(shù) 2.限制ip訪問次數(shù)
     * @return
     */
    LimitType limitType() default LimitType.DEFAULT;
}

3.創(chuàng)建lua腳本，保證redis中操作原子性（在resources/lua目錄下創(chuàng)建rateLimit.lua文件）

---定義變量：redis中key值,redis中過期時間,規(guī)定的時間段內(nèi)訪問次數(shù),當前訪問次數(shù)
local key = KEYS[1]
local time = tonumber(ARGV[1])
local count = tonumber(ARGV[2])
local current = redis.call('get',key)
if current and tonumber(current) > count then ---如果current不為空當前訪問次數(shù)大于規(guī)定時間段內(nèi)的訪問次數(shù)
    return tonumber(current)  ---返回當前訪問的次數(shù)
end
    current = redis.call('incr',key)    --- 若沒超過限制次數(shù),當前次數(shù)+1（自增）
if tonumber(current)==1 then
    redis.call('expire',key,time)       --- 如果當前訪問次數(shù)=1,表示第一次訪問,設(shè)置當前key的過期時間
end
    return tonumber(current)        --返回tonumber(current)

4.創(chuàng)建配置類，解決redis序列化與讀取lua腳本

@Configuration
public class RedisConfig {
    //Redis序列化
    @Bean
    public RedisTemplate<String, Object> redisTemplate(LettuceConnectionFactory redisConnectionFactory) {
        RedisTemplate<String, Object> redisTemplate = new RedisTemplate<>();
        redisTemplate.setConnectionFactory(redisConnectionFactory);
        //首先解決key的序列化方式
        StringRedisSerializer stringRedisSerializer = new StringRedisSerializer();
        redisTemplate.setKeySerializer(stringRedisSerializer);
        //解決value的序列化方式
        Jackson2JsonRedisSerializer<Object> jackson2JsonRedisSerializer = new Jackson2JsonRedisSerializer<>(Object.class);
        //序列化時將類的數(shù)據(jù)類型存入json，以便反序列化的時候轉(zhuǎn)換成正確的類型
        ObjectMapper objectMapper = new ObjectMapper();
        //objectMapper.enableDefaultTyping(ObjectMapper.DefaultTyping.NON_FINAL);
        objectMapper.activateDefaultTyping(LaissezFaireSubTypeValidator.instance, ObjectMapper.DefaultTyping.NON_FINAL);
        // 解決jackson2無法反序列化LocalDateTime的問題
        objectMapper.disable(SerializationFeature.WRITE_DATES_AS_TIMESTAMPS);
        objectMapper.registerModule(new JavaTimeModule());
        jackson2JsonRedisSerializer.setObjectMapper(objectMapper);
        redisTemplate.setValueSerializer(jackson2JsonRedisSerializer);
        return redisTemplate;
    }
    /**
     * 定義lua腳本
     */
    @Bean
    public DefaultRedisScript<Long> limitScript(){
        DefaultRedisScript<Long> script = new DefaultRedisScript<>();
        script.setResultType(Long.class);//設(shè)置lua腳本返回值類型 需要同lua腳本中返回值一致
        script.setScriptSource(new ResourceScriptSource(new ClassPathResource("lua\\rateLimit.lua")));//讀取lua文件
        return script;
    }
}

5.自定義限流異常與全局異常

//自定義限流異常
public class RateLimitException extends RuntimeException{
? ? private Integer code;
? ? private String msg;
? ? public RateLimitException(Integer code, String msg) {
? ? ? ? this.code = code;
? ? ? ? this.msg = msg;
? ? }
? ? public Integer getCode() {
? ? ? ? return code;
? ? }
? ? public void setCode(Integer code) {
? ? ? ? this.code = code;
? ? }
? ? public String getMsg() {
? ? ? ? return msg;
? ? }
? ? public void setMsg(String msg) {
? ? ? ? this.msg = msg;
? ? }
}
//全局異常
@RestControllerAdvice
public class GlobalException {
? ? @ExceptionHandler(RateLimitException.class)
? ? public Map rateLimit(RateLimitException e){
? ? ? ? Map<String, Object> map = new HashMap<>();
? ? ? ? map.put("code",e.getCode());
? ? ? ? map.put("msg",e.getMsg());
? ? ? ? return map;
? ? }
}

6.創(chuàng)建切面，進行限流操作

@Aspect
@Component
public class RateLimitAspect {
    private static final Logger logger= LoggerFactory.getLogger(RateLimitAspect.class);
    @Autowired
    RedisTemplate<String,Object> redisTemplate;
    @Autowired
    RedisScript<Long> redisScript; //實現(xiàn)類為DefaultRedisScript<Long> limitScript()
    @Pointcut("@annotation(com.xxx.annotation.RateLimiter)")
    public void pointCut(){}
    @Before("pointCut()")
    public void beforeRateLimit(JoinPoint jp){
        //獲取RateLimiter注解上的值
        MethodSignature methodSignature = (MethodSignature) jp.getSignature();
        RateLimiter rateLimiter = AnnotationUtils.findAnnotation(methodSignature.getMethod(), RateLimiter.class);
        int time = rateLimiter.time();
        int count = rateLimiter.count();
        //構(gòu)建redis中的key值
        String rateKey=getRateLimitKey(rateLimiter,methodSignature);
        System.out.println("redis中key值："+rateKey);
        try {
            Long current = redisTemplate.execute(redisScript, Collections.singletonList(rateKey), time, count);
            if(current==null||current.intValue()>count){
                logger.info("當前接口達到最大限流次數(shù)");
                throw new RateLimitException(500,"當前接口達到最大限流次數(shù)");
            }
            logger.info("一段時間內(nèi)允許的請求次數(shù):{},當前請求次數(shù):{},緩存的key為:{}",count,current,rateKey);
        } catch (Exception e) {
            throw e;
        }
    }
    /**
     * redis中key兩種類型格式為：
     * 1.  rate_limit:com.xxx.controller.HelloController-hello
     * 2.  rate_limit:127.0.0.1-com.xxx.controller.HelloController-hello
     * @param rateLimiter
     * @param methodSignature
     * @return
     */
    private String getRateLimitKey(RateLimiter rateLimiter, MethodSignature methodSignature) {
        StringBuffer key = new StringBuffer(rateLimiter.key());
        if(rateLimiter.limitType()== LimitType.IP){//如果參數(shù)類型為IP
            //獲取客戶端ip
            String clientIP = ServletUtil.getClientIP(((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest());
            key.append(clientIP+"-");
        }
        Method method = methodSignature.getMethod();
        //獲取全類名
        String className = method.getDeclaringClass().getName();
        //獲取方法名
        String methodName = method.getName();
        key.append(className)
                .append("-")
                .append(methodName);
        System.out.println("全類名+方法名 "+className+"-"+methodName);
        return key.toString();
    }
}

7.創(chuàng)建接口，進行驗證

@RestController
public class HelloController {
    @GetMapping("/hello")
    @RateLimiter(time = 10,count = 3,limitType = LimitType.DEFAULT) //10秒內(nèi)允許訪問三次
    public String hello(){
        return "hello";
    }
}