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

Example 6 with QuotaResource

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

the class NonBlockingRouterQuotaCallbackTest method testRouterWithQuotaCallback.

/**
 * Test Router with single scaling unit for correct accounting in {@link QuotaChargeCallback}.
 */
@Test
public void testRouterWithQuotaCallback() throws Exception {
    try {
        setRouter();
        assertExpectedThreadCounts(2, 1);
        AtomicLong listenerCalledCount = new AtomicLong(0);
        int expectedChargeCallbackCount = 0;
        // create a quota charge listener that increments an atomic counter everytime its called.
        // Also tests that in case quota if charged in tracking mode with throttleInProgress config set to false
        // then the requests go through even in case of exception.
        QuotaChargeCallback quotaChargeCallback = new QuotaChargeCallback() {

            @Override
            public void charge(long chunkSize) throws QuotaException {
                listenerCalledCount.addAndGet(chunkSize);
                throw new QuotaException("exception during check and charge", new RouterException("Quota exceeded.", RouterErrorCode.TooManyRequests), false);
            }

            @Override
            public void charge() throws QuotaException {
                charge(quotaAccountingSize);
            }

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

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

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

            @Override
            public QuotaMethod getQuotaMethod() {
                return null;
            }
        };
        // test for a composite blob.
        int blobSize = 3000;
        setOperationParams(blobSize, TTL_SECS);
        String compositeBlobId = router.putBlob(putBlobProperties, putUserMetadata, putChannel, PutBlobOptions.DEFAULT, null, quotaChargeCallback).get();
        expectedChargeCallbackCount += blobSize;
        assertEquals(expectedChargeCallbackCount, listenerCalledCount.get());
        RetainingAsyncWritableChannel retainingAsyncWritableChannel = new RetainingAsyncWritableChannel();
        router.getBlob(compositeBlobId, new GetBlobOptionsBuilder().build(), null, quotaChargeCallback).get().getBlobDataChannel().readInto(retainingAsyncWritableChannel, null).get();
        expectedChargeCallbackCount += blobSize;
        // read out all the chunks to make sure all the chunks are consumed and accounted for.
        retainingAsyncWritableChannel.consumeContentAsInputStream().close();
        assertEquals(expectedChargeCallbackCount, listenerCalledCount.get());
        // test for regular blobs.
        setOperationParams();
        List<String> blobIds = new ArrayList<>();
        for (int i = 0; i < 2; i++) {
            setOperationParams();
            String blobId = router.putBlob(putBlobProperties, putUserMetadata, putChannel, PutBlobOptions.DEFAULT, null, quotaChargeCallback).get();
            assertEquals(expectedChargeCallbackCount += PUT_CONTENT_SIZE, listenerCalledCount.get());
            logger.info("Put blob {}", blobId);
            blobIds.add(blobId);
        }
        setOperationParams();
        for (String blobId : blobIds) {
            router.getBlob(blobId, new GetBlobOptionsBuilder().build(), null, quotaChargeCallback).get();
            assertEquals(expectedChargeCallbackCount += PUT_CONTENT_SIZE, listenerCalledCount.get());
            router.updateBlobTtl(blobId, null, Utils.Infinite_Time, null, quotaChargeCallback).get();
            assertEquals(expectedChargeCallbackCount += quotaAccountingSize, listenerCalledCount.get());
            router.getBlob(blobId, new GetBlobOptionsBuilder().build(), null, quotaChargeCallback).get();
            assertEquals(expectedChargeCallbackCount += PUT_CONTENT_SIZE, listenerCalledCount.get());
            router.getBlob(blobId, new GetBlobOptionsBuilder().operationType(GetBlobOptions.OperationType.BlobInfo).build(), null, quotaChargeCallback).get();
            assertEquals(expectedChargeCallbackCount += quotaAccountingSize, listenerCalledCount.get());
            router.deleteBlob(blobId, null, null, quotaChargeCallback).get();
            assertEquals(expectedChargeCallbackCount += quotaAccountingSize, listenerCalledCount.get());
            try {
                router.getBlob(blobId, new GetBlobOptionsBuilder().build(), null, quotaChargeCallback).get();
                fail("Get blob should fail");
            } catch (ExecutionException e) {
                RouterException r = (RouterException) e.getCause();
                Assert.assertEquals("BlobDeleted error is expected", RouterErrorCode.BlobDeleted, r.getErrorCode());
                assertEquals(expectedChargeCallbackCount, listenerCalledCount.get());
            }
            router.getBlob(blobId, new GetBlobOptionsBuilder().getOption(GetOption.Include_Deleted_Blobs).build(), null, quotaChargeCallback).get();
            assertEquals(expectedChargeCallbackCount += PUT_CONTENT_SIZE, listenerCalledCount.get());
            router.getBlob(blobId, new GetBlobOptionsBuilder().getOption(GetOption.Include_All).build(), null, quotaChargeCallback).get();
            assertEquals(expectedChargeCallbackCount += PUT_CONTENT_SIZE, listenerCalledCount.get());
        }
        // test for stitched blobs.
        blobIds = new ArrayList<>();
        int stitchedBlobCount = 2;
        for (int i = 0; i < stitchedBlobCount; i++) {
            setOperationParams();
            String blobId = router.putBlob(putBlobProperties, putUserMetadata, putChannel, PutBlobOptions.DEFAULT, null, quotaChargeCallback).get();
            assertEquals(expectedChargeCallbackCount += PUT_CONTENT_SIZE, listenerCalledCount.get());
            logger.info("Put blob {}", blobId);
            blobIds.add(blobId);
        }
        String stitchedBlobId = router.stitchBlob(putBlobProperties, putUserMetadata, blobIds.stream().map(blobId -> new ChunkInfo(blobId, PUT_CONTENT_SIZE, Utils.Infinite_Time)).collect(Collectors.toList()), null, quotaChargeCallback).get();
        assertEquals(expectedChargeCallbackCount, listenerCalledCount.get());
        retainingAsyncWritableChannel = new RetainingAsyncWritableChannel();
        router.getBlob(stitchedBlobId, new GetBlobOptionsBuilder().build(), null, quotaChargeCallback).get().getBlobDataChannel().readInto(retainingAsyncWritableChannel, null).get();
        // read out all the chunks to make sure all the chunks are consumed and accounted for.
        retainingAsyncWritableChannel.consumeContentAsInputStream().close();
        assertEquals(expectedChargeCallbackCount += (PUT_CONTENT_SIZE * stitchedBlobCount), listenerCalledCount.get());
        router.updateBlobTtl(stitchedBlobId, null, Utils.Infinite_Time, null, quotaChargeCallback).get();
        assertEquals(expectedChargeCallbackCount += quotaAccountingSize, listenerCalledCount.get());
        router.deleteBlob(stitchedBlobId, null, null, quotaChargeCallback).get();
        assertEquals(expectedChargeCallbackCount + quotaAccountingSize, listenerCalledCount.get());
    } finally {
        router.close();
        assertExpectedThreadCounts(0, 0);
        // submission after closing should return a future that is already done.
        assertClosed();
    }
}

Example 7 with QuotaResource

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

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

the class AmbryCUQuotaEnforcer method charge.

@Override
public QuotaRecommendation charge(RestRequest restRequest, Map<QuotaName, Double> requestCostMap) throws QuotaException {
    final QuotaName quotaName = QuotaUtils.getCUQuotaName(restRequest);
    if (requestCostMap.isEmpty() || !requestCostMap.containsKey(quotaName)) {
        String errorMessage = String.format("No %s cost provided for request %s. Nothing to charge", quotaName.name(), RestUtils.convertToStr(restRequest));
        throw new QuotaException(errorMessage, true);
    }
    final QuotaResource quotaResource = QuotaResource.fromRestRequest(restRequest);
    doAndHandleException(() -> {
        quotaSource.chargeUsage(quotaResource, quotaName, requestCostMap.get(quotaName));
        return null;
    }, String.format("Could not charge for request %s due to", restRequest));
    return recommend(restRequest);
}

Example 9 with QuotaResource

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

the class AmbryCUQuotaEnforcer method recommend.

@Override
public QuotaRecommendation recommend(RestRequest restRequest) throws QuotaException {
    final QuotaName quotaName = QuotaUtils.getCUQuotaName(restRequest);
    float usage = doAndHandleException(() -> {
        final QuotaResource quotaResource = QuotaResource.fromRestRequest(restRequest);
        return quotaSource.getUsage(quotaResource, quotaName);
    }, String.format("Could not recommend for request %s due to", restRequest));
    return buildQuotaRecommendation(usage, quotaName);
}

Example 10 with QuotaResource

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

the class AmbryCUQuotaSourceTest method testCreation.

@Test
public void testCreation() throws QuotaException {
    Assert.assertEquals(4, ambryCUQuotaSource.getAllQuota().size());
    Assert.assertEquals(1024000000, (long) ambryCUQuotaSource.getQuota(new QuotaResource("101_1", QuotaResourceType.CONTAINER), QuotaName.READ_CAPACITY_UNIT).getQuotaValue());
    Assert.assertEquals(1024000000, (long) ambryCUQuotaSource.getQuota(new QuotaResource("101_1", QuotaResourceType.CONTAINER), QuotaName.WRITE_CAPACITY_UNIT).getQuotaValue());
    Assert.assertEquals(258438456, (long) ambryCUQuotaSource.getQuota(new QuotaResource("101_2", QuotaResourceType.CONTAINER), QuotaName.READ_CAPACITY_UNIT).getQuotaValue());
    Assert.assertEquals(258438456, (long) ambryCUQuotaSource.getQuota(new QuotaResource("101_2", QuotaResourceType.CONTAINER), QuotaName.WRITE_CAPACITY_UNIT).getQuotaValue());
    Assert.assertEquals(1024000000, (long) ambryCUQuotaSource.getQuota(new QuotaResource("102_1", QuotaResourceType.CONTAINER), QuotaName.READ_CAPACITY_UNIT).getQuotaValue());
    Assert.assertEquals(1024000000, (long) ambryCUQuotaSource.getQuota(new QuotaResource("102_1", QuotaResourceType.CONTAINER), QuotaName.WRITE_CAPACITY_UNIT).getQuotaValue());
    try {
        ambryCUQuotaSource.getQuota(new QuotaResource("101", QuotaResourceType.ACCOUNT), QuotaName.WRITE_CAPACITY_UNIT);
        Assert.fail("If quota is not present, it should throw an exception");
    } catch (QuotaException quotaException) {
    }
    try {
        ambryCUQuotaSource.getQuota(new QuotaResource("102", QuotaResourceType.ACCOUNT), QuotaName.WRITE_CAPACITY_UNIT);
        Assert.fail("If quota is not present, it should throw an exception");
    } catch (QuotaException quotaException) {
    }
    Assert.assertEquals(10737418240L, (long) ambryCUQuotaSource.getQuota(new QuotaResource("103", QuotaResourceType.ACCOUNT), QuotaName.READ_CAPACITY_UNIT).getQuotaValue());
    Assert.assertEquals(10737418240L, (long) ambryCUQuotaSource.getQuota(new QuotaResource("103", QuotaResourceType.ACCOUNT), QuotaName.WRITE_CAPACITY_UNIT).getQuotaValue());
    Assert.assertEquals(0, ambryCUQuotaSource.getSystemResourceUsage(QuotaName.READ_CAPACITY_UNIT), 0.01);
    Assert.assertEquals(0, ambryCUQuotaSource.getSystemResourceUsage(QuotaName.READ_CAPACITY_UNIT), 0.01);
    Assert.assertEquals(4, ambryCUQuotaSource.getAllQuotaUsage().size());
    Assert.assertEquals(0, (long) ambryCUQuotaSource.getUsage(new QuotaResource("101_1", QuotaResourceType.CONTAINER), QuotaName.READ_CAPACITY_UNIT));
    Assert.assertEquals(0, (long) ambryCUQuotaSource.getUsage(new QuotaResource("101_1", QuotaResourceType.CONTAINER), QuotaName.WRITE_CAPACITY_UNIT));
    Assert.assertEquals(0, (long) ambryCUQuotaSource.getUsage(new QuotaResource("101_2", QuotaResourceType.CONTAINER), QuotaName.READ_CAPACITY_UNIT));
    Assert.assertEquals(0, (long) ambryCUQuotaSource.getUsage(new QuotaResource("101_2", QuotaResourceType.CONTAINER), QuotaName.WRITE_CAPACITY_UNIT));
    Assert.assertEquals(0, (long) ambryCUQuotaSource.getUsage(new QuotaResource("102_1", QuotaResourceType.CONTAINER), QuotaName.READ_CAPACITY_UNIT));
    Assert.assertEquals(0, (long) ambryCUQuotaSource.getUsage(new QuotaResource("102_1", QuotaResourceType.CONTAINER), QuotaName.WRITE_CAPACITY_UNIT));
}