Examples with QuotaName com.github.ambry.quota.QuotaName used on opensource projects

Example 1 with QuotaName

use of com.github.ambry.quota.QuotaName in project ambry by linkedin.

the class AmbryCUQuotaEnforcerTest method testCharge.

@Test
public void testCharge() throws Exception {
    Container container = ACCOUNT.getAllContainers().iterator().next();
    Map<QuotaName, Double> readRequestCostMap = Collections.singletonMap(QuotaName.READ_CAPACITY_UNIT, 9.0);
    Map<QuotaName, Double> writeRequestCostMap = Collections.singletonMap(QuotaName.WRITE_CAPACITY_UNIT, 9.0);
    Map<String, CapacityUnit> usageMap = QUOTA_SOURCE.getAllQuotaUsage();
    // 1. Test that usage is updated and recommendation is serve when usage is within limit.
    RestRequest restRequest = QuotaTestUtils.createRestRequest(ACCOUNT, container, RestMethod.GET);
    AMBRY_QUOTA_ENFORCER.charge(restRequest, readRequestCostMap);
    QuotaRecommendation quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(restRequest);
    assertEquals(quotaRecommendation.getQuotaName(), QuotaName.READ_CAPACITY_UNIT);
    assertEquals(quotaRecommendation.getQuotaUsagePercentage(), 90, 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.ALLOW);
    // make sure that correct quota is charged.
    assertEquals(usageMap.get(String.valueOf(ACCOUNT.getId())).getWcu(), 0);
    restRequest = QuotaTestUtils.createRestRequest(ACCOUNT, container, RestMethod.POST);
    AMBRY_QUOTA_ENFORCER.charge(restRequest, writeRequestCostMap);
    quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(restRequest);
    assertEquals(quotaRecommendation.getQuotaName(), QuotaName.WRITE_CAPACITY_UNIT);
    assertEquals(quotaRecommendation.getQuotaUsagePercentage(), 90, 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.ALLOW);
    // make sure that correct quota is charged.
    assertEquals(usageMap.get(String.valueOf(ACCOUNT.getId())).getRcu(), 9);
    // 2. Test that retryable quota exception is thrown when quota not found.
    Account newAccount = ACCOUNT_SERVICE.generateRandomAccount(QuotaResourceType.ACCOUNT);
    restRequest = QuotaTestUtils.createRestRequest(newAccount, newAccount.getAllContainers().iterator().next(), RestMethod.GET);
    try {
        AMBRY_QUOTA_ENFORCER.charge(restRequest, readRequestCostMap);
        AMBRY_QUOTA_ENFORCER.recommend(restRequest);
        fail("if quota is not found we should see exception");
    } catch (QuotaException quotaException) {
        Assert.assertTrue(quotaException.isRetryable());
    }
    // 3. Test that retryable quota exception is thrown in case of any error.
    restRequest = QuotaTestUtils.createRestRequest(ACCOUNT, container, RestMethod.GET);
    QUOTA_SOURCE.throwException = true;
    try {
        AMBRY_QUOTA_ENFORCER.charge(restRequest, readRequestCostMap);
        fail("QuotaException should be thrown in case of any error.");
    } catch (QuotaException quotaException) {
        Assert.assertTrue(quotaException.isRetryable());
    }
    QUOTA_SOURCE.throwException = false;
    // 4. Test that usage is updated and recommendation is deny when usage >= quota.
    AMBRY_QUOTA_ENFORCER.charge(restRequest, readRequestCostMap);
    quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(restRequest);
    assertEquals(QuotaName.READ_CAPACITY_UNIT, quotaRecommendation.getQuotaName());
    assertEquals(180, quotaRecommendation.getQuotaUsagePercentage(), 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.DELAY);
    // make sure that correct quota is charged.
    assertEquals(9, usageMap.get(String.valueOf(ACCOUNT.getId())).getWcu());
    assertEquals(18, // make sure that correct quota is charged.
    usageMap.get(String.valueOf(ACCOUNT.getId())).getRcu());
    restRequest = QuotaTestUtils.createRestRequest(ACCOUNT, container, RestMethod.POST);
    AMBRY_QUOTA_ENFORCER.charge(restRequest, writeRequestCostMap);
    quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(restRequest);
    assertEquals(QuotaName.WRITE_CAPACITY_UNIT, quotaRecommendation.getQuotaName());
    assertEquals(180, quotaRecommendation.getQuotaUsagePercentage(), 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.DELAY);
    assertEquals(18, // make sure that correct quota is charged.
    usageMap.get(String.valueOf(ACCOUNT.getId())).getWcu());
    assertEquals(18, // make sure that correct quota is charged.
    usageMap.get(String.valueOf(ACCOUNT.getId())).getRcu());
    restRequest = QuotaTestUtils.createRestRequest(ACCOUNT, container, RestMethod.DELETE);
    AMBRY_QUOTA_ENFORCER.charge(restRequest, writeRequestCostMap);
    quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(restRequest);
    assertEquals(QuotaName.WRITE_CAPACITY_UNIT, quotaRecommendation.getQuotaName());
    assertEquals(270, quotaRecommendation.getQuotaUsagePercentage(), 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.DELAY);
    assertEquals(27, // make sure that correct quota is charged.
    usageMap.get(String.valueOf(ACCOUNT.getId())).getWcu());
    assertEquals(18, // make sure that correct quota is charged.
    usageMap.get(String.valueOf(ACCOUNT.getId())).getRcu());
    restRequest = QuotaTestUtils.createRestRequest(ACCOUNT, container, RestMethod.PUT);
    AMBRY_QUOTA_ENFORCER.charge(restRequest, writeRequestCostMap);
    quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(restRequest);
    assertEquals(QuotaName.WRITE_CAPACITY_UNIT, quotaRecommendation.getQuotaName());
    assertEquals(360, quotaRecommendation.getQuotaUsagePercentage(), 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.DELAY);
    assertEquals(36, // make sure that correct quota is charged.
    usageMap.get(String.valueOf(ACCOUNT.getId())).getWcu());
    assertEquals(18, // make sure that correct quota is charged.
    usageMap.get(String.valueOf(ACCOUNT.getId())).getRcu());
}

Example 2 with QuotaName

use of com.github.ambry.quota.QuotaName in project ambry by linkedin.

the class AmbryCUQuotaEnforcer method charge.

@Override
public QuotaRecommendation charge(RestRequest restRequest, Map<QuotaName, Double> requestCostMap) throws QuotaException {
    final QuotaName quotaName = QuotaUtils.getCUQuotaName(restRequest);
    if (requestCostMap.isEmpty() || !requestCostMap.containsKey(quotaName)) {
        String errorMessage = String.format("No %s cost provided for request %s. Nothing to charge", quotaName.name(), RestUtils.convertToStr(restRequest));
        throw new QuotaException(errorMessage, true);
    }
    final QuotaResource quotaResource = QuotaResource.fromRestRequest(restRequest);
    doAndHandleException(() -> {
        quotaSource.chargeUsage(quotaResource, quotaName, requestCostMap.get(quotaName));
        return null;
    }, String.format("Could not charge for request %s due to", restRequest));
    return recommend(restRequest);
}

Example 3 with QuotaName

use of com.github.ambry.quota.QuotaName in project ambry by linkedin.

the class AmbryCUQuotaEnforcer method recommend.

@Override
public QuotaRecommendation recommend(RestRequest restRequest) throws QuotaException {
    final QuotaName quotaName = QuotaUtils.getCUQuotaName(restRequest);
    float usage = doAndHandleException(() -> {
        final QuotaResource quotaResource = QuotaResource.fromRestRequest(restRequest);
        return quotaSource.getUsage(quotaResource, quotaName);
    }, String.format("Could not recommend for request %s due to", restRequest));
    return buildQuotaRecommendation(usage, quotaName);
}