java

关注公众号 jb51net

关闭
首页 > 软件编程 > java > Java实现令牌桶限流

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

作者：酸奶小肥阳

限流是在分布式系统中常用的一种策略,它可以有效地控制系统的访问流量,保证系统的稳定性和可靠性,在本文中,我将介绍如何使用Java自定义注解来实现一个简单的令牌桶限流器,需要的朋友可以参考下

什么是令牌桶限流？

令牌桶限流是一种常用的限流算法，它基于令牌桶的概念。在令牌桶中，令牌以固定的速率被生成并放置其中。当一个请求到达时，它必须获取一个令牌才能继续执行，否则将被阻塞或丢弃。

开始我们的实现

第一步：创建一个自定义注解

我们首先需要创建一个自定义注解，用于标识需要进行限流的方法。这个注解可以命名为@RateLimit，它可以带有以下几个参数

@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
    }
}

第二步：实现限流逻辑

接下来，我们需要编写一个类来处理限流逻辑。这个类可以命名为RateLimitAspect，它将会扫描所有被@RateLimit注解标记的方法，并在必要时进行限流。

/**
 * 令牌桶限流
 *
 * @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();
    }
}

SpEL解析工具类

@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;
    }
}

第三步：在方法上使用@RateLimit注解

现在，我们可以在需要进行限流的方法上使用@RateLimited注解，指定相应的限流速率。

示例1: 限制了令牌桶容量10，每10秒生成一个令牌，限制对象为当前用户。

@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);
    }
}

示例2: 限制了令牌桶容量1，每2秒生成一个令牌，限制对象为该项目。

@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();
    }
}

总结

通过使用Java自定义注解，我们成功地实现了一个简单的令牌桶限流器。这个限流器可以方便地应用于需要对访问速率进行控制的方法中，保证系统的稳定性和可靠性。

在实际项目中，我们可以根据需求对限流器进行进一步地扩展和优化，以满足不同场景下的限流需求。希望本文对你理解和实现限流算法有所帮助！

到此这篇关于基于Java自定义注解实现一个令牌桶限流的文章就介绍到这了,更多相关Java实现令牌桶限流内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家！

您可能感兴趣的文章:
阅读全文