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

Example 6 with QuotaChargeCallback

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

the class GetBlobOperationTest method testInstantiation.

/**
 * Test {@link GetBlobOperation} instantiation and validate the get methods.
 */
@Test
public void testInstantiation() {
    Callback<GetBlobResultInternal> getRouterCallback = new Callback<GetBlobResultInternal>() {

        @Override
        public void onCompletion(GetBlobResultInternal result, Exception exception) {
        // no op.
        }
    };
    blobId = new BlobId(routerConfig.routerBlobidCurrentVersion, BlobId.BlobIdType.NATIVE, mockClusterMap.getLocalDatacenterId(), Utils.getRandomShort(TestUtils.RANDOM), Utils.getRandomShort(TestUtils.RANDOM), mockClusterMap.getWritablePartitionIds(MockClusterMap.DEFAULT_PARTITION_CLASS).get(0), false, BlobId.BlobDataType.DATACHUNK);
    blobIdStr = blobId.getID();
    // test a good case
    // operationCount is not incremented here as this operation is not taken to completion.
    GetBlobOperation op = new GetBlobOperation(routerConfig, routerMetrics, mockClusterMap, responseHandler, blobId, new GetBlobOptionsInternal(new GetBlobOptionsBuilder().build(), false, routerMetrics.ageAtGet), getRouterCallback, routerCallback, blobIdFactory, kms, cryptoService, cryptoJobHandler, time, false, quotaChargeCallback);
    Assert.assertEquals("Callbacks must match", getRouterCallback, op.getCallback());
    Assert.assertEquals("Blob ids must match", blobIdStr, op.getBlobIdStr());
    // test the case where the tracker type is bad
    Properties properties = getDefaultNonBlockingRouterProperties(true);
    properties.setProperty("router.get.operation.tracker.type", "NonExistentTracker");
    RouterConfig badConfig = new RouterConfig(new VerifiableProperties(properties));
    try {
        new GetBlobOperation(badConfig, routerMetrics, mockClusterMap, responseHandler, blobId, new GetBlobOptionsInternal(new GetBlobOptionsBuilder().build(), false, routerMetrics.ageAtGet), getRouterCallback, routerCallback, blobIdFactory, kms, cryptoService, cryptoJobHandler, time, false, quotaChargeCallback);
        Assert.fail("Instantiation of GetBlobOperation with an invalid tracker type must fail");
    } catch (IllegalArgumentException e) {
    // expected. Nothing to do.
    }
}

Example 7 with QuotaChargeCallback

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

Example 8 with QuotaChargeCallback

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

the class OperationQuotaChargerTest method testCharge.

@Test
public void testCharge() throws Exception {
    // charge should return true if quotaChargeCallback is null.
    OperationQuotaCharger operationQuotaCharger = new OperationQuotaCharger(null, BLOBID, "GetOperation");
    Assert.assertTrue("charge should return true if quotaChargeCallback is null.", operationQuotaCharger.charge());
    QuotaChargeCallback quotaChargeCallback = Mockito.mock(QuotaChargeCallback.class);
    operationQuotaCharger = new OperationQuotaCharger(quotaChargeCallback, BLOBID, "GetOperation");
    // charge should return false if quotaChargeCallback throws exception and isCharged is false.
    Mockito.doThrow(new QuotaException("too many requests", new RouterException("", RouterErrorCode.TooManyRequests), true)).when(quotaChargeCallback).charge();
    Assert.assertFalse("charge should return true if quotaChargeCallback throws exception and isCharged is false.", operationQuotaCharger.charge());
    Mockito.verify(quotaChargeCallback, Mockito.times(1)).charge();
    // charge should return true if quotaChargeCallback.charge goes through.
    Mockito.doNothing().when(quotaChargeCallback).charge();
    Assert.assertTrue("charge should return true if quotaChargeCallback.charge goes through.", operationQuotaCharger.charge());
    Mockito.verify(quotaChargeCallback, Mockito.times(2)).charge();
    // Once isCharged is true, charge should never call quotaChargeCallback.charge.
    Mockito.doNothing().when(quotaChargeCallback).charge();
    Assert.assertTrue("Once isCharged is true, charge should never call quotaChargeCallback.charge.", operationQuotaCharger.charge());
    Mockito.verify(quotaChargeCallback, Mockito.times(2)).charge();
}

Example 9 with QuotaChargeCallback

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

the class GetManagerTest method testBadCallback.

/**
 * Test that a bad user defined callback will not crash the router.
 * @param getBlobCallback User defined callback to be called after getBlob operation.
 * @param getBlobCallbackCalled This latch should be at 0 after {@code getBlobCallback} has been called.
 * @param checkBadCallbackBlob {@code true} if the blob contents provided by the getBlob operation with the bad
 *                             callback should be inspected for correctness.
 * @throws Exception
 */
private void testBadCallback(Callback<GetBlobResult> getBlobCallback, CountDownLatch getBlobCallbackCalled, Boolean checkBadCallbackBlob) throws Exception {
    router = getNonBlockingRouter();
    setOperationParams(LARGE_BLOB_SIZE, new GetBlobOptionsBuilder().build());
    final CountDownLatch getBlobInfoCallbackCalled = new CountDownLatch(1);
    String blobId = router.putBlob(putBlobProperties, putUserMetadata, putChannel, new PutBlobOptionsBuilder().build()).get();
    List<Future<GetBlobResult>> getBlobInfoFutures = new ArrayList<>();
    List<Future<GetBlobResult>> getBlobDataFutures = new ArrayList<>();
    GetBlobOptions infoOptions = new GetBlobOptionsBuilder().operationType(GetBlobOptions.OperationType.BlobInfo).build();
    GetBlobOptions dataOptions = new GetBlobOptionsBuilder().operationType(GetBlobOptions.OperationType.Data).build();
    for (int i = 0; i < 5; i++) {
        if (i == 1) {
            getBlobInfoFutures.add(router.getBlob(blobId, infoOptions, new Callback<GetBlobResult>() {

                @Override
                public void onCompletion(GetBlobResult result, Exception exception) {
                    getBlobInfoCallbackCalled.countDown();
                    throw new RuntimeException("Throwing an exception in the user callback");
                }
            }, quotaChargeCallback));
            getBlobDataFutures.add(router.getBlob(blobId, dataOptions, getBlobCallback, quotaChargeCallback));
        } else {
            getBlobInfoFutures.add(router.getBlob(blobId, infoOptions));
            getBlobDataFutures.add(router.getBlob(blobId, dataOptions));
        }
    }
    options = dataOptions;
    for (int i = 0; i < getBlobDataFutures.size(); i++) {
        if (i != 1 || checkBadCallbackBlob) {
            compareContent(getBlobDataFutures.get(i).get().getBlobDataChannel());
        }
    }
    options = infoOptions;
    for (Future<GetBlobResult> future : getBlobInfoFutures) {
        compareBlobInfo(future.get().getBlobInfo(), -1);
    }
    Assert.assertTrue("getBlobInfo callback not called.", getBlobInfoCallbackCalled.await(2, TimeUnit.SECONDS));
    Assert.assertTrue("getBlob callback not called.", getBlobCallbackCalled.await(2, TimeUnit.SECONDS));
    Assert.assertEquals("All operations should be finished.", 0, router.getOperationsCount());
    Assert.assertTrue("Router should not be closed", router.isOpen());
    // Test that GetManager is still operational
    setOperationParams(CHUNK_SIZE, new GetBlobOptionsBuilder().build());
    blobId = router.putBlob(putBlobProperties, putUserMetadata, putChannel, new PutBlobOptionsBuilder().build()).get();
    getBlobAndCompareContent(blobId, -1);
    this.options = infoOptions;
    getBlobAndCompareContent(blobId, -1);
    router.close();
}

Example 10 with QuotaChargeCallback

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

the class NonBlockingRouterQuotaCallbackTest method testRouterWithDefaultQuotaCallback.

/**
 * Test default {@link QuotaChargeCallback} doesn't charge anything and doesn't error out when throttling is disabled.
 */
@Test
public void testRouterWithDefaultQuotaCallback() throws Exception {
    try {
        setRouter();
        assertExpectedThreadCounts(2, 1);
        AtomicInteger listenerCalledCount = new AtomicInteger(0);
        QuotaConfig quotaConfig = new QuotaConfig(new VerifiableProperties(new Properties()));
        QuotaManager quotaManager = new ChargeTesterQuotaManager(quotaConfig, new SimpleQuotaRecommendationMergePolicy(quotaConfig), accountService, null, new QuotaMetrics(new MetricRegistry()), listenerCalledCount);
        QuotaChargeCallback quotaChargeCallback = QuotaUtils.buildQuotaChargeCallback(null, quotaManager, true);
        int blobSize = 3000;
        setOperationParams(blobSize, TTL_SECS);
        String compositeBlobId = router.putBlob(putBlobProperties, putUserMetadata, putChannel, PutBlobOptions.DEFAULT, null, quotaChargeCallback).get();
        assertEquals(0, listenerCalledCount.get());
        RetainingAsyncWritableChannel retainingAsyncWritableChannel = new RetainingAsyncWritableChannel();
        router.getBlob(compositeBlobId, new GetBlobOptionsBuilder().build(), null, quotaChargeCallback).get().getBlobDataChannel().readInto(retainingAsyncWritableChannel, null).get();
        // read out all the chunks.
        retainingAsyncWritableChannel.consumeContentAsInputStream().close();
        assertEquals(0, listenerCalledCount.get());
    } finally {
        router.close();
        assertExpectedThreadCounts(0, 0);
        // submission after closing should return a future that is already done.
        assertClosed();
    }
}