Spring Boot 项目中使用自定义注解和切面获取客户端 IP 并限流

背景

上一篇【文章】中，为了基于 IP 限流，我们使用了 Sentinel 的热点参数限流功能，效果也能够达到预期，但是我们的代码侵入性太强，后期如果我们不想使用 Sentinel 来做限流，则需要修改大量代码，而且也不利于维护。

实际上，但我们使用 Sentinel 的热点参数限流功能时，很多地方的编码框架都是，都是先去获取资源，当遇到限流异常时，返回错误码 429 ，我们可以利用AOP切面编程思想，使用自定义的注解，对于需要限流的方法，统一编写限流逻辑，实现解耦。

优化限流

定义注解

定义注解前，我们需要思考一下，我们需要在使用注解时提供哪些参数，回想一下，使用 Sentinel 的热点参数限流时，我们需要编写的代码为：

try {
        // 尝试获取资源
        entry = SphU.entry(resourceName, entryType, batchCount, Object... args);
        // 被保护逻辑
        // TODO ...... 

    } catch (Throwable ex) {
        
        // 业务异常
        if (!BlockException.isBlockException(ex)) {
            Tracer.trace(ex);
        }
        // 降级操作
        if (ex instanceof DegradeException) {
            // 处理降级逻辑
        }
        // 限流操作，可以直接返回 HttpStatus.TOO_MANY_REQUESTS
        
    } finally {
        if (entry != null) {
            entry.exit(1, remoteAddr);
        }
    }

因此我们可以抽象出resourceName, entryType, batchCount是必须项，还有一个是 IP ，由于 IP 是通过HttpServletRequest对象来获取的，虽然说我们可以通过改造需要限流的接口方法都接受这个变量，但是又会导致侵入性太强，因此我们抛弃这种方法获取 IP ，而是通过RequestContextHolder来获取HttpServletRequest对象。

注解定义如下：

/**
 * @author Yalexin
 * 用于基于 IP 限流
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface SentinelIpLimiter {
    // 资源名字
    String value() default "";
    // 类型
    EntryType entryType() default EntryType.OUT;

    int batchCount () default 1;
}

定义切面处理逻辑

由于我们要在切面中返回 429 状态码，因此需要使用@Around来完全控制方法的执行流程：

/**
 * Author: Yalexin
 * SentinelIpLimiter 注解的切面逻辑
 */
@Aspect
@Component
public class SentinelLimiterAspect {
    private final Logger logger = LoggerFactory.getLogger(this.getClass());

    @Around("@annotation(sentinelLimiter)")
    public Object aroundAdvice(ProceedingJoinPoint joinPoint, SentinelIpLimiter sentinelLimiter) throws Throwable {
        // 获取注解的属性值
        String resourceName = sentinelLimiter.value();
        EntryType entryType = sentinelLimiter.entryType();
        int batchCount = sentinelLimiter.batchCount();

        // 获取当前请求的属性
        RequestAttributes requestAttributes = RequestContextHolder.getRequestAttributes();
        if (requestAttributes instanceof ServletRequestAttributes) {
            // 从请求属性中获取HttpServletRequest对象
            HttpServletRequest request = ((ServletRequestAttributes) requestAttributes).getRequest();
            // 获取 IP
            String remoteAddr = IPUtils.getIRealIPAddr(request);
            Entry entry = null;
            try {
                // 尝试获取资源
                entry = SphU.entry(resourceName, entryType, batchCount, remoteAddr);
                // 如果成功获得资源，则继续执行目标方法
                return joinPoint.proceed();

            } catch (Throwable ex) {
                logger.debug("block by "+ex.getMessage());
                // 业务异常
                if (!BlockException.isBlockException(ex)) {
                    Tracer.trace(ex);
                    logger.error(ex.getMessage());
                    return new ResponseEntity(null, HttpStatus.OK);
                }
                // 降级操作
                if (ex instanceof DegradeException) {
                    // 处理降级逻辑
                    logger.warn(ex.getMessage());
                    return new ResponseEntity(null, HttpStatus.OK);
                }
                // 限流操作，直接返回 HttpStatus.TOO_MANY_REQUESTS
                return new ResponseEntity(null, HttpStatus.TOO_MANY_REQUESTS);
            } finally {
                if (entry != null) {
                    entry.exit(1, remoteAddr);
                }
            }
        }
        // 如果没有匹配到Servlet请求，让其继续执行（或根据需要返回默认响应）
        return joinPoint.proceed();
    }
}

使用我们的自定义注解

在需要被定义为资源的接口方法上打上我们注解：

@SentinelIpLimiter(
        value = SentinelConstant.ADMIN_LOGIN_RULE,
        entryType = EntryType.IN,
        batchCount = 1
)
@PostMapping("/login")
public ResponseEntity login(HttpServletRequest request,
                            HttpServletResponse response,
                            @RequestBody JSONObject json) {
    return tryLogin(request, response, json);
}