Spring-aop +redssion 实现接口限流注解

添加依赖环境

        <dependency>
            <groupId>org.redisson</groupId>
            <artifactId>redisson-spring-boot-starter</artifactId>
            <version>3.21.1</version>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-aop</artifactId>
        </dependency>
        .....

yml 配置redis

  data:
    redis:
      host: 127.0.0.1
      port: 6379

配置RedssionClient

package com.chenkang.demo.config;
 
import jakarta.annotation.Resource;
import org.redisson.Redisson;
import org.redisson.api.RedissonClient;
import org.redisson.config.Config;
import org.redisson.config.SingleServerConfig;
import org.springframework.boot.autoconfigure.data.redis.RedisProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
 
/**
 * @author chenkang
 * @since 2024/5/29 11:18
 */
@Configuration
public class RedisConfig {
 
 
    @Resource
    private RedisProperties redisProperties;
 
    @Bean
    public RedissonClient redisson() {
        Config config = new Config();
        SingleServerConfig singleServerConfig = config.useSingleServer();
        String address = "redis://" + redisProperties.getHost() + ":" + redisProperties.getPort();
        singleServerConfig.setAddress(address);
        return Redisson.create(config);
    }
 
}

spring 加入注解

@EnableAspectJAutoProxy

定义注解类型 以及切面

package com.chenkang.demo.aop;
 
/**
 * 接口限制类型
 * @author chenkang
 * @since 2024/5/29 21:09
 */
public enum LimitType {
 
    /**
     * 请求者IP  -- 通常是某个IP  不过会有人进行虚拟IP进行访问
     */
    IP,
    /**
     * 方法
     */
    METHOD,
 
    /**
     * 用户id
     */
    USER_ID;
}

package com.chenkang.demo.aop;
 
import java.lang.annotation.*;
 
/**
 * @author chenkang
 * @since 2024/5/29 21:07
 */
@Target(value = {ElementType.METHOD})
@Retention(value = RetentionPolicy.RUNTIME)
public @interface RateLimiter {
 
    /**
     * 给定时间范围
     */
    int period() default 5;
 
    /**
     * 给定次数
     */
    int count() default 2;
 
    /**
     * 接口限流类型
     */
    LimitType limitType() default LimitType.METHOD;
 
    /**
     * 限流类型 限流的类型  1 超时等待  2 快速返回
     */
    int actor() default 2;
 
    /**
     *
     * 一超时时间秒，当actor=1是生效
     */
    int timeout() default 3;
 
 
 
 
}

package com.chenkang.demo.aop;
 
import org.aopalliance.aop.Advice;
import org.springframework.aop.Pointcut;
import org.springframework.aop.support.AbstractPointcutAdvisor;
import org.springframework.aop.support.annotation.AnnotationMatchingPointcut;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.BeanFactory;
import org.springframework.beans.factory.BeanFactoryAware;
import org.springframework.core.Ordered;
 
/**
 * @author chenkang
 * @since 2024/5/29 21:04
 */
 
public class RateLimiterAnnotationAdvisor extends AbstractPointcutAdvisor implements BeanFactoryAware, Ordered {
 
    private Pointcut pointcut;
 
    private Advice advice;
 
 
    public RateLimiterAnnotationAdvisor(RateLimiterAnnotationInterceptor annotationInterceptor){
        this.pointcut=buildPointcut();
        this.advice=annotationInterceptor;
    }
 
    @Override
    public Pointcut getPointcut() {
        return pointcut;
    }
 
    @Override
    public Advice getAdvice() {
        return advice;
    }
 
    @Override
    public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
        if (this.advice instanceof BeanFactoryAware) {
            ((BeanFactoryAware) this.advice).setBeanFactory(beanFactory);
        }
    }
 
    private Pointcut buildPointcut() {
        return AnnotationMatchingPointcut.forMethodAnnotation(RateLimiter.class);
    }
 
    @Override
    public int getOrder() {
        return 1;
    }
}

package com.chenkang.demo.aop;
 
import com.chenkang.demo.util.IpUtils;
import lombok.Setter;
import org.aopalliance.intercept.MethodInterceptor;
import org.aopalliance.intercept.MethodInvocation;
import org.redisson.api.*;
import org.springframework.aop.BeforeAdvice;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
 
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.lang.reflect.Method;
import java.util.Optional;
import java.util.concurrent.TimeUnit;
 
/**
 * @author chenkang
 * @since 2024/5/29 21:14
 */
@Setter
public class RateLimiterAnnotationInterceptor implements MethodInterceptor, BeforeAdvice {
 
    private RedissonClient redissonClient;
 
 
    @Nullable
    @Override
    public Object invoke(@Nonnull MethodInvocation invocation) throws Throwable {
        //代理执行的真正地方
        Method method = invocation.getMethod();
        RateLimiter rateLimiter = method.getAnnotation(RateLimiter.class);
        String formatKey = formatKey(rateLimiter.limitType(), method);
        redissonRateLimiter(rateLimiter,formatKey);
        return invocation.proceed();
    }
 
    /**
     * 获取Redisson的RRateLimiter
     */
    private void redissonRateLimiter(RateLimiter limit, String key) {
        RRateLimiter rateLimiter = redissonClient.getRateLimiter(key);
        if (!rateLimiter.isExists()) {
            rateLimiter.trySetRate(RateType.OVERALL, limit.count(),limit.period(),  RateIntervalUnit.SECONDS);
        }
        RateLimiterConfig rateLimiterConfig = rateLimiter.getConfig();
        Long rateInterval = rateLimiterConfig.getRateInterval();
        Long rate = rateLimiterConfig.getRate();
        if (TimeUnit.MILLISECONDS.convert(limit.period(), TimeUnit.SECONDS) != rateInterval || limit.count() != rate) {
            rateLimiter.delete();
            rateLimiter.setRate(RateType.OVERALL, limit.count(), limit.period(), RateIntervalUnit.SECONDS);
        }
        if (limit.actor() == 1) {
            rateLimiter.acquire();
        }
        if (limit.actor() == 2) {
            if(rateLimiter.tryAcquire(1,limit.timeout(),TimeUnit.SECONDS)) {return;}
            throw new RuntimeException("请求频繁，请稍后再试！");
        }
 
    }
 
    private String formatKey(LimitType limitType,Method method) {
        if (limitType.equals(LimitType.IP)) {
            ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
            return Optional.ofNullable(attributes).map(ServletRequestAttributes::getRequest).map(IpUtils::getRealClientIp).orElse("");
        }else if(limitType.equals(LimitType.METHOD)){
            return  "limit:" + method.getClass().getSimpleName()+method.getName();
        }else {
            //TODO 获取当前登录人
            return "User";
        }
    }
 
 
}

package com.chenkang.demo.config;
 
import com.chenkang.demo.aop.RateLimiterAnnotationAdvisor;
import com.chenkang.demo.aop.RateLimiterAnnotationInterceptor;
import org.redisson.api.RedissonClient;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
 
/**
 * @author chenkang
 * @since 2024/5/29 21:21
 */
@Configuration
public class RateLimiterConfig {
 
    @Bean
    public RateLimiterAnnotationAdvisor lockAnnotationAdvisor(RateLimiterAnnotationInterceptor limiterAnnotationInterceptor) {
        return new RateLimiterAnnotationAdvisor(limiterAnnotationInterceptor);
    }
 
    @Bean
    public RateLimiterAnnotationInterceptor lockInterceptor(RedissonClient redisClient) {
        RateLimiterAnnotationInterceptor limiterAnnotationInterceptor = new RateLimiterAnnotationInterceptor();
        limiterAnnotationInterceptor.setRedissonClient(redisClient);
        return limiterAnnotationInterceptor;
    }
}

测试：

@GetMapping("rateLimiter")
@RateLimiter()
public String  rateLimiter() {
    System.out.println(111);
    return "hhah";
}

全局异常捕获返回就可以了