使用Java自定義注解實現(xiàn)一個簡單的令牌桶限流器

更新時間：2023年10月08日 08:30:59 作者：酸奶小肥陽

限流是在分布式系統(tǒng)中常用的一種策略,它可以有效地控制系統(tǒng)的訪問流量,保證系統(tǒng)的穩(wěn)定性和可靠性,在本文中,我將介紹如何使用Java自定義注解來實現(xiàn)一個簡單的令牌桶限流器,需要的朋友可以參考下

什么是令牌桶限流？

令牌桶限流是一種常用的限流算法，它基于令牌桶的概念。在令牌桶中，令牌以固定的速率被生成并放置其中。當(dāng)一個請求到達(dá)時，它必須獲取一個令牌才能繼續(xù)執(zhí)行，否則將被阻塞或丟棄。

開始我們的實現(xiàn)

第一步：創(chuàng)建一個自定義注解

我們首先需要創(chuàng)建一個自定義注解，用于標(biāo)識需要進(jìn)行限流的方法。這個注解可以命名為@RateLimit，它可以帶有以下幾個參數(shù)

rate: 表示該方法的限流速率，單位可以是每秒請求數(shù)（QPS）。
prefixKey: 針對不同方法上對同一個資源做限流的情況。
target: 限流的對象，默認(rèn)使用spEl表達(dá)式對入?yún)⑦M(jìn)行獲取
capacity: 令牌桶容量，滿了之后令牌不再增加

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimit {
    /**
     * key的前綴,默認(rèn)取方法全限定名，除非我們在不同方法上對同一個資源做頻控，就自己指定
     */
    String prefixKey() default "";
    /**
     * 選擇目標(biāo)類型: 1.EL，需要在spEl()中指定限流資源。2.USER，針對用戶進(jìn)行限流
     */
    Target target() default Target.EL;
    /**
     * springEl表達(dá)式 指定頻控對象
     */
    String spEl() default "";
    /**
     * 令牌桶容量
     */
    double capacity() default 10;
    /**
     * 令牌生成速率 n/秒
     */
    double rate() default 1;
    enum Target {
        EL, USER
    }
}

第二步：實現(xiàn)限流邏輯

接下來，我們需要編寫一個類來處理限流邏輯。這個類可以命名為RateLimitAspect，它將會掃描所有被@RateLimit注解標(biāo)記的方法，并在必要時進(jìn)行限流。

/**
 * 令牌桶限流
 *
 * @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();
            // 如果生成的令牌超過的桶最大容量，那么令牌數(shù)取桶最大容量
            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過期時間設(shè)置大于 容量/速率
            redisTemplate.opsForValue().set(k, var2, timeout, TimeUnit.SECONDS);
        });
        return joinPoint.proceed();
    }
}

SpEL解析工具類

@Component
public class SpELUtil {
    /**
     * 獲取表達(dá)式中的參數(shù)值
     *
     * @param expr      表達(dá)式
     * @param joinPoint 切點
     * @return 參數(shù)值
     */
    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);
    }
    /**
     * 獲取方法的參數(shù)名稱
     *
     * @param joinPoint 切點
     * @return 參數(shù)名稱
     */
    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注解

現(xiàn)在，我們可以在需要進(jìn)行限流的方法上使用@RateLimited注解，指定相應(yīng)的限流速率。

示例1: 限制了令牌桶容量10，每10秒生成一個令牌，限制對象為當(dāng)前用戶。

@Api(tags = "項目服務(wù)")
@Validated
@Slf4j
@RestController
@RequestMapping("/api/v1")
@RequiredArgsConstructor
public class ProjectController {
    @Resource
    private IProjectService projectService;
    @ApiOperation("創(chuàng)建項目")
    @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("數(shù)據(jù)導(dǎo)出")
    @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();
    }
}