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

Example 1 with QuotaResourceType

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

the class StorageQuotaEnforcerTest method testStorageUsageRefresherListener.

/**
 * Test when storage usage updates.
 */
@Test
public void testStorageUsageRefresherListener() throws Exception {
    int initNumAccounts = 10;
    Map<String, Map<String, Long>> containerUsage = TestUtils.makeStorageMap(initNumAccounts, 10, 10000, 1000);
    InMemAccountService accountService = new InMemAccountService(false, false);
    Map<QuotaResource, Long> expectedStorageUsages = 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) {
            expectedStorageUsages.put(QuotaResource.fromAccountId((short) i), containerUsage.get(String.valueOf(i)).values().stream().mapToLong(Long::longValue).sum());
        } else {
            for (Map.Entry<String, Long> containerEntry : containerUsage.get(String.valueOf(i)).entrySet()) {
                expectedStorageUsages.put(QuotaResource.fromContainerId((short) i, Short.valueOf(containerEntry.getKey())), containerEntry.getValue());
            }
        }
    }
    StorageQuotaEnforcer enforcer = new StorageQuotaEnforcer(config, new JSONStringStorageQuotaSource(new HashMap<>(), accountService), (StorageUsageRefresher) null);
    enforcer.initStorageUsage(containerUsage);
    StorageUsageRefresher.Listener listener = enforcer.getUsageRefresherListener();
    int numUpdates = 10;
    for (int i = 1; i <= numUpdates; i++) {
        if (i % 2 == 0) {
            // add new storage usage
            Map<String, Map<String, Long>> additionalUsage = TestUtils.makeStorageMap(1, 10, 10000, 1000);
            short accountId = (short) (initNumAccounts + i);
            containerUsage.put(String.valueOf(accountId), additionalUsage.remove("1"));
            QuotaResourceType resourceType = i % 4 == 0 ? QuotaResourceType.CONTAINER : QuotaResourceType.ACCOUNT;
            AccountBuilder accountBuilder = new AccountBuilder(accountId, String.valueOf(accountId), Account.AccountStatus.ACTIVE, resourceType);
            for (int j = 1; j <= 10; j++) {
                accountBuilder.addOrUpdateContainer(new ContainerBuilder((short) j, String.valueOf(j), Container.ContainerStatus.ACTIVE, "", (short) accountId).build());
            }
            accountService.updateAccounts(Collections.singleton(accountBuilder.build()));
            if (resourceType == QuotaResourceType.ACCOUNT) {
                expectedStorageUsages.put(QuotaResource.fromAccountId(accountId), containerUsage.get(String.valueOf(accountId)).values().stream().mapToLong(Long::longValue).sum());
            } else {
                for (Map.Entry<String, Long> containerEntry : containerUsage.get(String.valueOf(accountId)).entrySet()) {
                    expectedStorageUsages.put(QuotaResource.fromContainerId(accountId, Short.valueOf(containerEntry.getKey())), containerEntry.getValue());
                }
            }
        } else {
            // change existing storage usage
            Random random = new Random();
            int accountId = random.nextInt(initNumAccounts) + 1;
            int containerId = random.nextInt(10) + 1;
            long newValue = random.nextLong();
            long oldValue = containerUsage.get(String.valueOf(accountId)).get(String.valueOf(containerId));
            containerUsage.get(String.valueOf(accountId)).put(String.valueOf(containerId), newValue);
            if (accountService.getAccountById((short) accountId).getQuotaResourceType() == QuotaResourceType.ACCOUNT) {
                QuotaResource resource = QuotaResource.fromAccountId((short) accountId);
                expectedStorageUsages.put(resource, expectedStorageUsages.get(resource) - oldValue + newValue);
            } else {
                expectedStorageUsages.put(QuotaResource.fromContainerId((short) accountId, (short) containerId), newValue);
            }
        }
        listener.onNewContainerStorageUsage(containerUsage);
        assertEquals(expectedStorageUsages, enforcer.getStorageUsages());
    }
}

Example 2 with QuotaResourceType

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

the class AccountContainerTest method accountFromJson.

/**
 * Deserialize a {@link JSONObject} to an {@link Account}.
 * @param metadata The {@link JSONObject}.
 * @return A {@link Account}.
 * @throws JSONException
 */
private Account accountFromJson(JSONObject metadata) throws JSONException {
    if (metadata == null) {
        throw new IllegalArgumentException("metadata cannot be null.");
    }
    short metadataVersion = (short) metadata.getInt(Account.JSON_VERSION_KEY);
    switch(metadataVersion) {
        case Account.JSON_VERSION_1:
            short id = (short) metadata.getInt(ACCOUNT_ID_KEY);
            String name = metadata.getString(ACCOUNT_NAME_KEY);
            AccountStatus status = AccountStatus.valueOf(metadata.getString(Account.STATUS_KEY));
            int snapshotVersion = metadata.optInt(Account.SNAPSHOT_VERSION_KEY, Account.SNAPSHOT_VERSION_DEFAULT_VALUE);
            long lastModifiedTime = metadata.optLong(Account.LAST_MODIFIED_TIME_KEY, Account.LAST_MODIFIED_TIME_DEFAULT_VALUE);
            boolean aclInheritedByContainer = metadata.optBoolean(ACL_INHERITED_BY_CONTAINER_KEY, ACL_INHERITED_BY_CONTAINER_DEFAULT_VALUE);
            if (name == null || status == null) {
                throw new IllegalStateException("Either of required fields name=" + name + " or status=" + status + " is null");
            }
            JSONArray containerArray = metadata.optJSONArray(CONTAINERS_KEY);
            List<Container> containers = new ArrayList<>();
            if (containerArray != null) {
                for (int index = 0; index < containerArray.length(); index++) {
                    containers.add(containerFromJson(containerArray.getJSONObject(index), id));
                }
            }
            QuotaResourceType quotaResourceType = metadata.optEnum(QuotaResourceType.class, QUOTA_RESOURCE_TYPE_KEY, QUOTA_RESOURCE_TYPE_DEFAULT_VALUE);
            return new AccountBuilder(id, name, status, quotaResourceType).snapshotVersion(snapshotVersion).lastModifiedTime(lastModifiedTime).aclInheritedByContainer(aclInheritedByContainer).containers(containers).build();
        default:
            throw new IllegalStateException("Unsupported account json version=" + metadataVersion);
    }
}

Example 3 with QuotaResourceType

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

the class StorageQuotaEnforcerTest method testInitStorageUsage.

/**
 * Test to initialize the storage usage with non-empty map.
 */
@Test
public void testInitStorageUsage() throws Exception {
    Map<String, Map<String, Long>> containerUsage = TestUtils.makeStorageMap(10, 10, 1000, 100);
    InMemAccountService accountService = new InMemAccountService(false, false);
    Map<QuotaResource, Long> expectedStorageUsages = new HashMap<>();
    // Account and container id's base is 1, not 0
    for (int i = 1; i <= 10; 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) {
            expectedStorageUsages.put(QuotaResource.fromAccountId((short) i), containerUsage.get(String.valueOf(i)).values().stream().mapToLong(Long::longValue).sum());
        } else {
            for (Map.Entry<String, Long> containerEntry : containerUsage.get(String.valueOf(i)).entrySet()) {
                expectedStorageUsages.put(QuotaResource.fromContainerId((short) i, Short.valueOf(containerEntry.getKey())), containerEntry.getValue());
            }
        }
    }
    StorageQuotaEnforcer enforcer = new StorageQuotaEnforcer(config, new JSONStringStorageQuotaSource(new HashMap<>(), accountService), (StorageUsageRefresher) null);
    enforcer.initStorageUsage(containerUsage);
    assertEquals(expectedStorageUsages, enforcer.getStorageUsages());
}

Example 4 with QuotaResourceType

use of com.github.ambry.quota.QuotaResourceType 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 5 with QuotaResourceType

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

the class JSONStringStorageQuotaSource method getQuota.

@Override
public Quota getQuota(QuotaResource quotaResource, QuotaName quotaName) {
    if (quotaName != QuotaName.STORAGE_IN_GB) {
        return null;
    }
    QuotaResourceType resourceType = quotaResource.getQuotaResourceType();
    if (resourceType != QuotaResourceType.ACCOUNT && resourceType != QuotaResourceType.CONTAINER) {
        throw new IllegalArgumentException("Unsupported quota resource type: " + resourceType);
    }
    String resourceId = quotaResource.getResourceId();
    if (storageQuota.containsKey(resourceId)) {
        long quotaValue = storageQuota.get(resourceId);
        return new Quota(quotaName, quotaValue, quotaResource);
    }
    return null;
}