自定义注解之单机版限流

自定义注解之单机版限流

0、前言

1、实现自定义注解

1.1、创建自定义注解

import java.lang.annotation.Documented;
import java.lang.annotation.ElementType;
import java.lang.annotation.Inherited;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

/**
 * @author ：
 * @date ：Created in 15:11 2022/9/1
 * @description ：自定义限流注解
 * @version: 1.0
 */
@Inherited
@Documented
@Target({ElementType.METHOD, ElementType.FIELD, ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
public @interface AccessLimit {
    // 限流方法名称
    String name() default "";
    // 每秒允许访问的次数10，底层令牌
    double permitsPerSecond() default 10;
}

1.2、通过AOP切面实现

import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.Signature;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.context.annotation.Scope;
import org.springframework.stereotype.Component;
import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Component
@Scope
@Aspect
@Slf4j
public class AccessLimitAspect {

    @Resource
    protected HttpServletRequest httpServletRequest;

    //设置令牌生成的频率
    private Map<String, RateLimiter> rateLimiters = new ConcurrentHashMap<>();

    /**
     *  定义切入点
     * @Author 
     * @Date Created in 16:44 2022/9/1
     * @Description
     * @param
     * @return void
     **/
    @Pointcut("@annotation(com.*.annotation.AccessLimit)")
    public void limit() {
    }

    @Around("limit()")
    public Object around(ProceedingJoinPoint proceedingJoinPoint) {
        Signature signature = proceedingJoinPoint.getSignature();
        MethodSignature methodSignature = (MethodSignature) signature;
        AccessLimit accessLimit = methodSignature.getMethod().getDeclaredAnnotation(AccessLimit.class);
        double permitsPerSecond = accessLimit.permitsPerSecond();
        String name = accessLimit.name();
        RateLimiter rateLimiter = rateLimiters.get(name);
        if (rateLimiter == null) {
            rateLimiter = RateLimiter.create(permitsPerSecond);
            rateLimiters.put(name, rateLimiter);
        }
        boolean flag = rateLimiter.tryAcquire();
        Object obj = null;
        if (flag) {
            try {
                obj = proceedingJoinPoint.proceed();
            } catch (Throwable throwable) {
                throwable.printStackTrace();
            }
        } else {
            log.warn("方法:{},{}", name, ExceptionCodeEnum.ACCESS_ERROR.getDesc());
            throw new BusinessException(ExceptionCodeEnum.ACCESS_ERROR);
        }
        return obj;
    }

}

1.3、使用

/**
 * @author ：
 * @date ：Created in 17:02 2022/9/2
 * @description ：
 * @version: 1.0
 */
@Slf4j
@RestController
@RequestMapping(value = "/user")
@RequiredArgsConstructor
@Api(tags = "用户管理")
public class UserController {

    private final IUserService userService;

    @GetMapping
    @Operation(summary = "条件查询")
    @AccessLimit(name = "findUsers", permitsPerSecond = 40)
    public Result<PageResult<UserResponse>> findUser(UserConditionReq userReq) {
        log.info("[user--通过条件查询用户信息，data:{}]", userReq);
        return Result.success(userService.findUsersByCondition(userReq));
    }
}