use of com.github.ambry.rest.RestRequest in project ambry by linkedin.
the class AmbryQuotaManagerTest method testChargeAndRecommendForForcedCharge.
/**
* Test {@link QuotaManager#chargeAndRecommend} when quota exceed check is set to true and force charge is set to true.
* @throws Exception in case of any exception.
*/
private void testChargeAndRecommendForForcedCharge() throws Exception {
RestRequest restRequest = QuotaTestUtils.createRestRequest(ACCOUNT, ACCOUNT.getAllContainers().iterator().next(), RestMethod.GET);
AmbryQuotaManager ambryQuotaManager = buildAmbryQuotaManagerWithNoEnforcers();
ambryQuotaManager.init();
// 1. test that chargeAndRecommend returns ALLOW if there are no enforcers.
QuotaAction quotaAction = ambryQuotaManager.chargeAndRecommend(restRequest, null, true, true);
Assert.assertEquals(QuotaAction.ALLOW, quotaAction);
// 2. test that enforcers are called and usage is charged.
ambryQuotaManager = buildAmbryQuotaManagerWithAmbryCUEnforcer();
testAmbryCUQuotaEnforcer.resetDefaults();
ambryQuotaManager.init();
Map<QuotaName, Double> costMap = Collections.singletonMap(QuotaName.READ_CAPACITY_UNIT, 11.0);
quotaAction = ambryQuotaManager.chargeAndRecommend(restRequest, costMap, true, true);
Assert.assertEquals(QuotaAction.ALLOW, quotaAction);
Assert.assertEquals(1, testAmbryCUQuotaEnforcer.initCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer.recommendCallCount);
Assert.assertEquals(1, testAmbryCUQuotaEnforcer.chargeCallCount);
Assert.assertEquals(110, testAmbryCUQuotaEnforcer.getQuotaSource().getUsage(QuotaResource.fromRestRequest(restRequest), QuotaName.READ_CAPACITY_UNIT), 0.0001);
Assert.assertEquals(11, testAmbryCUQuotaEnforcer.getQuotaSource().getSystemResourceUsage(QuotaName.READ_CAPACITY_UNIT), 0.0001);
// 3. test that enforcers are called and usage is charged even if usage is outside quota limits.
testAmbryCUQuotaEnforcer.resetDefaults();
costMap = Collections.singletonMap(QuotaName.READ_CAPACITY_UNIT, 100.0);
quotaAction = ambryQuotaManager.chargeAndRecommend(restRequest, costMap, true, true);
Assert.assertEquals(QuotaAction.ALLOW, quotaAction);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer.initCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer.recommendCallCount);
Assert.assertEquals(1, testAmbryCUQuotaEnforcer.chargeCallCount);
Assert.assertEquals(1110, testAmbryCUQuotaEnforcer.getQuotaSource().getUsage(QuotaResource.fromRestRequest(restRequest), QuotaName.READ_CAPACITY_UNIT), 0.0001);
Assert.assertEquals(111, testAmbryCUQuotaEnforcer.getQuotaSource().getSystemResourceUsage(QuotaName.READ_CAPACITY_UNIT), 0.0001);
// 3. test that enforcers are called and usage is charged even if usage is outside quota limits and exceed is not allowed.
testAmbryCUQuotaEnforcer.resetDefaults();
costMap = Collections.singletonMap(QuotaName.READ_CAPACITY_UNIT, 100.0);
quotaAction = ambryQuotaManager.chargeAndRecommend(restRequest, costMap, true, true);
Assert.assertEquals(QuotaAction.ALLOW, quotaAction);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer.initCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer.recommendCallCount);
Assert.assertEquals(1, testAmbryCUQuotaEnforcer.chargeCallCount);
Assert.assertEquals(2110, testAmbryCUQuotaEnforcer.getQuotaSource().getUsage(QuotaResource.fromRestRequest(restRequest), QuotaName.READ_CAPACITY_UNIT), 0.0001);
Assert.assertEquals(211, testAmbryCUQuotaEnforcer.getQuotaSource().getSystemResourceUsage(QuotaName.READ_CAPACITY_UNIT), 0.0001);
}
use of com.github.ambry.rest.RestRequest in project ambry by linkedin.
the class AmbryQuotaManagerTest method testChargeAndRecommendForNoExceedNoForce.
/**
* Test {@link QuotaManager#chargeAndRecommend} when quota exceed check and force charge are set to false.
* @throws Exception in case of any exception.
*/
private void testChargeAndRecommendForNoExceedNoForce() throws Exception {
RestRequest restRequest = QuotaTestUtils.createRestRequest(ACCOUNT, ACCOUNT.getAllContainers().iterator().next(), RestMethod.GET);
AmbryQuotaManager ambryQuotaManager = buildAmbryQuotaManagerWithNoEnforcers();
ambryQuotaManager.init();
// 1. test that chargeAndRecommend returns ALLOW with not enforcers.
QuotaAction quotaAction = ambryQuotaManager.chargeAndRecommend(restRequest, null, false, false);
Assert.assertEquals(QuotaAction.ALLOW, quotaAction);
// 2. test that enforcers are called.
ambryQuotaManager = buildAmbryQuotaManagerWithAmbryCUEnforcer();
testAmbryCUQuotaEnforcer.resetDefaults();
ambryQuotaManager.init();
Map<QuotaName, Double> costMap = Collections.singletonMap(QuotaName.READ_CAPACITY_UNIT, 11.0);
quotaAction = ambryQuotaManager.chargeAndRecommend(restRequest, costMap, false, false);
Assert.assertEquals(QuotaAction.ALLOW, quotaAction);
Assert.assertEquals(1, testAmbryCUQuotaEnforcer.initCallCount);
Assert.assertEquals(1, testAmbryCUQuotaEnforcer.recommendCallCount);
Assert.assertEquals(1, testAmbryCUQuotaEnforcer.chargeCallCount);
Assert.assertEquals(110, testAmbryCUQuotaEnforcer.getQuotaSource().getUsage(QuotaResource.fromRestRequest(restRequest), QuotaName.READ_CAPACITY_UNIT), 0.0001);
// 3. test that if one enforcer recommends delay then no charge is done and delay is returned.
testAmbryCUQuotaEnforcer.resetDefaults();
ambryQuotaManager.init();
costMap = Collections.singletonMap(QuotaName.READ_CAPACITY_UNIT, 1.0);
quotaAction = ambryQuotaManager.chargeAndRecommend(restRequest, costMap, false, false);
Assert.assertEquals(QuotaAction.DELAY, quotaAction);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer.initCallCount);
Assert.assertEquals(1, testAmbryCUQuotaEnforcer.recommendCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer.chargeCallCount);
Assert.assertEquals(110, testAmbryCUQuotaEnforcer.getQuotaSource().getUsage(QuotaResource.fromRestRequest(restRequest), QuotaName.READ_CAPACITY_UNIT), 0.0001);
// 4. test that if one enforcer recommends reject then no charge is done and reject is returned.
ambryQuotaManager = buildAmbryQuotaManager();
List<QuotaEnforcer> quotaEnforcers = new ArrayList<>(ambryQuotaManager.quotaEnforcers);
TestCUQuotaEnforcerFactory.TestAmbryCUQuotaEnforcer testAmbryCUQuotaEnforcer1 = (TestCUQuotaEnforcerFactory.TestAmbryCUQuotaEnforcer) quotaEnforcers.get(0);
TestCUQuotaEnforcerFactory.TestAmbryCUQuotaEnforcer testAmbryCUQuotaEnforcer2 = (TestCUQuotaEnforcerFactory.TestAmbryCUQuotaEnforcer) quotaEnforcers.get(1);
testAmbryCUQuotaEnforcer1.resetDefaults();
testAmbryCUQuotaEnforcer2.resetDefaults();
ambryQuotaManager.init();
testAmbryCUQuotaEnforcer1.recommendReturnVal = TestCUQuotaEnforcerFactory.TestAmbryCUQuotaEnforcer.REJECT_QUOTA_RECOMMENDATION;
quotaAction = ambryQuotaManager.chargeAndRecommend(restRequest, costMap, false, false);
Assert.assertEquals(QuotaAction.REJECT, quotaAction);
Assert.assertEquals(1, testAmbryCUQuotaEnforcer1.initCallCount);
Assert.assertEquals(1, testAmbryCUQuotaEnforcer1.recommendCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer1.chargeCallCount);
Assert.assertEquals(1, testAmbryCUQuotaEnforcer2.initCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer2.recommendCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer2.chargeCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer1.isQuotaExceedAllowedCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer2.isQuotaExceedAllowedCallCount);
// 4a. flip the enforcer that returns reject.
testAmbryCUQuotaEnforcer1.resetDefaults();
testAmbryCUQuotaEnforcer2.resetDefaults();
testAmbryCUQuotaEnforcer2.recommendReturnVal = TestCUQuotaEnforcerFactory.TestAmbryCUQuotaEnforcer.REJECT_QUOTA_RECOMMENDATION;
quotaAction = ambryQuotaManager.chargeAndRecommend(restRequest, costMap, false, false);
Assert.assertEquals(QuotaAction.REJECT, quotaAction);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer1.initCallCount);
Assert.assertEquals(1, testAmbryCUQuotaEnforcer1.recommendCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer1.chargeCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer2.initCallCount);
Assert.assertEquals(1, testAmbryCUQuotaEnforcer2.recommendCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer2.chargeCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer1.isQuotaExceedAllowedCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer2.isQuotaExceedAllowedCallCount);
// 5. test that if one enforcer recommends reject and another returns delay then no charge is done and reject is returned.
testAmbryCUQuotaEnforcer1.resetDefaults();
testAmbryCUQuotaEnforcer2.resetDefaults();
testAmbryCUQuotaEnforcer1.recommendReturnVal = TestCUQuotaEnforcerFactory.TestAmbryCUQuotaEnforcer.REJECT_QUOTA_RECOMMENDATION;
testAmbryCUQuotaEnforcer2.recommendReturnVal = TestCUQuotaEnforcerFactory.TestAmbryCUQuotaEnforcer.DELAY_QUOTA_RECOMMENDATION;
quotaAction = ambryQuotaManager.chargeAndRecommend(restRequest, costMap, false, false);
Assert.assertEquals(QuotaAction.REJECT, quotaAction);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer1.initCallCount);
Assert.assertEquals(1, testAmbryCUQuotaEnforcer1.recommendCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer1.chargeCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer2.initCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer2.recommendCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer2.chargeCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer1.isQuotaExceedAllowedCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer2.isQuotaExceedAllowedCallCount);
// 5a. flip the enforcer that returns reject.
testAmbryCUQuotaEnforcer1.resetDefaults();
testAmbryCUQuotaEnforcer2.resetDefaults();
testAmbryCUQuotaEnforcer1.recommendReturnVal = TestCUQuotaEnforcerFactory.TestAmbryCUQuotaEnforcer.DELAY_QUOTA_RECOMMENDATION;
testAmbryCUQuotaEnforcer2.recommendReturnVal = TestCUQuotaEnforcerFactory.TestAmbryCUQuotaEnforcer.REJECT_QUOTA_RECOMMENDATION;
quotaAction = ambryQuotaManager.chargeAndRecommend(restRequest, costMap, false, false);
Assert.assertEquals(QuotaAction.REJECT, quotaAction);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer1.initCallCount);
Assert.assertEquals(1, testAmbryCUQuotaEnforcer1.recommendCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer1.chargeCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer2.initCallCount);
Assert.assertEquals(1, testAmbryCUQuotaEnforcer2.recommendCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer2.chargeCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer1.isQuotaExceedAllowedCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer2.isQuotaExceedAllowedCallCount);
// 6. test that if any enforcer's recommend throws exception, the exception is thrown from chargeAndRecommend.
testAmbryCUQuotaEnforcer1.resetDefaults();
testAmbryCUQuotaEnforcer2.resetDefaults();
testAmbryCUQuotaEnforcer1.throwException = true;
try {
ambryQuotaManager.chargeAndRecommend(restRequest, costMap, false, false);
Assert.fail("If any enforcer's recommend throws exception, then chargeAndRecommend should throw exception.");
} catch (QuotaException ignored) {
}
Assert.assertEquals(0, testAmbryCUQuotaEnforcer1.initCallCount);
Assert.assertEquals(1, testAmbryCUQuotaEnforcer1.recommendCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer1.chargeCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer2.initCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer2.recommendCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer2.chargeCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer1.isQuotaExceedAllowedCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer2.isQuotaExceedAllowedCallCount);
// 6a. flip the enforcer that throws exception
testAmbryCUQuotaEnforcer1.resetDefaults();
testAmbryCUQuotaEnforcer2.resetDefaults();
testAmbryCUQuotaEnforcer2.throwException = true;
try {
ambryQuotaManager.chargeAndRecommend(restRequest, costMap, false, false);
Assert.fail("If any enforcer's recommend throws exception, then chargeAndRecommend should throw exception.");
} catch (QuotaException ignored) {
}
Assert.assertEquals(0, testAmbryCUQuotaEnforcer1.initCallCount);
Assert.assertEquals(1, testAmbryCUQuotaEnforcer1.recommendCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer1.chargeCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer2.initCallCount);
Assert.assertEquals(1, testAmbryCUQuotaEnforcer2.recommendCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer2.chargeCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer1.isQuotaExceedAllowedCallCount);
Assert.assertEquals(0, testAmbryCUQuotaEnforcer2.isQuotaExceedAllowedCallCount);
}
use of com.github.ambry.rest.RestRequest in project ambry by linkedin.
the class RejectingQuotaChargeCallback method charge.
@Override
public void charge(long chunkSize) throws QuotaException {
try {
Map<QuotaName, Double> requestCost = requestCostPolicy.calculateRequestQuotaCharge(restRequest, chunkSize).entrySet().stream().collect(Collectors.toMap(entry -> QuotaName.valueOf(entry.getKey()), Map.Entry::getValue));
QuotaAction quotaAction = quotaManager.chargeAndRecommend(restRequest, requestCost, false, true);
if (QuotaUtils.shouldThrottle(quotaAction) && isQuotaEnforcedOnRequest) {
if (quotaManager.getQuotaMode() == QuotaMode.THROTTLING && quotaManager.getQuotaConfig().throttleInProgressRequests) {
throw new QuotaException("Exception while charging quota", new RouterException("RequestQuotaExceeded", RouterErrorCode.TooManyRequests), false);
} else {
LOGGER.debug("Quota exceeded for an in progress request.");
}
}
} catch (Exception ex) {
if (ex.getCause() instanceof RouterException && ((RouterException) ex.getCause()).getErrorCode().equals(RouterErrorCode.TooManyRequests)) {
throw ex;
}
LOGGER.error("Unexpected exception while charging quota.", ex);
}
}
use of com.github.ambry.rest.RestRequest in project ambry by linkedin.
the class SimpleRequestQuotaCostPolicyTest method testCalculateRequestQuotaCharge.
@Test
public void testCalculateRequestQuotaCharge() throws Exception {
QuotaConfig quotaConfig = new QuotaConfig(new VerifiableProperties(new Properties()));
SimpleRequestQuotaCostPolicy quotaRequestCostPolicy = new SimpleRequestQuotaCostPolicy(quotaConfig);
RestResponseChannel restResponseChannel = mock(RestResponseChannel.class);
when(restResponseChannel.getHeader(anyString())).thenReturn(0);
String blobUri = "/AAYIAQSSAAgAAQAAAAAAABpFymbGwe7sRBWYa5OPlkcNHQ.bin";
// test for a 4 MB GET request.
long blobSize = 4 * MB;
RestRequest restRequest = createMockRequestWithMethod(RestMethod.GET, blobUri, -1);
Map<String, Double> costMap = quotaRequestCostPolicy.calculateRequestQuotaCharge(restRequest, blobSize);
verifyReadCost(costMap, Math.ceil(blobSize / (double) quotaConfig.quotaAccountingUnit));
// test for a small GET request (fractional CU).
blobSize = 6 * MB;
costMap = quotaRequestCostPolicy.calculateRequestQuotaCharge(restRequest, blobSize);
verifyReadCost(costMap, Math.ceil(blobSize / (double) quotaConfig.quotaAccountingUnit));
// test for a GET request of blob of size 0.
restRequest = createMockRequestWithMethod(RestMethod.GET, blobUri, -1);
costMap = quotaRequestCostPolicy.calculateRequestQuotaCharge(restRequest, 0);
verifyReadCost(costMap, 1);
// test for a GET request of blob of size 512.
blobSize = 512;
restRequest = createMockRequestWithMethod(RestMethod.GET, blobUri, -1);
costMap = quotaRequestCostPolicy.calculateRequestQuotaCharge(restRequest, blobSize);
verifyReadCost(costMap, 1);
// test for a small POST request (fractional storage cost).
blobSize = 8 * MB;
restRequest = createMockRequestWithMethod(RestMethod.POST, blobUri, blobSize);
costMap = quotaRequestCostPolicy.calculateRequestQuotaCharge(restRequest, blobSize);
verifyWriteCost(costMap, Math.ceil(blobSize / (double) quotaConfig.quotaAccountingUnit), 8 * 1024 * 1024 / (double) QuotaUtils.BYTES_IN_GB);
// test for a large POST request.
blobSize = 4 * QuotaUtils.BYTES_IN_GB;
restRequest = createMockRequestWithMethod(RestMethod.POST, blobUri, blobSize);
costMap = quotaRequestCostPolicy.calculateRequestQuotaCharge(restRequest, blobSize);
verifyWriteCost(costMap, Math.ceil(blobSize / (double) quotaConfig.quotaAccountingUnit), 4);
// test for a POST request of blob of size 0.
restRequest = createMockRequestWithMethod(RestMethod.POST, blobUri, 0);
costMap = quotaRequestCostPolicy.calculateRequestQuotaCharge(restRequest, 0);
verifyWriteCost(costMap, 1, 0);
// test for a POST request of blob of size 512.
restRequest = createMockRequestWithMethod(RestMethod.POST, blobUri, 0);
costMap = quotaRequestCostPolicy.calculateRequestQuotaCharge(restRequest, 512);
verifyWriteCost(costMap, 1, 0);
// test for a HEAD request.
restRequest = createMockRequestWithMethod(RestMethod.HEAD, blobUri, -1);
costMap = quotaRequestCostPolicy.calculateRequestQuotaCharge(restRequest, quotaConfig.quotaAccountingUnit);
verifyReadCost(costMap, 1);
// test for a DELETE request.
restRequest = createMockRequestWithMethod(RestMethod.DELETE, blobUri, -1);
costMap = quotaRequestCostPolicy.calculateRequestQuotaCharge(restRequest, quotaConfig.quotaAccountingUnit);
verifyWriteCost(costMap, 1, 0.0);
// test for a PUT request.
restRequest = createMockRequestWithMethod(RestMethod.PUT, blobUri, -1);
costMap = quotaRequestCostPolicy.calculateRequestQuotaCharge(restRequest, quotaConfig.quotaAccountingUnit);
verifyWriteCost(costMap, 1, 0.0);
}
use of com.github.ambry.rest.RestRequest in project ambry by linkedin.
the class NonBlockingRouterTest method testRouterBasicForRegularBlob.
/**
* Test Router with a single scaling unit for regular blob operations.
*/
void testRouterBasicForRegularBlob() throws Exception {
// More extensive test for puts present elsewhere - these statements are here just to exercise the flow within the
// NonBlockingRouter class, and to ensure that operations submitted to a router eventually completes.
List<String> blobIds = new ArrayList<>();
for (int i = 0; i < 2; i++) {
setOperationParams();
RestRequest restRequest = createRestRequestForPutOperation();
PutBlobOptions putBlobOptions = new PutBlobOptionsBuilder().restRequest(restRequest).build();
String blobId = router.putBlob(putBlobProperties, putUserMetadata, putChannel, putBlobOptions).get();
if (testEncryption) {
Assert.assertEquals(restRequest, kms.getCurrentRestRequest());
}
logger.debug("Put blob {}", blobId);
blobIds.add(blobId);
}
for (String blobId : blobIds) {
RestRequest restRequest = createRestRequestForGetOperation();
GetBlobOptions getBlobOptions = new GetBlobOptionsBuilder().restRequest(restRequest).build();
router.getBlob(blobId, getBlobOptions).get();
if (testEncryption) {
Assert.assertEquals(restRequest, kms.getCurrentRestRequest());
}
router.updateBlobTtl(blobId, null, Utils.Infinite_Time).get();
router.getBlob(blobId, new GetBlobOptionsBuilder().build()).get();
router.getBlob(blobId, new GetBlobOptionsBuilder().operationType(GetBlobOptions.OperationType.BlobInfo).build()).get();
router.deleteBlob(blobId, null).get();
try {
router.getBlob(blobId, new GetBlobOptionsBuilder().build()).get();
fail("Get blob should fail");
} catch (ExecutionException e) {
RouterException r = (RouterException) e.getCause();
Assert.assertEquals("BlobDeleted error is expected", RouterErrorCode.BlobDeleted, r.getErrorCode());
}
router.getBlob(blobId, new GetBlobOptionsBuilder().getOption(GetOption.Include_Deleted_Blobs).build()).get();
router.getBlob(blobId, new GetBlobOptionsBuilder().getOption(GetOption.Include_All).build()).get();
}
}
Aggregations