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

Example 1 with QuotaException

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

the class AmbrySecurityService method processRequestCharges.

@Override
public void processRequestCharges(RestRequest restRequest, RestResponseChannel responseChannel, BlobInfo blobInfo) {
    Map<String, Double> requestCost = ambryCostModelPolicy.calculateRequestCost(restRequest, responseChannel, blobInfo);
    setRequestCostHeader(requestCost, responseChannel);
    if (QuotaUtils.isRequestResourceQuotaManaged(restRequest) && quotaManager != null) {
        ThrottlingRecommendation throttlingRecommendation = null;
        try {
            throttlingRecommendation = quotaManager.recommend(restRequest);
        } catch (QuotaException quotaException) {
            LOGGER.warn("Exception {} in while processing request charges.", quotaException);
        }
        if (throttlingRecommendation != null) {
            RestUtils.buildUserQuotaHeadersMap(throttlingRecommendation).entrySet().stream().forEach(headerEntry -> responseChannel.setHeader(headerEntry.getKey(), headerEntry.getValue()));
        }
    }
}

Example 2 with QuotaException

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

the class AmbryCUQuotaEnforcerTest method testCharge.

@Test
public void testCharge() throws Exception {
    Container container = ACCOUNT.getAllContainers().iterator().next();
    Map<QuotaName, Double> readRequestCostMap = Collections.singletonMap(QuotaName.READ_CAPACITY_UNIT, 9.0);
    Map<QuotaName, Double> writeRequestCostMap = Collections.singletonMap(QuotaName.WRITE_CAPACITY_UNIT, 9.0);
    Map<String, CapacityUnit> usageMap = QUOTA_SOURCE.getAllQuotaUsage();
    // 1. Test that usage is updated and recommendation is serve when usage is within limit.
    RestRequest restRequest = QuotaTestUtils.createRestRequest(ACCOUNT, container, RestMethod.GET);
    AMBRY_QUOTA_ENFORCER.charge(restRequest, readRequestCostMap);
    QuotaRecommendation quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(restRequest);
    assertEquals(quotaRecommendation.getQuotaName(), QuotaName.READ_CAPACITY_UNIT);
    assertEquals(quotaRecommendation.getQuotaUsagePercentage(), 90, 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.ALLOW);
    // make sure that correct quota is charged.
    assertEquals(usageMap.get(String.valueOf(ACCOUNT.getId())).getWcu(), 0);
    restRequest = QuotaTestUtils.createRestRequest(ACCOUNT, container, RestMethod.POST);
    AMBRY_QUOTA_ENFORCER.charge(restRequest, writeRequestCostMap);
    quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(restRequest);
    assertEquals(quotaRecommendation.getQuotaName(), QuotaName.WRITE_CAPACITY_UNIT);
    assertEquals(quotaRecommendation.getQuotaUsagePercentage(), 90, 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.ALLOW);
    // make sure that correct quota is charged.
    assertEquals(usageMap.get(String.valueOf(ACCOUNT.getId())).getRcu(), 9);
    // 2. Test that retryable quota exception is thrown when quota not found.
    Account newAccount = ACCOUNT_SERVICE.generateRandomAccount(QuotaResourceType.ACCOUNT);
    restRequest = QuotaTestUtils.createRestRequest(newAccount, newAccount.getAllContainers().iterator().next(), RestMethod.GET);
    try {
        AMBRY_QUOTA_ENFORCER.charge(restRequest, readRequestCostMap);
        AMBRY_QUOTA_ENFORCER.recommend(restRequest);
        fail("if quota is not found we should see exception");
    } catch (QuotaException quotaException) {
        Assert.assertTrue(quotaException.isRetryable());
    }
    // 3. Test that retryable quota exception is thrown in case of any error.
    restRequest = QuotaTestUtils.createRestRequest(ACCOUNT, container, RestMethod.GET);
    QUOTA_SOURCE.throwException = true;
    try {
        AMBRY_QUOTA_ENFORCER.charge(restRequest, readRequestCostMap);
        fail("QuotaException should be thrown in case of any error.");
    } catch (QuotaException quotaException) {
        Assert.assertTrue(quotaException.isRetryable());
    }
    QUOTA_SOURCE.throwException = false;
    // 4. Test that usage is updated and recommendation is deny when usage >= quota.
    AMBRY_QUOTA_ENFORCER.charge(restRequest, readRequestCostMap);
    quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(restRequest);
    assertEquals(QuotaName.READ_CAPACITY_UNIT, quotaRecommendation.getQuotaName());
    assertEquals(180, quotaRecommendation.getQuotaUsagePercentage(), 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.DELAY);
    // make sure that correct quota is charged.
    assertEquals(9, usageMap.get(String.valueOf(ACCOUNT.getId())).getWcu());
    assertEquals(18, // make sure that correct quota is charged.
    usageMap.get(String.valueOf(ACCOUNT.getId())).getRcu());
    restRequest = QuotaTestUtils.createRestRequest(ACCOUNT, container, RestMethod.POST);
    AMBRY_QUOTA_ENFORCER.charge(restRequest, writeRequestCostMap);
    quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(restRequest);
    assertEquals(QuotaName.WRITE_CAPACITY_UNIT, quotaRecommendation.getQuotaName());
    assertEquals(180, quotaRecommendation.getQuotaUsagePercentage(), 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.DELAY);
    assertEquals(18, // make sure that correct quota is charged.
    usageMap.get(String.valueOf(ACCOUNT.getId())).getWcu());
    assertEquals(18, // make sure that correct quota is charged.
    usageMap.get(String.valueOf(ACCOUNT.getId())).getRcu());
    restRequest = QuotaTestUtils.createRestRequest(ACCOUNT, container, RestMethod.DELETE);
    AMBRY_QUOTA_ENFORCER.charge(restRequest, writeRequestCostMap);
    quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(restRequest);
    assertEquals(QuotaName.WRITE_CAPACITY_UNIT, quotaRecommendation.getQuotaName());
    assertEquals(270, quotaRecommendation.getQuotaUsagePercentage(), 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.DELAY);
    assertEquals(27, // make sure that correct quota is charged.
    usageMap.get(String.valueOf(ACCOUNT.getId())).getWcu());
    assertEquals(18, // make sure that correct quota is charged.
    usageMap.get(String.valueOf(ACCOUNT.getId())).getRcu());
    restRequest = QuotaTestUtils.createRestRequest(ACCOUNT, container, RestMethod.PUT);
    AMBRY_QUOTA_ENFORCER.charge(restRequest, writeRequestCostMap);
    quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(restRequest);
    assertEquals(QuotaName.WRITE_CAPACITY_UNIT, quotaRecommendation.getQuotaName());
    assertEquals(360, quotaRecommendation.getQuotaUsagePercentage(), 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.DELAY);
    assertEquals(36, // make sure that correct quota is charged.
    usageMap.get(String.valueOf(ACCOUNT.getId())).getWcu());
    assertEquals(18, // make sure that correct quota is charged.
    usageMap.get(String.valueOf(ACCOUNT.getId())).getRcu());
}

Example 3 with QuotaException

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

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

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

the class AmbryCUQuotaEnforcerTest method testRecommend.

@Test
public void testRecommend() throws Exception {
    // 1. Test that recommendation is serve when usage is within limit and correct quota is used based on rest method.
    QuotaRecommendation quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(QuotaTestUtils.createRestRequest(ACCOUNT, ACCOUNT.getAllContainers().iterator().next(), RestMethod.GET));
    assertEquals(QuotaName.READ_CAPACITY_UNIT, quotaRecommendation.getQuotaName());
    assertEquals(0, quotaRecommendation.getQuotaUsagePercentage(), 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.ALLOW);
    quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(QuotaTestUtils.createRestRequest(ACCOUNT, ACCOUNT.getAllContainers().iterator().next(), RestMethod.POST));
    assertEquals(QuotaName.WRITE_CAPACITY_UNIT, quotaRecommendation.getQuotaName());
    assertEquals(0, quotaRecommendation.getQuotaUsagePercentage(), 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.ALLOW);
    quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(QuotaTestUtils.createRestRequest(ACCOUNT, ACCOUNT.getAllContainers().iterator().next(), RestMethod.PUT));
    assertEquals(QuotaName.WRITE_CAPACITY_UNIT, quotaRecommendation.getQuotaName());
    assertEquals(0, quotaRecommendation.getQuotaUsagePercentage(), 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.ALLOW);
    quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(QuotaTestUtils.createRestRequest(ACCOUNT, ACCOUNT.getAllContainers().iterator().next(), RestMethod.DELETE));
    assertEquals(QuotaName.WRITE_CAPACITY_UNIT, quotaRecommendation.getQuotaName());
    assertEquals(0, quotaRecommendation.getQuotaUsagePercentage(), 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.ALLOW);
    // 2. Test that retryable QuotaException is thrown when quota not found.
    Account newAccount = ACCOUNT_SERVICE.generateRandomAccount(QuotaResourceType.ACCOUNT);
    try {
        AMBRY_QUOTA_ENFORCER.recommend(QuotaTestUtils.createRestRequest(newAccount, newAccount.getAllContainers().iterator().next(), RestMethod.GET));
        fail("QuotaException should be thrown when quota not found");
    } catch (QuotaException quotaException) {
        Assert.assertTrue(quotaException.isRetryable());
    }
    // 3. Test that recommendation is serve in case of any error.
    QUOTA_SOURCE.throwException = true;
    try {
        AMBRY_QUOTA_ENFORCER.recommend(QuotaTestUtils.createRestRequest(ACCOUNT, ACCOUNT.getAllContainers().iterator().next(), RestMethod.GET));
        fail("QuotaException should be thrown.");
    } catch (QuotaException quotaException) {
        Assert.assertTrue(quotaException.isRetryable());
    }
    QUOTA_SOURCE.throwException = false;
    // 4. Test that recommendation is deny when usage >= quota.
    Map<String, CapacityUnit> usageMap = QUOTA_SOURCE.getAllQuotaUsage();
    Map<String, CapacityUnit> quotaMap = QUOTA_SOURCE.getAllQuota();
    String id = String.valueOf(ACCOUNT.getId());
    usageMap.put(id, new CapacityUnit(quotaMap.get(id).getRcu() + 1, quotaMap.get(id).getWcu()));
    float usagePercentage = (float) (usageMap.get(id).getRcu() * 100) / quotaMap.get(id).getRcu();
    quotaRecommendation = AMBRY_QUOTA_ENFORCER.recommend(QuotaTestUtils.createRestRequest(ACCOUNT, ACCOUNT.getAllContainers().iterator().next(), RestMethod.GET));
    assertEquals(QuotaName.READ_CAPACITY_UNIT, quotaRecommendation.getQuotaName());
    assertEquals(usagePercentage, quotaRecommendation.getQuotaUsagePercentage(), 0.1);
    assertEquals(quotaRecommendation.getQuotaAction(), QuotaAction.DELAY);
}