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

Example 11 with QuotaResource

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

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

the class StorageQuotaEnforcer method getQuotaAndUsage.

/**
 * Return quota and current usage for the account/container carried in the given {@code restRequest}.
 * If there is no account and container found in the {@code restRequest}, this method would return -1
 * for quota. If there is no quota found for the account/container, this method would return -1 for
 * quota as well.
 * @param restRequest the {@link RestRequest} that carries account and container in the header.
 * @return A {@link Pair} whose first element is quota the second element is current storage usage.
 */
Pair<Long, Long> getQuotaAndUsage(RestRequest restRequest) {
    long quotaValue = -1L;
    long currentUsage = 0L;
    try {
        Account account = RestUtils.getAccountFromArgs(restRequest.getArgs());
        Container container = RestUtils.getContainerFromArgs(restRequest.getArgs());
        QuotaResource quotaResource = account.getQuotaResourceType() == QuotaResourceType.ACCOUNT ? QuotaResource.fromAccount(account) : QuotaResource.fromContainer(container);
        quotaValue = getQuotaValueForResource(quotaResource);
        if (quotaValue != -1L) {
            currentUsage = storageUsages.getOrDefault(quotaResource, 0L);
        }
    } catch (Exception e) {
        logger.error("Failed to getQuotaAndUsage for RestRequest {}", restRequest, e);
    }
    return new Pair<>(quotaValue, currentUsage);
}

Example 13 with QuotaResource

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

the class QuotaAwareOperationController method pollQuotaExceedAllowedRequestsIfAny.

/**
 * Poll for out of quota requests that are allowed to exceed quota.
 * @param requestsToSend {@link List} of {@link RequestInfo} to be sent.
 * @param requestQueue {@link Map} of {@link QuotaResource} to {@link List} of {@link RequestInfo} from which the requests to be sent will be polled.
 */
private void pollQuotaExceedAllowedRequestsIfAny(List<RequestInfo> requestsToSend, Map<QuotaResource, LinkedList<RequestInfo>> requestQueue) {
    List<QuotaResource> quotaResources = new ArrayList<>(requestQueue.keySet());
    Collections.shuffle(quotaResources);
    while (!requestQueue.isEmpty()) {
        for (QuotaResource quotaResource : quotaResources) {
            RequestInfo requestInfo = requestQueue.get(quotaResource).getFirst();
            if (!requestInfo.getChargeable().quotaExceedAllowed()) {
                // If quota exceeded requests aren't allowed, then there is nothing more to do.
                return;
            }
            requestInfo.getChargeable().charge();
            requestsToSend.add(requestInfo);
            requestQueue.get(quotaResource).removeFirst();
            if (requestQueue.get(quotaResource).isEmpty()) {
                requestQueue.remove(quotaResource);
            }
        }
    }
}

Example 14 with QuotaResource

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

the class OperationQuotaChargerTest method testGetQuotaResource.

@Test
public void testGetQuotaResource() throws Exception {
    // getQuotaResource should return null if quotaChargeCallback is null.
    OperationQuotaCharger operationQuotaCharger = new OperationQuotaCharger(null, BLOBID, "GetOperation");
    Assert.assertNull("getQuotaResource should return null if quotaChargeCallback is null.", operationQuotaCharger.getQuotaResource());
    QuotaChargeCallback quotaChargeCallback = Mockito.mock(QuotaChargeCallback.class);
    operationQuotaCharger = new OperationQuotaCharger(quotaChargeCallback, BLOBID, "GetOperation");
    // getQuotaResource should return what quotaChargeCallback.getQuotaResource returns.
    QuotaResource quotaResource = new QuotaResource("test", QuotaResourceType.ACCOUNT);
    Mockito.when(quotaChargeCallback.getQuotaResource()).thenReturn(quotaResource);
    Assert.assertEquals("getQuotaResource should return what quotaChargeCallback.getQuotaResource returns.", quotaResource, quotaChargeCallback.getQuotaResource());
    // getQuotaResource should return null if quotaChargeCallback throws exception.
    Mockito.when(quotaChargeCallback.getQuotaResource()).thenThrow(new QuotaException("", new RestServiceException("", RestServiceErrorCode.InternalServerError), false));
    Assert.assertNull("getQuotaResource should return null if quotaChargeCallback is null.", operationQuotaCharger.getQuotaResource());
}