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

Example 1 with QuotaMethod

use of com.github.ambry.quota.QuotaMethod 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);
}

Example 2 with QuotaMethod

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

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

the class QuotaAwareOperationController method drainRequestQueue.

/**
 * Drain the request queue based on quota and system resources and update the requests to be send to {@code requestsToSend}.
 * @param requestsToSend a list of {@link RequestInfo} that will contain the requests to be sent out.
 */
private void drainRequestQueue(List<RequestInfo> requestsToSend) {
    for (QuotaMethod quotaMethod : QuotaMethod.values()) {
        pollQuotaCompliantRequests(requestsToSend, getRequestQueue(quotaMethod));
        pollQuotaExceedAllowedRequestsIfAny(requestsToSend, getRequestQueue(quotaMethod));
    }
}