springboot项目利用redis脚本实现请求的限流

目的：限流 - 访问某请求的频次达到一定次数时，拒绝访问

1. 我们利用redis来记录频次，频次是value，key我们自定义，这里我们自定义考虑了两种场景，以Ip为key，或者，以请求的methodName为key，每次访问时，都以相同的key去redis中取value频次，当发现频次大于指定值时，抛出异常，拒绝执行后续逻辑。
2. 限流是针对请求的，也就是controller，所以我们采用aop包裹controller，这里以自定义注解的形式，标记controller作为pointCut。
3. 简单分析下key的两种自定义模式，如果以ip为key的话，限流作用点是具体的某一个用户（一定时间段内，这个ip的用户访问频次过高则限流），如果以方法名为key的话，限流的作用点是这个方法（一定时间段内，多个用户访问频次总和过高则限流）
4. 上面说的“一定时间段内”，其实就是在redis中给了这个key一个有效时间（或者叫生存时间），当这个key被建立，也就是key的value=1时，给他一个有效时间，这个时间，就是key的生命时间段。

@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.CUSTOM;

}

public enum LimitType {
    // 默认
    CUSTOM,
    // ip限流
    IP
}

@Aspect
@Component
public class LimitAspect {

    private final RedisTemplate<Object,Object> redisTemplate;
    private static final Logger logger = LoggerFactory.getLogger(LimitAspect.class);

    public LimitAspect(RedisTemplate<Object, Object> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    @Pointcut("@annotation(com.deacy.shop.annotation.Limit)")
    public void pointcut(){
    }

    public Object around(ProceedingJoinPoint joinPoint) throws Throwable{
        HttpServletRequest request = RequestHolderUtil.getHttpServletRequest();
        MethodSignature methodSignature = (MethodSignature) joinPoint.getSignature();
        Method method = methodSignature.getMethod();
        Limit limit = method.getAnnotation(Limit.class);
        LimitType limitType = limit.limitType();
        String key = limit.key();
        // 若key为空，对其赋值 IP类型赋ip CUSTOM类型赋方法名
        if(StringUtils.isEmpty(key)){
            if(limitType == LimitType.IP){
                key = IpUtil.getIp(request);
            }else {
                key = method.getName();
            }
        }
        // 准备执行redis脚本
        String el = org.apache.commons.lang3.StringUtils.join(limit.prefix(), "_", key, "_", request.getRequestURI().replaceAll("/", "_"));
        List<Object> keys = Arrays.asList(el);
        String luaScript = buildLuaScript();
        DefaultRedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
        // 执行redis脚本，period时间段内访问频次count过高会被限流
        Number frequency = redisTemplate.execute(redisScript, keys, limit.count(), limit.period());
        if(frequency!=null && frequency.intValue()<=limit.count()){
            logger.info("第{}次访问key为 {}，描述为 [{}] 的接口", frequency, keys.get(0), limit.name());
            return joinPoint.proceed();
        }else {
            throw new BadRequestException("请求访问次数受限");
        }
    }

    /**
     * 限流脚本（一定时间段内访问频次过高会被限流）
     * 如果指定key的值比限流次数大，就返回c
     * 将key值增加1（incr没有的话会自动新建的）
     * 如果key值==1的话，设置其有效时间
     * 最后返回c
     */
    private 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;";
    }
}

@Api("限流demo")
@RestController
@RequestMapping("/limit")
public class LimitController {

    private static final AtomicInteger ATOMIC_INTEGER = new AtomicInteger();

    // limet注解中key如果是空的则会根据limitType类型自动对key进行赋值，这部分逻辑在LimitAspect中
    @ApiOperation("测试")
    @GetMapping("/test")
    @Limit(period=60,count=3,name="testLimit",prefix = "p",limitType = LimitType.IP)
    public int test(){
        return ATOMIC_INTEGER.incrementAndGet();
    }

}

本文地址：https://blog.csdn.net/DX_dixi/article/details/110230388