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

Example 1 with QuotaChargeCallback

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

the class MockRouterCallback method testInstantiation.

/**
 * Test {@link GetBlobInfoOperation} instantiation and validate the get methods.
 */
@Test
public void testInstantiation() {
    BlobId blobId = new BlobId(routerConfig.routerBlobidCurrentVersion, BlobId.BlobIdType.NATIVE, ClusterMap.UNKNOWN_DATACENTER_ID, Utils.getRandomShort(random), Utils.getRandomShort(random), mockClusterMap.getWritablePartitionIds(MockClusterMap.DEFAULT_PARTITION_CLASS).get(0), false, BlobId.BlobDataType.DATACHUNK);
    Callback<GetBlobResultInternal> getOperationCallback = (result, exception) -> {
    // no op.
    };
    // test a good case
    GetBlobInfoOperation op = new GetBlobInfoOperation(routerConfig, routerMetrics, mockClusterMap, responseHandler, blobId, options, getOperationCallback, routerCallback, kms, cryptoService, cryptoJobHandler, time, false, quotaChargeCallback);
    Assert.assertEquals("Callback must match", getOperationCallback, op.getCallback());
    Assert.assertEquals("Blob ids must match", blobId.getID(), op.getBlobIdStr());
    // test the case where the tracker type is bad
    Properties properties = getNonBlockingRouterProperties(true);
    properties.setProperty("router.get.operation.tracker.type", "NonExistentTracker");
    RouterConfig badConfig = new RouterConfig(new VerifiableProperties(properties));
    try {
        new GetBlobInfoOperation(badConfig, routerMetrics, mockClusterMap, responseHandler, blobId, options, getOperationCallback, routerCallback, kms, cryptoService, cryptoJobHandler, time, false, quotaChargeCallback);
        Assert.fail("Instantiation of GetBlobInfoOperation with an invalid tracker type must fail");
    } catch (IllegalArgumentException e) {
    // expected. Nothing to do.
    }
}

Example 2 with QuotaChargeCallback

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

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

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

the class OperationQuotaChargerTest method testCheck.

@Test
public void testCheck() throws Exception {
    // check should return true if quotaChargeCallback is null.
    OperationQuotaCharger operationQuotaCharger = new OperationQuotaCharger(null, BLOBID, "GetOperation");
    Assert.assertTrue("check should return true if quotaChargeCallback is null.", operationQuotaCharger.check());
    QuotaChargeCallback quotaChargeCallback = Mockito.mock(QuotaChargeCallback.class);
    operationQuotaCharger = new OperationQuotaCharger(quotaChargeCallback, BLOBID, "GetOperation");
    // check should return true if quotaChargeCallback returns true.
    Mockito.when(quotaChargeCallback.check()).thenReturn(true);
    Assert.assertTrue("check should return true if quotaChargeCallback returns true.", operationQuotaCharger.check());
    // check should return false if quotaChargeCallback returns false.
    Mockito.when(quotaChargeCallback.check()).thenReturn(false);
    Assert.assertFalse("check should return false if quotaChargeCallback returns false.", operationQuotaCharger.check());
    // check should return false even if quotaChargeCallback returns false when isCharged is true.
    Mockito.doNothing().when(quotaChargeCallback).charge();
    // sets isCharged to true.
    operationQuotaCharger.charge();
    Mockito.when(quotaChargeCallback.check()).thenReturn(false);
    Assert.assertTrue("check should return true even if quotaChargeCallback returns false when isCharged is true.", operationQuotaCharger.check());
}

Example 5 with QuotaChargeCallback

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

the class OperationQuotaChargerTest method testQuotaExceedAllowed.

@Test
public void testQuotaExceedAllowed() throws Exception {
    // quotaExceedAllowed should return true if quotaChargeCallback is null.
    OperationQuotaCharger operationQuotaCharger = new OperationQuotaCharger(null, BLOBID, "GetOperation");
    Assert.assertTrue("quotaExceedAllowed should return true if quotaChargeCallback is null.", operationQuotaCharger.quotaExceedAllowed());
    QuotaChargeCallback quotaChargeCallback = Mockito.mock(QuotaChargeCallback.class);
    operationQuotaCharger = new OperationQuotaCharger(quotaChargeCallback, BLOBID, "GetOperation");
    // quotaExceedAllowed should return true if quotaChargeCallback.quotaExceedAllowed returns true.
    Mockito.when(quotaChargeCallback.quotaExceedAllowed()).thenReturn(true);
    Assert.assertTrue("quotaExceedAllowed should return true if quotaChargeCallback.quotaExceedAllowed returns true.", operationQuotaCharger.quotaExceedAllowed());
    // quotaExceedAllowed should return false if quotaChargeCallback.quotaExceedAllowed returns false.
    Mockito.when(quotaChargeCallback.quotaExceedAllowed()).thenReturn(false);
    Assert.assertFalse("quotaExceedAllowed should return false if quotaChargeCallback.quotaExceedAllowed returns false.", operationQuotaCharger.quotaExceedAllowed());
}