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

Example 1 with QuotaResource

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

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

the class StorageQuotaEnforcer method charge.

/**
 * Add given {@code usage} to the current storage usage of account/container carried in {@code restRequest} even
 * if the result exceeds quota for the target account/container. If there is no account and container found in
 * {@code restRequest}, then this is a no-op. If there is no quota found for the account/container, then this is
 * a no-op. A {@link Pair} whose first element is quota and second element is the storage usage after charge.
 * @param restRequest the {@link RestRequest} that carries account and container in the header.
 * @param usage the usage to charge
 * @return A {@link Pair} whose first element is quota and second element is the storage usage after charge.
 */
Pair<Long, Long> charge(RestRequest restRequest, long usage) {
    long quotaValue = -1L;
    long usageAfterCharge = 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) {
            AtomicLong existingUsage = new AtomicLong();
            storageUsages.compute(quotaResource, (k, v) -> {
                existingUsage.set(v == null ? 0 : v);
                if (v == null) {
                    return usage;
                }
                return v + usage;
            });
            usageAfterCharge = existingUsage.addAndGet(usage);
        }
    } catch (Exception e) {
        logger.error("Failed to charge for RestRequest {}", restRequest, e);
    }
    return new Pair<>(quotaValue, usageAfterCharge);
}

Example 3 with QuotaResource

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

the class QuotaAwareOperationController method addToRequestQueue.

/**
 * Add the specified {@link List} of {@link RequestInfo}s to the request queue.
 * @param requestInfos {@link List} of {@link RequestInfo} objects.
 */
private void addToRequestQueue(List<RequestInfo> requestInfos) {
    for (RequestInfo requestInfo : requestInfos) {
        QuotaResource quotaResource = requestInfo.getChargeable().getQuotaResource();
        if (quotaResource == null) {
            quotaResource = UNKNOWN_QUOTA_RESOURCE;
        }
        getRequestQueue(requestInfo.getChargeable().getQuotaMethod()).putIfAbsent(quotaResource, new LinkedList<>());
        getRequestQueue(requestInfo.getChargeable().getQuotaMethod()).get(quotaResource).add(requestInfo);
    }
}

Example 4 with QuotaResource

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

the class QuotaAwareOperationController method pollQuotaCompliantRequests.

/**
 * Drain the request queue based on resource quota only and update the requests to be sent to {@code requestsToSend}.
 * @param requestsToSend a list of {@link RequestInfo} that will contain the requests to be sent out.
 * @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 pollQuotaCompliantRequests(List<RequestInfo> requestsToSend, Map<QuotaResource, LinkedList<RequestInfo>> requestQueue) {
    for (QuotaResource quotaResource : requestQueue.keySet()) {
        if (quotaResource.equals(UNKNOWN_QUOTA_RESOURCE)) {
            // If there are requests for which QuotaResource couldn't be found, this is most likely a bug.
            // As a temporary hack, we will consider all those requests to be quota compliant.
            requestsToSend.addAll(requestQueue.get(UNKNOWN_QUOTA_RESOURCE));
            requestQueue.remove(UNKNOWN_QUOTA_RESOURCE);
            continue;
        }
        while (!requestQueue.get(quotaResource).isEmpty()) {
            RequestInfo requestInfo = requestQueue.get(quotaResource).getFirst();
            if (requestInfo.getChargeable().check()) {
                requestsToSend.add(requestInfo);
                requestQueue.get(quotaResource).removeFirst();
                requestInfo.getChargeable().charge();
            } else {
                break;
            }
        }
        if (requestQueue.get(quotaResource).isEmpty()) {
            requestQueue.remove(quotaResource);
        }
    }
}

Example 5 with QuotaResource

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

the class MockReadableStreamChannel method testBadCallback.

/**
 * Test that a bad user defined callback will not crash the router.
 * @throws Exception
 */
@Test
public void testBadCallback() throws Exception {
    RequestAndResult req = new RequestAndResult(chunkSize * 5 + random.nextInt(chunkSize - 1) + 1);
    router = getNonBlockingRouter();
    final CountDownLatch callbackCalled = new CountDownLatch(1);
    requestAndResultsList.clear();
    for (int i = 0; i < 4; i++) {
        requestAndResultsList.add(new RequestAndResult(chunkSize + random.nextInt(5) * random.nextInt(chunkSize)));
    }
    instantiateNewRouterForPuts = false;
    ReadableStreamChannel putChannel = new ByteBufferReadableStreamChannel(ByteBuffer.wrap(req.putContent));
    Future future = router.putBlob(req.putBlobProperties, req.putUserMetadata, putChannel, req.options, (result, exception) -> {
        callbackCalled.countDown();
        throw new RuntimeException("Throwing an exception in the user callback");
    }, new QuotaChargeCallback() {

        @Override
        public void charge(long chunkSize) {
        }

        @Override
        public void charge() {
        }

        @Override
        public boolean check() {
            return false;
        }

        @Override
        public boolean quotaExceedAllowed() {
            return false;
        }

        @Override
        public QuotaResource getQuotaResource() throws QuotaException {
            return null;
        }

        @Override
        public QuotaMethod getQuotaMethod() {
            return null;
        }
    });
    submitPutsAndAssertSuccess(false);
    // future.get() for operation with bad callback should still succeed
    future.get();
    Assert.assertTrue("Callback not called.", callbackCalled.await(MAX_WAIT_MS, TimeUnit.MILLISECONDS));
    assertEquals("All operations should be finished.", 0, router.getOperationsCount());
    Assert.assertTrue("Router should not be closed", router.isOpen());
    // Test that PutManager is still operational
    requestAndResultsList.clear();
    requestAndResultsList.add(new RequestAndResult(chunkSize + random.nextInt(5) * random.nextInt(chunkSize)));
    instantiateNewRouterForPuts = false;
    submitPutsAndAssertSuccess(true);
}