JAVA自定义注解实现接口/ip限流

基本步骤：

1.自定义注解 @RateLimiter，添加属性：key（redis中key前缀）、time（redis中key过期时间，单位秒）、count（time中定义时间段内可以访问的次数），limitType（限流类型：1.限制接口；2.限制ip）

2.创建AOP，定义前置通知，获取@RateLimiter注解上的各项值，客户端每访问一次接口，redis中存储的current（当前时间段内已访问的次数）值+1，若current值大于count规定的值，抛出异常，提示到达允许访问的次数；

具体代码如下：

        <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.创建枚举类，定义限流类型

public enum LimitType {
    /**
     * 默认限流策略，针对某一个接口进行限流
     */
    DEFAULT,
    /**
     * 根据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;

    /**
     * 限流时间内限流次数
     * @return
     */
    int count() default 100;

    /**
     * 限流类型：1.限制接口访问次数 2.限制ip访问次数
     * @return
     */
    LimitType limitType() default LimitType.DEFAULT;
}

3.创建lua脚本，保证redis中操作原子性（在resources/lua目录下创建rateLimit.lua文件）

---定义变量：redis中key值,redis中过期时间,规定的时间段内访问次数,当前访问次数
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不为空当前访问次数大于规定时间段内的访问次数
    return tonumber(current)  ---返回当前访问的次数
end
    current = redis.call('incr',key)    --- 若没超过限制次数,当前次数+1（自增）
if tonumber(current)==1 then
    redis.call('expire',key,time)       --- 如果当前访问次数=1,表示第一次访问,设置当前key的过期时间
end
    return tonumber(current)        --返回tonumber(current)

4.创建配置类，解决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);
        //序列化时将类的数据类型存入json，以便反序列化的时候转换成正确的类型
        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);//设置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.创建切面，进行限流操作

@Aspect
@Component
public class RateLimitAspect {
    private static final Logger logger= LoggerFactory.getLogger(RateLimitAspect.class);

    @Autowired
    RedisTemplate<String,Object> redisTemplate;
    @Autowired
    RedisScript<Long> redisScript; //实现类为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();
        //构建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("当前接口达到最大限流次数");
                throw new RateLimitException(500,"当前接口达到最大限流次数");
            }
            logger.info("一段时间内允许的请求次数:{},当前请求次数:{},缓存的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){//如果参数类型为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.创建接口，进行验证

@RestController
public class HelloController {

    @GetMapping("/hello")
    @RateLimiter(time = 10,count = 3,limitType = LimitType.DEFAULT) //10秒内允许访问三次
    public String hello(){
        return "hello";
    }
}

gitee开源地址：

RateLimiter: JAVA自定义注解实现接口/ip限流