使用Java自定义注解实现一个简单的令牌桶限流器

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimit {
    /**
     * key的前缀,默认取方法全限定名，除非我们在不同方法上对同一个资源做频控，就自己指定
     */
    String prefixKey() default "";
    /**
     * 选择目标类型: 1.EL，需要在spEl()中指定限流资源。2.USER，针对用户进行限流
     */
    Target target() default Target.EL;
    /**
     * springEl表达式 指定频控对象
     */
    String spEl() default "";
    /**
     * 令牌桶容量
     */
    double capacity() default 10;
    /**
     * 令牌生成速率 n/秒
     */
    double rate() default 1;
    enum Target {
        EL, USER
    }
}

/**
 * 令牌桶限流
 *
 * @date 2023/07/07
 */
@Aspect
@Component
@Slf4j
@RequiredArgsConstructor
public class RateLimitAspect {
    @Resource
    private RedisTemplate<String, BucketLog> redisTemplate;
    @Resource
    private RbacUserService rbacUserService;
    private final SpELUtil spELUtil;
    @Around("@annotation(com.netease.fuxi.config.annotation.RateLimit)")
    public Object doAround(ProceedingJoinPoint joinPoint) throws Throwable {
        Method method = ((MethodSignature) joinPoint.getSignature()).getMethod();
        RateLimit[] annotations = method.getAnnotationsByType(RateLimit.class);
        var var1 = new HashMap<String, RateLimit>();
        for (int i = 0; i < annotations.length; i++) {
            RateLimit annotation = annotations[i];
            final String prefix = StringUtils.isBlank(annotation.prefixKey()) ?
                    method.getDeclaringClass() + "#" + method.getName() + ":index:" + i : annotation.prefixKey();
            String key = "";
            switch (annotation.target()) {
                case EL:
                    key = spELUtil.getArgValue(annotation.spEl(), joinPoint);
                    break;
                case USER:
                    key = rbacUserService.getCurrentUser();
            }
            var1.put(prefix + ":" + key, annotation);
        }
        var1.forEach((k, v) -> {
            var var2 = Boolean.TRUE.equals(redisTemplate.hasKey(k)) ? redisTemplate.opsForValue().get(k) :
                    new BucketLog(v.capacity(), Instant.now().getEpochSecond());
            long nowTime = Instant.now().getEpochSecond();
            double addTokens = (nowTime - var2.getLastRefillTime()) * v.rate();
            // 如果生成的令牌超过的桶最大容量，那么令牌数取桶最大容量
            var2.setTokens(Math.min(var2.getTokens() + addTokens, v.capacity()));
            var2.setLastRefillTime(nowTime);
            double remain = var2.getTokens() - 1;
            if (remain < 0) {
                throw new BusinessException("操作太频繁，请稍后重试", 42901);
            }
            var2.setTokens(remain);
            long timeout = (long) Math.ceil(v.capacity() / v.rate());// redis过期时间设置大于 容量/速率
            redisTemplate.opsForValue().set(k, var2, timeout, TimeUnit.SECONDS);
        });
        return joinPoint.proceed();
    }
}

@Component
public class SpELUtil {
    /**
     * 获取表达式中的参数值
     *
     * @param expr      表达式
     * @param joinPoint 切点
     * @return 参数值
     */
    public String getArgValue(String expr, JoinPoint joinPoint) {
        Object[] args = joinPoint.getArgs();
        String[] parameterNames = getParameterNames(joinPoint);
        EvaluationContext context = new StandardEvaluationContext();
        for (int i = 0; i < parameterNames.length; i++) {
            context.setVariable(parameterNames[i], args[i]);
        }
        ExpressionParser parser = new SpelExpressionParser();
        return parser.parseExpression(expr).getValue(context, String.class);
    }
    /**
     * 获取方法的参数名称
     *
     * @param joinPoint 切点
     * @return 参数名称
     */
    private String[] getParameterNames(JoinPoint joinPoint) {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        String[] parameterNames = signature.getParameterNames();
        if (parameterNames == null || parameterNames.length == 0) {
            ParameterNameDiscoverer parameterNameDiscoverer = new LocalVariableTableParameterNameDiscoverer();
            parameterNames = parameterNameDiscoverer.getParameterNames(signature.getMethod());
        }
        return parameterNames;
    }
}

@Api(tags = "项目服务")
@Validated
@Slf4j
@RestController
@RequestMapping("/api/v1")
@RequiredArgsConstructor
public class ProjectController {
    @Resource
    private IProjectService projectService;
    @ApiOperation("创建项目")
    @PostMapping("/project")
    @RateLimit(capacity = 10, rate = 0.1, target = RateLimit.Target.USER)
    public Result<ProjectVO> createProject() {
        ProjectVO projectVO = projectService.createProject();
        return Result.ok(projectVO);
    }
}

@Slf4j
@RestController
@RequestMapping("/api/v1")
public class ProjectController {
    @Resource
    private IProjectService projectService;
    @ApiOperation("数据导出")
    @PostMapping("/project/{projectId}/export")
    @RateLimit(capacity = 1, rate = 0.5, spEl = "#projectId")
    public Result<Void> export(@PathVariable String projectId,
                               @RequestBody @Valid ExportDTO exportDTO) {
        projectService.export(projectId, exportDTO);
        return Result.ok();
    }
}