两种白名单限流方案（redis lua限流，guava方案）

两种白名单限流方案

1、redis lua方案

创建注解类

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface Limit {
    

    // 资源名称，用于描述接口功能
    String name() default "";

    // 资源 key
    String key() default "";

    // key 前缀
    String prefix() default "";

    // 时间单位秒
    int period();

    // 限制单位时间访问的次数
    int count();

    // 限制类型
    LimitType limitType() default LimitType.CUSTOMER;
}

编写AOP拦截类

@Slf4j
@Aspect
@Component
public class LimitAspect {
    

    @Autowired
    @Qualifier("jacksonRedisTemplate")
    private RedisTemplate<String,Object> redisTemplate;


    @Pointcut("@annotation(com.xx.common.annotation.Limit)")
    public void pointcut() {
    
    }

    @Around("pointcut()")
    public Object around(ProceedingJoinPoint point) throws Throwable {
    
        HttpServletRequest request = ((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder.getRequestAttributes())).getRequest();

        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        Limit limitAnnotation = method.getAnnotation(Limit.class);
        LimitType limitType = limitAnnotation.limitType();
        String name = limitAnnotation.name();
        String key;
        String ip = IpUtil.getIpAddr(request);
        int limitPeriod = limitAnnotation.period();
        int limitCount = limitAnnotation.count();
        switch (limitType) {
    
            case IP:
                key = ip;
                break;
            case CUSTOMER:
                key = limitAnnotation.key();
                break;
            default:
                key = StringUtils.upperCase(method.getName());
        }
        ImmutableList<String> keys = ImmutableList.of(StringUtils.join(limitAnnotation.prefix() + "_", key, ip));
        String luaScript = buildLuaScript();
        RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
        Number count = redisTemplate.execute(redisScript, keys, limitCount, limitPeriod);
        log.info("限流拦截IP:{} 第 {} 次访问key={}，接口描述:{}", ip, count, keys, name);
        if (count != null && count.intValue() <= limitCount) {
    
            return point.proceed();
        } else {
    
            throw new LimitAccessException("接口访问超出频率限制");
        }
    }

    /** * 限流脚本 * 调用的时候不超过阈值，则执行计算器自加。 * * @return lua脚本 */
    private static String buildLuaScript() {
    
        return "local c" +
                "\nc = redis.call('get',KEYS[1])" +
                "\nif c and tonumber(c) > tonumber(ARGV[1]) then" +
                "\nreturn c;" +
                "\nend" +
                "\nc = redis.call('incr',KEYS[1])" +
                "\nif tonumber(c) == 1 then" +
                "\nredis.call('expire',KEYS[1],ARGV[2])" +
                "\nend" +
                "\nreturn c;";
    }
}

枚举

 public enum LimitType {
    
    // 传统类型
    CUSTOMER,
    // 根据 IP 限制
    IP;
}
 

获取IP工具类

@Slf4j
public class IpUtil {
    

    private static final String UNKNOWN = "unknown";
    private static final String LOCAL_IP = "127.0.0.1";
    private static final Integer IP_LENGTH = 15;

    protected IpUtil() {
    

    }

    /** * 获取 IP地址 * 使用 Nginx等反向代理软件， 则不能通过 request.getRemoteAddr()获取 IP地址 * 如果使用了多级反向代理的话，X-Forwarded-For的值并不止一个，而是一串IP地址， * X-Forwarded-For中第一个非 unknown的有效IP字符串，则为真实IP地址 */
    public static String getIpAddr(HttpServletRequest request) {
    
        String ip = request.getHeader("x-forwarded-for");
        if (ip == null || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ip)) {
    
            ip = request.getHeader("Proxy-Client-IP");
        }
        if (ip == null || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ip)) {
    
            ip = request.getHeader("WL-Proxy-Client-IP");
        }
        if (ip == null || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ip)) {
    
            ip = request.getRemoteAddr();
            if (LOCAL_IP.equals(ip)) {
    
                // 根据网卡取本机配置的IP
                InetAddress inet = null;
                try {
    
                    inet = InetAddress.getLocalHost();
                } catch (UnknownHostException e) {
    
                    log.error("获取ip错误:{}", e);
                }
                if (inet != null) {
    
                    ip = inet.getHostAddress();
                }
            }
        }
        if (ip != null && ip.length() > IP_LENGTH) {
    
            if (ip.indexOf(StrUtil.COMMA) > 0) {
    
                ip = ip.substring(0, ip.indexOf(","));
            }
        }
        return "0:0:0:0:0:0:0:1".equals(ip) ? LOCAL_IP : ip;
    }
}

2、guava方案

引入guava包

			<!--guava 限流-->
			<dependency>
				<groupId>com.google.guava</groupId>
				<artifactId>guava</artifactId>
				<version>31.0.1-jre</version>
			</dependency>

自定义注解

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimit {
    

    int NOT_LIMITED = 0;

    /** * qps */
    double qps() default NOT_LIMITED;

    /** * 超时时长 */
    int timeout() default 0;

    /** * 超时时间单位 */
    TimeUnit timeUnit() default TimeUnit.MILLISECONDS;

    /** * 超时时间单位 */
    LimitType limitType() default LimitType.CUSTOMER;
}  

AOP拦截：

@Slf4j
@Aspect
@Component
public class RateLimitAspect {
    
    /** * map */
    private static final ConcurrentMap<String, RateLimiter> RATE_LIMITER_CACHE = new ConcurrentHashMap<>();


    @Pointcut("@annotation(com.xx.common.annotation.RateLimit)")
    public void pointcut() {
    
    }

    @Around("pointcut()")
    public Object around(ProceedingJoinPoint point) throws Throwable {
    
        log.info("guava限流拦截到了{}方法...", point.getSignature().getName());
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();

        if (method.isAnnotationPresent(RateLimit.class)) {
    
            //获取方法上的注解
            RateLimit rateLimit = method.getAnnotation(RateLimit.class);

            if (rateLimit != null && rateLimit.qps() > 0) {
    
                double qps = rateLimit.qps();
                if (RATE_LIMITER_CACHE.get(method.getName()) == null) {
    
                    // 初始化 QPS
                    RATE_LIMITER_CACHE.put(method.getName(), com.google.common.util.concurrent.RateLimiter.create(qps));
                }

                log.debug("【{}】的QPS设置为: {}", method.getName(), RATE_LIMITER_CACHE.get(method.getName()).getRate());
                // 尝试获取令牌
                if (RATE_LIMITER_CACHE.get(method.getName()) != null && !RATE_LIMITER_CACHE.get(method.getName()).tryAcquire(rateLimit.timeout(), rateLimit.timeUnit())) {
    
                    throw new LimitAccessException("接口访问超出频率限制");
                }
                return point.proceed();
            }
        }
        return point.proceed();
    }
}