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

Example 16 with ContainerBuilder

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

the class StorageQuotaEnforcerTest method testGetQuotaAndUsageAndCharge.

/**
 * Test {@link StorageQuotaEnforcer#getQuotaAndUsage} and {@link StorageQuotaEnforcer#charge} methods.
 * @throws Exception
 */
@Test
public void testGetQuotaAndUsageAndCharge() throws Exception {
    int initNumAccounts = 10;
    Map<String, Map<String, Long>> containerUsage = TestUtils.makeStorageMap(initNumAccounts, 10, 10000, 1000);
    InMemAccountService accountService = new InMemAccountService(false, false);
    Map<String, Long> storageQuota = new HashMap<>();
    // Account and container id's base is 1, not 0
    for (int i = 1; i <= initNumAccounts; i++) {
        QuotaResourceType resourceType = i <= containerUsage.size() / 2 ? QuotaResourceType.CONTAINER : QuotaResourceType.ACCOUNT;
        AccountBuilder accountBuilder = new AccountBuilder((short) i, String.valueOf(i), Account.AccountStatus.ACTIVE, resourceType);
        for (int j = 1; j <= 10; j++) {
            accountBuilder.addOrUpdateContainer(new ContainerBuilder((short) j, String.valueOf(j), Container.ContainerStatus.ACTIVE, "", (short) i).build());
        }
        accountService.updateAccounts(Collections.singleton(accountBuilder.build()));
        if (resourceType == QuotaResourceType.ACCOUNT) {
            storageQuota.put(QuotaResource.fromAccount(accountService.getAccountById((short) i)).getResourceId(), containerUsage.get(String.valueOf(i)).values().stream().mapToLong(Long::longValue).sum());
        } else {
            accountService.getAccountById((short) i).getAllContainers().forEach(c -> storageQuota.put(QuotaResource.fromContainer(c).getResourceId(), containerUsage.get(String.valueOf(c.getParentAccountId())).get(String.valueOf(c.getId()))));
        }
    }
    JSONStringStorageQuotaSource quotaSource = new JSONStringStorageQuotaSource(storageQuota, accountService);
    StorageQuotaEnforcer enforcer = new StorageQuotaEnforcer(config, quotaSource, (StorageUsageRefresher) null);
    enforcer.initStorageUsage(Collections.EMPTY_MAP);
    for (Map.Entry<String, Map<String, Long>> accountEntry : containerUsage.entrySet()) {
        short accountId = Short.valueOf(accountEntry.getKey());
        if (accountService.getAccountById(accountId).getQuotaResourceType() == QuotaResourceType.ACCOUNT) {
            long quota = (long) (quotaSource.getQuota(QuotaResource.fromAccount(accountService.getAccountById(accountId)), QuotaName.STORAGE_IN_GB).getQuotaValue()) * BYTES_IN_GB;
            RestRequest restRequest = createRestRequest(accountService, accountId, (short) 1);
            Pair<Long, Long> quotaAndUsage = enforcer.getQuotaAndUsage(restRequest);
            assertEquals("Account id: " + accountEntry.getKey(), quota, quotaAndUsage.getFirst().longValue());
            assertEquals(0L, quotaAndUsage.getSecond().longValue());
            quotaAndUsage = enforcer.charge(restRequest, quota / 2);
            assertEquals(quota, quotaAndUsage.getFirst().longValue());
            assertEquals(quota / 2, quotaAndUsage.getSecond().longValue());
            quotaAndUsage = enforcer.charge(restRequest, quota);
            assertEquals(quota, quotaAndUsage.getFirst().longValue());
            assertEquals(quota / 2 + quota, quotaAndUsage.getSecond().longValue());
        } else {
            for (Map.Entry<String, Long> containerEntry : accountEntry.getValue().entrySet()) {
                short containerId = Short.valueOf(containerEntry.getKey());
                long quota = containerEntry.getValue() * BYTES_IN_GB;
                RestRequest restRequest = createRestRequest(accountService, accountId, containerId);
                Pair<Long, Long> quotaAndUsage = enforcer.getQuotaAndUsage(restRequest);
                assertEquals(quota, quotaAndUsage.getFirst().longValue());
                assertEquals(0L, quotaAndUsage.getSecond().longValue());
                quotaAndUsage = enforcer.charge(restRequest, quota / 2);
                assertEquals(quota, quotaAndUsage.getFirst().longValue());
                assertEquals(quota / 2, quotaAndUsage.getSecond().longValue());
                quotaAndUsage = enforcer.charge(restRequest, quota);
                assertEquals(quota, quotaAndUsage.getFirst().longValue());
                assertEquals(quota / 2 + quota, quotaAndUsage.getSecond().longValue());
            }
        }
    }
    // Now create a restRequest that doesn't carry account and container
    RestRequest restRequest = createRestRequest();
    Pair<Long, Long> quotaAndUsage = enforcer.getQuotaAndUsage(restRequest);
    assertEquals(-1L, quotaAndUsage.getFirst().longValue());
    assertEquals(0L, quotaAndUsage.getSecond().longValue());
    quotaAndUsage = enforcer.charge(restRequest, 100L);
    assertEquals(-1L, quotaAndUsage.getFirst().longValue());
    assertEquals(0L, quotaAndUsage.getSecond().longValue());
    Account account = new AccountBuilder((short) 1000, String.valueOf(1000), Account.AccountStatus.ACTIVE, QuotaResourceType.CONTAINER).addOrUpdateContainer(new ContainerBuilder((short) 10000, String.valueOf(10000), Container.ContainerStatus.ACTIVE, "", (short) 1000).build()).build();
    accountService.updateAccounts(Collections.singleton(account));
    restRequest = createRestRequest(accountService, (short) 1000, (short) 10000);
    quotaAndUsage = enforcer.getQuotaAndUsage(restRequest);
    assertEquals(-1L, quotaAndUsage.getFirst().longValue());
    assertEquals(0L, quotaAndUsage.getSecond().longValue());
    quotaAndUsage = enforcer.charge(restRequest, 100L);
    assertEquals(-1L, quotaAndUsage.getFirst().longValue());
    assertEquals(0L, quotaAndUsage.getSecond().longValue());
}

Example 17 with ContainerBuilder

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

the class AmbryCUQuotaSourceTest method createContainer.

/**
 * Create a {@link Container} with the specified containerId.
 * @param containerId id of the container.
 * @return Container object.
 */
private static Container createContainer(String containerId) {
    ContainerBuilder containerBuilder = new ContainerBuilder();
    containerBuilder.setId(Short.parseShort(containerId));
    containerBuilder.setName(containerId);
    containerBuilder.setStatus(Container.ContainerStatus.ACTIVE);
    return containerBuilder.build();
}

Example 18 with ContainerBuilder

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

the class JSONStringStorageQuotaSourceTest method testJSONStringStorageQuotaSource.

@Test
public void testJSONStringStorageQuotaSource() throws Exception {
    // Trick to create a string literal without escape.
    String json = "{`10`: {`1`: 1000, `2`: 3000}, `20`: {`4`: 2000, `5`: 1000}, `30`: 4000, `40`: 5000, `50`: {`6`: 6000}}".replace("`", "\"");
    Properties properties = new Properties();
    properties.setProperty(StorageQuotaConfig.STORAGE_QUOTA_IN_JSON, json);
    StorageQuotaConfig config = new StorageQuotaConfig(new VerifiableProperties(properties));
    // Setting up accounts and account service
    InMemAccountService accountService = new InMemAccountService(false, false);
    Account account = new AccountBuilder((short) 10, "10", Account.AccountStatus.ACTIVE, QuotaResourceType.CONTAINER).addOrUpdateContainer(new ContainerBuilder((short) 1, "1", Container.ContainerStatus.ACTIVE, "", (short) 10).build()).addOrUpdateContainer(new ContainerBuilder((short) 2, "2", Container.ContainerStatus.ACTIVE, "", (short) 10).build()).build();
    accountService.updateAccounts(Collections.singleton(account));
    account = new AccountBuilder((short) 20, "20", Account.AccountStatus.ACTIVE, QuotaResourceType.CONTAINER).addOrUpdateContainer(new ContainerBuilder((short) 4, "4", Container.ContainerStatus.ACTIVE, "", (short) 20).build()).addOrUpdateContainer(new ContainerBuilder((short) 5, "5", Container.ContainerStatus.ACTIVE, "", (short) 20).build()).build();
    accountService.updateAccounts(Collections.singleton(account));
    account = new AccountBuilder((short) 30, "30", Account.AccountStatus.ACTIVE, QuotaResourceType.ACCOUNT).addOrUpdateContainer(new ContainerBuilder((short) 4, "4", Container.ContainerStatus.ACTIVE, "", (short) 30).build()).build();
    accountService.updateAccounts(Collections.singleton(account));
    account = new AccountBuilder((short) 40, "40", Account.AccountStatus.ACTIVE, QuotaResourceType.ACCOUNT).addOrUpdateContainer(new ContainerBuilder((short) 4, "4", Container.ContainerStatus.ACTIVE, "", (short) 40).build()).build();
    accountService.updateAccounts(Collections.singleton(account));
    account = new AccountBuilder((short) 50, "50", Account.AccountStatus.ACTIVE, QuotaResourceType.CONTAINER).addOrUpdateContainer(new ContainerBuilder((short) 6, "6", Container.ContainerStatus.ACTIVE, "", (short) 50).build()).build();
    accountService.updateAccounts(Collections.singleton(account));
    JSONStringStorageQuotaSource source = new JSONStringStorageQuotaSource(config, accountService);
    QuotaResourceType resourceType = QuotaResourceType.CONTAINER;
    Quota quota = source.getQuota(new QuotaResource("1000_1", resourceType), QuotaName.STORAGE_IN_GB);
    assertNull(quota);
    quota = source.getQuota(new QuotaResource("10_1", resourceType), QuotaName.STORAGE_IN_GB);
    assertEquals(1000L, (long) quota.getQuotaValue());
    quota = source.getQuota(new QuotaResource("10_2", resourceType), QuotaName.STORAGE_IN_GB);
    assertEquals(3000L, (long) quota.getQuotaValue());
    quota = source.getQuota(new QuotaResource("20_4", resourceType), QuotaName.STORAGE_IN_GB);
    assertEquals(2000L, (long) quota.getQuotaValue());
    quota = source.getQuota(new QuotaResource("20_5", resourceType), QuotaName.STORAGE_IN_GB);
    assertEquals(1000L, (long) quota.getQuotaValue());
    quota = source.getQuota(new QuotaResource("30", QuotaResourceType.ACCOUNT), QuotaName.STORAGE_IN_GB);
    assertEquals(4000L, (long) quota.getQuotaValue());
    quota = source.getQuota(new QuotaResource("40", QuotaResourceType.ACCOUNT), QuotaName.STORAGE_IN_GB);
    assertEquals(5000L, (long) quota.getQuotaValue());
    quota = source.getQuota(new QuotaResource("50_6", resourceType), QuotaName.STORAGE_IN_GB);
    assertEquals(6000L, (long) quota.getQuotaValue());
}

Example 19 with ContainerBuilder

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

the class AccountTestUtils method generateContainerBuilders.

/**
 * Generate {@link ContainerBuilder}s for specified {@code accountId}.
 * @param numContainers number of {@link ContainerBuilder}s to generate.
 * @param accountId accountId for container.
 * @return {@link List} of {@link ContainerBuilder}s.
 */
public static List<ContainerBuilder> generateContainerBuilders(int numContainers, short accountId) {
    List<ContainerBuilder> containerBuilders = new ArrayList<>();
    Set<Short> containerIdSet = new HashSet<>();
    for (int j = 0; j < numContainers; j++) {
        short containerId = Utils.getRandomShort(random);
        if (!containerIdSet.add(containerId)) {
            j--;
            continue;
        }
        String containerName = UUID.randomUUID().toString();
        Container.ContainerStatus containerStatus = random.nextBoolean() ? Container.ContainerStatus.ACTIVE : Container.ContainerStatus.INACTIVE;
        String containerDescription = UUID.randomUUID().toString();
        boolean containerCaching = random.nextBoolean();
        boolean containerEncryption = random.nextBoolean();
        boolean containerPreviousEncryption = containerEncryption || random.nextBoolean();
        boolean mediaScanDisabled = random.nextBoolean();
        String replicationPolicy = TestUtils.getRandomString(10);
        boolean ttlRequired = random.nextBoolean();
        ContainerBuilder containerBuilder = new ContainerBuilder(containerId, containerName, containerStatus, containerDescription, accountId).setEncrypted(containerEncryption).setPreviouslyEncrypted(containerPreviousEncryption).setCacheable(containerCaching).setMediaScanDisabled(mediaScanDisabled).setReplicationPolicy(replicationPolicy).setTtlRequired(ttlRequired);
        containerBuilders.add(containerBuilder);
    }
    return containerBuilders;
}