Examples with Container com.github.ambry.account.Container used on opensource projects

Example 31 with Container

use of com.github.ambry.account.Container in project ambry by linkedin.

the class AmbryCUQuotaSourceTest method testChargeForValidResource.

@Test
public void testChargeForValidResource() throws Exception {
    Account account = inMemAccountService.getAccountById((short) 102);
    Container container = new ArrayList<>(account.getAllContainers()).get(0);
    ambryCUQuotaSource.chargeUsage(QuotaResource.fromContainer(container), QuotaName.READ_CAPACITY_UNIT, 102400000);
    Assert.assertEquals(10.0, ambryCUQuotaSource.getUsage(QuotaResource.fromContainer(container), QuotaName.READ_CAPACITY_UNIT), 0.01);
    Assert.assertEquals(0, ambryCUQuotaSource.getUsage(QuotaResource.fromContainer(container), QuotaName.WRITE_CAPACITY_UNIT), 0.01);
    ambryCUQuotaSource.chargeUsage(QuotaResource.fromContainer(container), QuotaName.WRITE_CAPACITY_UNIT, 102400000);
    Assert.assertEquals(10.0, ambryCUQuotaSource.getUsage(QuotaResource.fromContainer(container), QuotaName.READ_CAPACITY_UNIT), 0.01);
    Assert.assertEquals(10.0, ambryCUQuotaSource.getUsage(QuotaResource.fromContainer(container), QuotaName.WRITE_CAPACITY_UNIT), 0.01);
}

Example 32 with Container

use of com.github.ambry.account.Container in project ambry by linkedin.

the class AmbryCUQuotaSourceTest method createAccountForQuota.

/**
 * Create {@link Account} object with specified quota and accountId.
 * @param mapOrQuota quota of the account.
 * @param accountId id of the account.
 * @return Account object.
 */
private static Account createAccountForQuota(JsonCUQuotaDataProviderUtil.MapOrQuota mapOrQuota, String accountId) {
    AccountBuilder accountBuilder = new AccountBuilder();
    accountBuilder.id(Short.parseShort(accountId));
    accountBuilder.name(accountId);
    List<Container> containers = new ArrayList<>();
    if (!mapOrQuota.isQuota()) {
        for (String containerId : mapOrQuota.getContainerQuotas().keySet()) {
            containers.add(createContainer(containerId));
        }
    }
    accountBuilder.containers(containers);
    accountBuilder.status(Account.AccountStatus.ACTIVE);
    if (mapOrQuota.isQuota()) {
        accountBuilder.quotaResourceType(QuotaResourceType.ACCOUNT);
    } else {
        accountBuilder.quotaResourceType(QuotaResourceType.CONTAINER);
    }
    return accountBuilder.build();
}

Example 33 with Container

use of com.github.ambry.account.Container in project ambry by linkedin.

the class AmbryCUQuotaSourceTest method testChargeForNonExistentResource.

@Test
public void testChargeForNonExistentResource() throws Exception {
    Account account = createAccountForQuota(new JsonCUQuotaDataProviderUtil.MapOrQuota(new CapacityUnit(10, 10)), "106");
    Container container = createContainer("1");
    verifyFailsWithQuotaException(() -> ambryCUQuotaSource.getUsage(QuotaResource.fromAccount(account), QuotaName.READ_CAPACITY_UNIT), "Get usage for non existent resource");
    verifyFailsWithQuotaException(() -> ambryCUQuotaSource.getUsage(QuotaResource.fromAccount(account), QuotaName.WRITE_CAPACITY_UNIT), "Get usage for non existent resource");
    verifyFailsWithQuotaException(() -> {
        ambryCUQuotaSource.chargeUsage(QuotaResource.fromAccount(account), QuotaName.READ_CAPACITY_UNIT, 10.0);
        return null;
    }, "Charge usage for non existent resource");
    verifyFailsWithQuotaException(() -> ambryCUQuotaSource.getUsage(QuotaResource.fromAccount(account), QuotaName.READ_CAPACITY_UNIT), "Get usage for non existent resource");
    verifyFailsWithQuotaException(() -> ambryCUQuotaSource.getUsage(QuotaResource.fromAccount(account), QuotaName.WRITE_CAPACITY_UNIT), "Get usage for non existent resource");
    verifyFailsWithQuotaException(() -> {
        ambryCUQuotaSource.chargeUsage(QuotaResource.fromContainer(container), QuotaName.READ_CAPACITY_UNIT, 10.0);
        return null;
    }, "Charge usage for non existent resource");
    verifyFailsWithQuotaException(() -> ambryCUQuotaSource.getUsage(QuotaResource.fromContainer(container), QuotaName.READ_CAPACITY_UNIT), "Get usage for non existent resource");
    verifyFailsWithQuotaException(() -> ambryCUQuotaSource.getUsage(QuotaResource.fromContainer(container), QuotaName.WRITE_CAPACITY_UNIT), "Get usage for non existent resource");
}

Example 34 with Container

use of com.github.ambry.account.Container 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 35 with Container

use of com.github.ambry.account.Container in project ambry by linkedin.

the class JsonCUQuotaDataProviderUtil method getCUQuotasFromJson.

/**
 * Util to parse the provided json string into resource quotas and create a {@link Map} of QuotaResource id and
 * {@link CapacityUnit} from json configs.
 * @param resourceCUQuotaInJson Json string containing resource quota values.
 * @param accountService The {@link AccountService} to use.
 * @throws IOException in case of any exception.
 */
public static Map<String, CapacityUnit> getCUQuotasFromJson(String resourceCUQuotaInJson, AccountService accountService) throws IOException {
    Map<String, CapacityUnit> quota = new HashMap<>();
    ObjectMapper objectMapper = new ObjectMapper();
    if (resourceCUQuotaInJson != null && !resourceCUQuotaInJson.trim().isEmpty()) {
        Map<String, MapOrQuota> tempQuotas = objectMapper.readValue(resourceCUQuotaInJson, new TypeReference<Map<String, MapOrQuota>>() {
        });
        for (Map.Entry<String, MapOrQuota> entry : tempQuotas.entrySet()) {
            final Account account = accountService.getAccountById(Short.parseShort(entry.getKey()));
            if (account == null) {
                throw new IllegalStateException("No account id " + entry.getKey() + " is found in the account service");
            }
            if (account.getQuotaResourceType() == QuotaResourceType.ACCOUNT && !entry.getValue().isQuota || account.getQuotaResourceType() != QuotaResourceType.ACCOUNT && entry.getValue().isQuota) {
                throw new IllegalStateException("Account " + entry.getKey() + " quota enforcement type is different from account metadata: " + account.getQuotaResourceType());
            }
            if (entry.getValue().isQuota()) {
                quota.put(QuotaResource.fromAccount(account).getResourceId(), entry.getValue().getQuota());
            } else {
                for (Map.Entry<String, CapacityUnit> containerQuotaEntry : entry.getValue().getContainerQuotas().entrySet()) {
                    String containerIdStr = containerQuotaEntry.getKey();
                    Container container = account.getContainerById(Short.parseShort(containerIdStr));
                    if (container == null) {
                        throw new IllegalStateException("No container id " + containerIdStr + " is found in the account service under account " + entry.getKey());
                    }
                    quota.put(QuotaResource.fromContainer(container).getResourceId(), containerQuotaEntry.getValue());
                }
            }
        }
    }
    return quota;
}