使用Aop+Redis+lua限流，优化高并发问题

有一只柴犬

发布于 2024-01-25 10:51:11

1870

发布于 2024-01-25 10:51:11

文章被收录于专栏：JAVA体系

限流的方式有很多： 1、单机模式下，可以使用AtomicInteger、RateLimiter、Semaphore。 2、分布式下，可以使用队列（如Kafka等），但是编码比较繁杂；也可以使用Nginx限流，但是属于网关层面，不能解决所有问题（如内部服务接口）。所以，应用层也是需要做限流操作的。这里简单结合Aop+redis+lua来实现。注：如果是需要接入层先流的话，建议还是要使用nginx自带的连接数限流模块和请求限流模块。 Lua脚本：

    /**
     * 限流脚本
     */
    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;";
    }

1、KEYS[1]获取传入的keys参数，（这里为redis的键key） 2、ARGV[1]获取到传入的limit参数，（这里为请求的token数量，也可以理解为次数） 3、ARGV[2]获取到传入的limit参数，（这里为使用的限流key的过期时间） 限流注解，Limit

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

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

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

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

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

    // 限制访问次数
    int count();

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

}

切面Aspect

@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.xx.xx.Limit)")
    public void pointcut() {
    }

    @Around("pointcut()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        HttpServletRequest request = getHttpServletRequest();
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method signatureMethod = signature.getMethod();
        Limit limit = signatureMethod.getAnnotation(Limit.class);
        String key = limit.key();

        ImmutableList<Object> keys = ImmutableList.of(StringUtils.join(limit.prefix(), "_", key, "_", request.getRequestURI().replaceAll("/","_")));

        String luaScript = buildLuaScript();
        RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
        Number count = redisTemplate.execute(redisScript, keys, limit.count(), limit.period());
        if (null != count && count.intValue() <= limit.count()) {
            logger.info("第{}次访问key为 {}，描述为 [{}] 的接口", count, keys, limit.name());
            return joinPoint.proceed();
        } else {
            throw new BadRequestException("访问次数受限制");
        }
    }
    
    public HttpServletRequest getHttpServletRequest() {
        return ((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder.getRequestAttributes())).getRequest();
    }
}

测试

/**
 * @author shamee
 * 限流测试
 */
@RestController
@RequestMapping("/api/limit")
public class LimitController {

    private static final AtomicInteger ATOMIC_INTEGER = new AtomicInteger();

    /**
     * 测试限流注解， 
     * @param period 过期时间
     * @param count 次数
     * @param name 接口描述
     * 
     * 60秒内最多只能访问 5次，保存到redis的键名为 limitP_testK，
     */
    @GetMapping
    @Limit(key = "testK", period = 60, count = 5, name = "test_limit", prefix = "limitP")
    public int test_limit() {
        return ATOMIC_INTEGER.incrementAndGet();
    }
}

简单备注，以便日后查阅。

本文参与腾讯云自媒体同步曝光计划，分享自作者个人站点/博客。

原始发表：2024-01-25，如有侵权请联系 cloudcommunity@tencent.com 删除

高并发