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

Example 6 with RetainingAsyncWritableChannel

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

the class GetAccountsHandlerTest method validRequestsTest.

/**
 * Test valid request cases.
 * @throws Exception
 */
@Test
public void validRequestsTest() throws Exception {
    Account account = accountService.createAndAddRandomAccount();
    ThrowingBiConsumer<RestRequest, Collection<Account>> testAction = (request, expectedAccounts) -> {
        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("Accounts do not match", new HashSet<>(expectedAccounts), new HashSet<>(AccountCollectionSerde.accountsFromInputStreamInJson(asyncWritableChannel.consumeContentAsInputStream())));
    };
    testAction.accept(createRestRequest(null, null, null, Operations.ACCOUNTS), accountService.getAllAccounts());
    testAction.accept(createRestRequest(account.getName(), null, null, Operations.ACCOUNTS), Collections.singleton(account));
    testAction.accept(createRestRequest(null, Short.toString(account.getId()), null, Operations.ACCOUNTS), Collections.singleton(account));
}

Example 7 with RetainingAsyncWritableChannel

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

the class NettyMultipartRequestTest method doMultipartDecodeTest.

// multipartRequestDecodeTest() helpers.
/**
 * Does a multipart decode test.
 * 1. Creates a {@link NettyMultipartRequest}.
 * 2. Adds the {@link HttpRequest} and {@link HttpContent} generated by encoding the {@code files} as a multipart
 * request to the {@link NettyMultipartRequest}.
 * 3. Reads data from the {@link NettyMultipartRequest} via read operations and
 * {@link NettyMultipartRequest#getArgs()} and verifies them against the source data ({@code files}).
 * @param expectedRequestSize the value expected on a call to {@link NettyMultipartRequest#getSize()}.
 * @param files the {@link InMemoryFile}s that form the parts of the multipart request.
 * @param digestAlgorithm the digest algorithm to use. Can be empty or {@code null} if digest checking is not
 *                        required.
 * @throws Exception
 */
private void doMultipartDecodeTest(int expectedRequestSize, InMemoryFile[] files, String digestAlgorithm) throws Exception {
    HttpHeaders httpHeaders = new DefaultHttpHeaders();
    httpHeaders.set(RestUtils.Headers.BLOB_SIZE, expectedRequestSize);
    NettyMultipartRequest request = createRequest(httpHeaders, files);
    assertEquals("Request size does not match", expectedRequestSize, request.getSize());
    request.prepare();
    RetainingAsyncWritableChannel asyncWritableChannel;
    byte[] readOutput;
    Map<String, Object> args = request.getArgs();
    ByteBuffer blobData = ByteBuffer.allocate(0);
    byte[] wholeDigest = null;
    if (files != null) {
        for (InMemoryFile file : files) {
            if (file.name.equals(RestUtils.MultipartPost.BLOB_PART)) {
                blobData = file.content;
                if (digestAlgorithm != null && !digestAlgorithm.isEmpty()) {
                    MessageDigest digest = MessageDigest.getInstance(digestAlgorithm);
                    digest.update(blobData);
                    wholeDigest = digest.digest();
                    blobData.rewind();
                    request.setDigestAlgorithm(digestAlgorithm);
                }
            } else {
                Object value = args.get(file.name);
                assertNotNull("Request does not contain " + file, value);
                assertTrue("Argument value is not ByteBuffer", value instanceof ByteBuffer);
                readOutput = new byte[((ByteBuffer) value).remaining()];
                ((ByteBuffer) value).get(readOutput);
                assertArrayEquals(file.name + " content does not match", file.content.array(), readOutput);
            }
        }
    }
    asyncWritableChannel = new RetainingAsyncWritableChannel(expectedRequestSize);
    request.readInto(asyncWritableChannel, null).get();
    try (InputStream is = asyncWritableChannel.consumeContentAsInputStream()) {
        readOutput = Utils.readBytesFromStream(is, (int) asyncWritableChannel.getBytesWritten());
    }
    assertArrayEquals(RestUtils.MultipartPost.BLOB_PART + " content does not match", blobData.array(), readOutput);
    assertArrayEquals("Part by part digest should match digest of whole", wholeDigest, request.getDigest());
    closeRequestAndValidate(request);
}

Example 8 with RetainingAsyncWritableChannel

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

Example 9 with RetainingAsyncWritableChannel

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

the class ServerTestUtil method checkBlobId.

private static void checkBlobId(Router router, BlobId blobId, byte[] data) throws Exception {
    GetBlobResult result = router.getBlob(blobId.getID(), new GetBlobOptionsBuilder().build()).get(20, TimeUnit.SECONDS);
    ReadableStreamChannel blob = result.getBlobDataChannel();
    assertEquals("Size does not match that of data", data.length, result.getBlobInfo().getBlobProperties().getBlobSize());
    RetainingAsyncWritableChannel channel = new RetainingAsyncWritableChannel();
    blob.readInto(channel, null).get(1, TimeUnit.SECONDS);
    try (InputStream is = channel.consumeContentAsInputStream()) {
        assertArrayEquals(data, Utils.readBytesFromStream(is, (int) channel.getBytesWritten()));
    }
}