Examples with QuotaResult com.tencent.polaris.api.plugin.ratelimiter.QuotaResult used on opensource projects

Example 6 with QuotaResult

use of com.tencent.polaris.api.plugin.ratelimiter.QuotaResult in project polaris-java by polarismesh.

the class RemoteAwareBucket method allocateQuota.

@Override
public QuotaResult allocateQuota(long curTimeMs, int token) throws PolarisException {
    int stopIdx = -1;
    TokenBucketMode mode = TokenBucketMode.UNKNOWN;
    UpdateIdentifier[] identifiers = new UpdateIdentifier[tokenBuckets.size()];
    AllocateResult[] results = new AllocateResult[tokenBuckets.size()];
    for (int i = 0; i < tokenBuckets.size(); i++) {
        identifiers[i] = flowCache.borrowThreadCacheObject(UpdateIdentifier.class);
        results[i] = flowCache.borrowThreadCacheObject(AllocateResult.class);
    }
    for (int i = 0; i < tokenBuckets.size(); i++) {
        TokenBucket tokenBucket = tokenBuckets.get(i);
        AllocateResult allocateResult = tokenBucket.tryAllocateToken(mode, token, curTimeMs, identifiers[i], results[i]);
        mode = allocateResult.getMode();
        if (allocateResult.getLeft() < 0) {
            stopIdx = i;
            break;
        }
    }
    QuotaResult response;
    boolean usedRemoteQuota = mode == TokenBucketMode.REMOTE;
    if (stopIdx >= 0) {
        // 有一个扣除不成功，则进行限流
        TokenBucket tokenBucket = tokenBuckets.get(stopIdx);
        if (usedRemoteQuota) {
            // 远程才记录滑窗, 滑窗用于上报
            tokenBucket.confirmLimited(token, curTimeMs);
        }
        // 归还配额
        for (int i = 0; i < stopIdx; i++) {
            TokenBucket bucket = tokenBuckets.get(i);
            bucket.giveBackToken(identifiers[i], token, mode);
        }
        response = new QuotaResult(QuotaResult.Code.QuotaResultLimited, 0, "");
    } else {
        // 记录分配的配额
        for (TokenBucket tokenBucket : tokenBuckets) {
            if (usedRemoteQuota) {
                tokenBucket.confirmPassed(token, curTimeMs);
            }
        }
        response = new QuotaResult(QuotaResult.Code.QuotaResultOk, 0, "");
    }
    // 归还对象
    for (int i = 0; i < tokenBuckets.size(); i++) {
        flowCache.giveBackThreadCacheObject(identifiers[i]);
        flowCache.giveBackThreadCacheObject(results[i]);
    }
    return response;
}