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

Example 1 with QuotaRecommendation

use of com.github.ambry.quota.QuotaRecommendation 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 QuotaRecommendation

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

the class AmbryCUQuotaEnforcerTest method testRecommend.

@Test
public void testRecommend() throws Exception {
    // 1. Test that recommendation is serve when usage is within limit and correct quota is used based on rest method.
    QuotaRecommendation quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(QuotaTestUtils.createRestRequest(ACCOUNT, ACCOUNT.getAllContainers().iterator().next(), RestMethod.GET));
    assertEquals(QuotaName.READ_CAPACITY_UNIT, quotaRecommendation.getQuotaName());
    assertEquals(0, quotaRecommendation.getQuotaUsagePercentage(), 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.ALLOW);
    quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(QuotaTestUtils.createRestRequest(ACCOUNT, ACCOUNT.getAllContainers().iterator().next(), RestMethod.POST));
    assertEquals(QuotaName.WRITE_CAPACITY_UNIT, quotaRecommendation.getQuotaName());
    assertEquals(0, quotaRecommendation.getQuotaUsagePercentage(), 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.ALLOW);
    quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(QuotaTestUtils.createRestRequest(ACCOUNT, ACCOUNT.getAllContainers().iterator().next(), RestMethod.PUT));
    assertEquals(QuotaName.WRITE_CAPACITY_UNIT, quotaRecommendation.getQuotaName());
    assertEquals(0, quotaRecommendation.getQuotaUsagePercentage(), 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.ALLOW);
    quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(QuotaTestUtils.createRestRequest(ACCOUNT, ACCOUNT.getAllContainers().iterator().next(), RestMethod.DELETE));
    assertEquals(QuotaName.WRITE_CAPACITY_UNIT, quotaRecommendation.getQuotaName());
    assertEquals(0, quotaRecommendation.getQuotaUsagePercentage(), 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.ALLOW);
    // 2. Test that retryable QuotaException is thrown when quota not found.
    Account newAccount = ACCOUNT_SERVICE.generateRandomAccount(QuotaResourceType.ACCOUNT);
    try {
        AMBRY_QUOTA_ENFORCER.recommend(QuotaTestUtils.createRestRequest(newAccount, newAccount.getAllContainers().iterator().next(), RestMethod.GET));
        fail("QuotaException should be thrown when quota not found");
    } catch (QuotaException quotaException) {
        Assert.assertTrue(quotaException.isRetryable());
    }
    // 3. Test that recommendation is serve in case of any error.
    QUOTA_SOURCE.throwException = true;
    try {
        AMBRY_QUOTA_ENFORCER.recommend(QuotaTestUtils.createRestRequest(ACCOUNT, ACCOUNT.getAllContainers().iterator().next(), RestMethod.GET));
        fail("QuotaException should be thrown.");
    } catch (QuotaException quotaException) {
        Assert.assertTrue(quotaException.isRetryable());
    }
    QUOTA_SOURCE.throwException = false;
    // 4. Test that recommendation is deny when usage >= quota.
    Map<String, CapacityUnit> usageMap = QUOTA_SOURCE.getAllQuotaUsage();
    Map<String, CapacityUnit> quotaMap = QUOTA_SOURCE.getAllQuota();
    String id = String.valueOf(ACCOUNT.getId());
    usageMap.put(id, new CapacityUnit(quotaMap.get(id).getRcu() + 1, quotaMap.get(id).getWcu()));
    float usagePercentage = (float) (usageMap.get(id).getRcu() * 100) / quotaMap.get(id).getRcu();
    quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(QuotaTestUtils.createRestRequest(ACCOUNT, ACCOUNT.getAllContainers().iterator().next(), RestMethod.GET));
    assertEquals(QuotaName.READ_CAPACITY_UNIT, quotaRecommendation.getQuotaName());
    assertEquals(usagePercentage, quotaRecommendation.getQuotaUsagePercentage(), 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.DELAY);
}

Example 3 with QuotaRecommendation

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

the class StorageQuotaEnforcer method recommendBasedOnQuotaAndUsage.

/**
 * Return a {@link QuotaRecommendation} based on the given pair of quota and current usage.
 * @param pair The {@link Pair} of quota and current usage.
 * @return A {@link QuotaRecommendation}.
 */
private QuotaRecommendation recommendBasedOnQuotaAndUsage(Pair<Long, Long> pair) {
    long quotaValue = pair.getFirst();
    long currentUsage = pair.getSecond();
    if (quotaValue == -1L) {
        // There is no quota set for the given account/container
        return NO_QUOTA_VALUE_RECOMMENDATION;
    }
    QuotaAction quotaAction = (currentUsage >= quotaValue) ? QuotaAction.REJECT : QuotaAction.ALLOW;
    float usagePercentage = currentUsage >= quotaValue ? 100f : ((float) currentUsage) / quotaValue * 100f;
    return new QuotaRecommendation(quotaAction, usagePercentage, QuotaName.STORAGE_IN_GB, NO_RETRY);
}