Examples with RetainingAsyncWritableChannel com.github.ambry.commons.RetainingAsyncWritableChannel used on opensource projects

Example 1 with RetainingAsyncWritableChannel

use of com.github.ambry.commons.RetainingAsyncWritableChannel in project ambry by linkedin.

the class PostAccountContainersHandlerTest method validRequestsTest.

/**
 * Test valid request cases.
 * @throws Exception
 */
@Test
public void validRequestsTest() throws Exception {
    String accountName = theAccount.getName();
    short accountId = theAccount.getId();
    ThrowingConsumer<Collection<Container>> testAction = inputContainers -> {
        String requestBody = new String(AccountCollectionSerde.serializeContainersInJson(inputContainers));
        RestResponseChannel restResponseChannel = new MockRestResponseChannel();
        RestRequest request = createRestRequest(requestBody, accountName, null);
        ReadableStreamChannel responseChannel = sendRequestGetResponse(request, restResponseChannel);
        assertNotNull("Date has not been set", restResponseChannel.getHeader(RestUtils.Headers.DATE));
        assertEquals("Content-length is not as expected", responseChannel.getSize(), Integer.parseInt((String) restResponseChannel.getHeader(RestUtils.Headers.CONTENT_LENGTH)));
        assertEquals("Account id in response header is not as expected", accountId, Short.parseShort((String) restResponseChannel.getHeader(RestUtils.Headers.TARGET_ACCOUNT_ID)));
        RetainingAsyncWritableChannel asyncWritableChannel = new RetainingAsyncWritableChannel((int) responseChannel.getSize());
        responseChannel.readInto(asyncWritableChannel, null).get();
        Collection<Container> outputContainers = AccountCollectionSerde.containersFromInputStreamInJson(asyncWritableChannel.consumeContentAsInputStream(), accountId);
        assertEquals("Unexpected count returned", inputContainers.size(), outputContainers.size());
        for (Container container : outputContainers) {
            assertEquals("Container in account service not as expected", container, accountService.getContainerByName(accountName, container.getName()));
        }
    };
    // add new container
    testAction.accept(Collections.singleton(accountService.getRandomContainer(accountId)));
    // add multiple containers
    List<Container> containerList = new ArrayList<>();
    for (int j = 0; j < 10; j++) {
        containerList.add(new ContainerBuilder(Container.UNKNOWN_CONTAINER_ID, "Test-" + j, Container.ContainerStatus.ACTIVE, "", accountId).build());
    }
    testAction.accept(containerList);
// TODO: update existing containers when support is added
}

Example 2 with RetainingAsyncWritableChannel

use of com.github.ambry.commons.RetainingAsyncWritableChannel in project ambry by linkedin.

the class GetAccountsHandlerTest method getSingleContainerSuccessTest.

/**
 * Test success case of getting single container.
 * @throws Exception
 */
@Test
public void getSingleContainerSuccessTest() throws Exception {
    Account existingAccount = accountService.createAndAddRandomAccount();
    Container existingContainer = existingAccount.getAllContainers().iterator().next();
    ThrowingBiConsumer<RestRequest, Container> testAction = (request, expectedContainer) -> {
        RestResponseChannel restResponseChannel = new MockRestResponseChannel();
        ReadableStreamChannel channel = sendRequestGetResponse(request, restResponseChannel);
        assertNotNull("There should be a response", channel);
        Assert.assertNotNull("Date has not been set", restResponseChannel.getHeader(RestUtils.Headers.DATE));
        assertEquals("Content-type is not as expected", RestUtils.JSON_CONTENT_TYPE, restResponseChannel.getHeader(RestUtils.Headers.CONTENT_TYPE));
        assertEquals("Content-length is not as expected", channel.getSize(), Integer.parseInt((String) restResponseChannel.getHeader(RestUtils.Headers.CONTENT_LENGTH)));
        RetainingAsyncWritableChannel asyncWritableChannel = new RetainingAsyncWritableChannel((int) channel.getSize());
        channel.readInto(asyncWritableChannel, null).get();
        assertEquals("Container does not match", Collections.singletonList(expectedContainer), AccountCollectionSerde.containersFromInputStreamInJson(asyncWritableChannel.consumeContentAsInputStream(), existingAccount.getId()));
    };
    testAction.accept(createRestRequest(existingAccount.getName(), null, existingContainer.getName(), Operations.ACCOUNTS_CONTAINERS), existingContainer);
}

Example 3 with RetainingAsyncWritableChannel

use of com.github.ambry.commons.RetainingAsyncWritableChannel in project ambry by linkedin.

the class RestTestUtils method getJsonizedResponseBody.

/**
 * Reads the response received from the {@code channel} and decodes it into a {@link JSONObject}.
 * @param channel the {@link ReadableStreamChannel} that contains the response
 * @return the response decoded into a {@link JSONObject}.
 * @throws Exception
 */
public static JSONObject getJsonizedResponseBody(ReadableStreamChannel channel) throws Exception {
    RetainingAsyncWritableChannel asyncWritableChannel = new RetainingAsyncWritableChannel((int) channel.getSize());
    channel.readInto(asyncWritableChannel, null).get();
    try (InputStream is = asyncWritableChannel.consumeContentAsInputStream()) {
        return new JSONObject(new String(Utils.readBytesFromStream(is, (int) asyncWritableChannel.getBytesWritten())));
    }
}

Example 4 with RetainingAsyncWritableChannel

use of com.github.ambry.commons.RetainingAsyncWritableChannel in project ambry by linkedin.

the class RestTestUtils method getResponseBody.

/**
 * Reads the response received from the {@code channel}.
 * @param channel the {@link ReadableStreamChannel} that contains the response
 * @return the response in byte array.
 * @throws Exception
 */
public static byte[] getResponseBody(ReadableStreamChannel channel) throws Exception {
    RetainingAsyncWritableChannel asyncWritableChannel = new RetainingAsyncWritableChannel((int) channel.getSize());
    channel.readInto(asyncWritableChannel, null).get();
    byte[] result = new byte[(int) asyncWritableChannel.getBytesWritten()];
    asyncWritableChannel.consumeContentAsByteBuf().readBytes(result);
    return result;
}

Example 5 with RetainingAsyncWritableChannel

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