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

Example 16 with ByteBufferAsyncWritableChannel

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

the class CopyForcingByteBuf method doContentAddAndReadTest.

/**
 * Does the content addition and read verification based on the arguments provided.
 * @param digestAlgorithm the digest algorithm to use. Can be empty or {@code null} if digest checking is not
 *                        required.
 * @param content the complete content.
 * @param httpContents {@code content} in parts and as {@link HttpContent}. Should contain all the data in
 * {@code content}.
 * @param numChunksToAddBeforeRead the number of {@link HttpContent} to add before making the
 * {@link NettyRequest#readInto(AsyncWritableChannel, Callback)} call.
 * @param method Http method
 * @throws Exception
 */
private void doContentAddAndReadTest(String digestAlgorithm, ByteBuffer content, List<HttpContent> httpContents, int numChunksToAddBeforeRead, HttpMethod method) throws Exception {
    if (numChunksToAddBeforeRead < 0 || numChunksToAddBeforeRead > httpContents.size()) {
        throw new IllegalArgumentException("Illegal value of numChunksToAddBeforeRead");
    }
    Channel channel = new MockChannel();
    NettyRequest nettyRequest = createNettyRequest(method, "/", null, channel);
    byte[] wholeDigest = null;
    if (digestAlgorithm != null && !digestAlgorithm.isEmpty()) {
        MessageDigest digest = MessageDigest.getInstance(digestAlgorithm);
        digest.update(content);
        wholeDigest = digest.digest();
        content.rewind();
        nettyRequest.setDigestAlgorithm(digestAlgorithm);
    }
    int bytesToVerify = 0;
    int addedCount = 0;
    for (; addedCount < numChunksToAddBeforeRead; addedCount++) {
        HttpContent httpContent = httpContents.get(addedCount);
        bytesToVerify += httpContent.content().readableBytes();
        nettyRequest.addContent(httpContent);
        // ref count always 2 when added before calling readInto()
        assertEquals("Reference count is not as expected", 2, httpContent.refCnt());
    }
    ByteBufferAsyncWritableChannel writeChannel = new ByteBufferAsyncWritableChannel();
    ReadIntoCallback callback = new ReadIntoCallback();
    Future<Long> future = nettyRequest.readInto(writeChannel, callback);
    readAndVerify(bytesToVerify, writeChannel, content);
    bytesToVerify = 0;
    for (; addedCount < httpContents.size(); addedCount++) {
        HttpContent httpContent = httpContents.get(addedCount);
        bytesToVerify += httpContent.content().readableBytes();
        nettyRequest.addContent(httpContent);
        assertEquals("Reference count is not as expected", 2, httpContent.refCnt());
    }
    readAndVerify(bytesToVerify, writeChannel, content);
    verifyRefCnts(httpContents);
    writeChannel.close();
    callback.awaitCallback();
    if (callback.exception != null) {
        throw callback.exception;
    }
    long futureBytesRead = future.get();
    assertEquals("Total bytes read does not match (callback)", content.limit(), callback.bytesRead);
    assertEquals("Total bytes read does not match (future)", content.limit(), futureBytesRead);
    assertEquals("nettyRequest.getBytesReceived() does not match expected", content.limit(), nettyRequest.getBytesReceived());
    // check twice to make sure the same digest is returned every time
    for (int i = 0; i < 2; i++) {
        assertArrayEquals("Part by part digest should match digest of whole", wholeDigest, nettyRequest.getDigest());
    }
    closeRequestAndValidate(nettyRequest, channel);
}

Example 17 with ByteBufferAsyncWritableChannel

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

the class CopyForcingByteBuf method getDigestWithoutSettingAlgorithmTest.

/**
 * Tests for failure when {@link NettyRequest#getDigest()} is called without a call to
 * {@link NettyRequest#setDigestAlgorithm(String)}.
 * @throws RestServiceException
 */
private void getDigestWithoutSettingAlgorithmTest() throws RestServiceException {
    List<HttpContent> httpContents = new ArrayList<HttpContent>();
    generateContent(httpContents);
    Channel channel = new MockChannel();
    NettyRequest nettyRequest = createNettyRequest(HttpMethod.POST, "/", null, channel);
    ByteBufferAsyncWritableChannel writeChannel = new ByteBufferAsyncWritableChannel();
    ReadIntoCallback callback = new ReadIntoCallback();
    nettyRequest.readInto(writeChannel, callback);
    for (HttpContent httpContent : httpContents) {
        nettyRequest.addContent(httpContent);
    }
    assertNull("Digest should be null because no digest algorithm was set", nettyRequest.getDigest());
    closeRequestAndValidate(nettyRequest, channel);
}

Example 18 with ByteBufferAsyncWritableChannel

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

the class CopyForcingByteBuf method getDigestBeforeAllContentProcessedTest.

/**
 * Tests for failure when {@link NettyRequest#getDigest()} is called before
 * 1. All content is added.
 * 2. All content is processed (i.e. before a call to {@link NettyRequest#readInto(AsyncWritableChannel, Callback)}).
 * @throws NoSuchAlgorithmException
 * @throws RestServiceException
 */
private void getDigestBeforeAllContentProcessedTest() throws NoSuchAlgorithmException, RestServiceException {
    // all content not added test.
    List<HttpContent> httpContents = new ArrayList<HttpContent>();
    generateContent(httpContents);
    Channel channel = new MockChannel();
    NettyRequest nettyRequest = createNettyRequest(HttpMethod.POST, "/", null, channel);
    nettyRequest.setDigestAlgorithm("MD5");
    // add all except the LastHttpContent
    for (int i = 0; i < httpContents.size() - 1; i++) {
        nettyRequest.addContent(httpContents.get(i));
    }
    ByteBufferAsyncWritableChannel writeChannel = new ByteBufferAsyncWritableChannel();
    ReadIntoCallback callback = new ReadIntoCallback();
    nettyRequest.readInto(writeChannel, callback);
    try {
        nettyRequest.getDigest();
        fail("Getting a digest should have failed because all the content has not been added");
    } catch (IllegalStateException e) {
    // expected. Nothing to do.
    }
    closeRequestAndValidate(nettyRequest, channel);
    // content not processed test.
    httpContents.clear();
    generateContent(httpContents);
    nettyRequest = createNettyRequest(HttpMethod.POST, "/", null, channel);
    nettyRequest.setDigestAlgorithm("MD5");
    for (HttpContent httpContent : httpContents) {
        nettyRequest.addContent(httpContent);
    }
    try {
        nettyRequest.getDigest();
        fail("Getting a digest should have failed because the content has not been processed (readInto() not called)");
    } catch (IllegalStateException e) {
    // expected. Nothing to do.
    }
    closeRequestAndValidate(nettyRequest, channel);
}

Example 19 with ByteBufferAsyncWritableChannel

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

the class CloudOperationTest method getBlobAndAssertSuccess.

/**
 * Construct GetBlob operations with appropriate callbacks, then poll those operations until they complete,
 * and ensure that the whole blob data is read out and the contents match.
 * @param blobId id of the blob to get
 * @param expectedLifeVersion the expected lifeVersion from get operation.
 * @param expectedBlobSize the expected blob size
 * @param expectedBlobProperties  the expected {@link BlobProperties} for the blob.
 * @param expectedUserMetadata the expected user meta data
 * @param expectPutContent the expected blob content
 * @param options options of the get blob operation
 * @throws Exception Any unexpected exception
 */
private void getBlobAndAssertSuccess(final BlobId blobId, final short expectedLifeVersion, final int expectedBlobSize, final BlobProperties expectedBlobProperties, final byte[] expectedUserMetadata, final byte[] expectPutContent, final GetBlobOptionsInternal options) throws Exception {
    final CountDownLatch readCompleteLatch = new CountDownLatch(1);
    final AtomicLong readCompleteResult = new AtomicLong(0);
    // callback to compare the data
    Callback<GetBlobResultInternal> callback = (result, exception) -> {
        Assert.assertNull("Shouldn't have exception", exception);
        try {
            BlobInfo blobInfo;
            switch(options.getBlobOptions.getOperationType()) {
                case All:
                    Assert.assertFalse("not supposed to be raw mode", options.getBlobOptions.isRawMode());
                    blobInfo = result.getBlobResult.getBlobInfo();
                    Assert.assertTrue("Blob properties must be the same", RouterTestHelpers.arePersistedFieldsEquivalent(expectedBlobProperties, blobInfo.getBlobProperties()));
                    Assert.assertEquals("Blob size should in received blobProperties should be the same as actual", expectedBlobSize, blobInfo.getBlobProperties().getBlobSize());
                    Assert.assertArrayEquals("User metadata must be the same", expectedUserMetadata, blobInfo.getUserMetadata());
                    Assert.assertEquals("LifeVersion mismatch", expectedLifeVersion, blobInfo.getLifeVersion());
                    break;
                case Data:
                    Assert.assertNull("Unexpected blob info in operation result", result.getBlobResult.getBlobInfo());
                    break;
                case BlobInfo:
                    blobInfo = result.getBlobResult.getBlobInfo();
                    Assert.assertTrue("Blob properties must be the same", RouterTestHelpers.arePersistedFieldsEquivalent(expectedBlobProperties, blobInfo.getBlobProperties()));
                    Assert.assertEquals("Blob size should in received blobProperties should be the same as actual", expectedBlobSize, blobInfo.getBlobProperties().getBlobSize());
                    Assert.assertNull("Unexpected blob data in operation result", result.getBlobResult.getBlobDataChannel());
                    Assert.assertEquals("LifeVersion mismatch", expectedLifeVersion, blobInfo.getLifeVersion());
            }
        } catch (Throwable e) {
            Assert.fail("Shouldn't receive exception here");
        }
        if (options.getBlobOptions.getOperationType() != GetBlobOptions.OperationType.BlobInfo) {
            final ByteBufferAsyncWritableChannel asyncWritableChannel = new ByteBufferAsyncWritableChannel();
            Utils.newThread(() -> {
                Future<Long> readIntoFuture = result.getBlobResult.getBlobDataChannel().readInto(asyncWritableChannel, null);
                assertBlobReadSuccess(options.getBlobOptions, readIntoFuture, asyncWritableChannel, result.getBlobResult.getBlobDataChannel(), readCompleteLatch, readCompleteResult, expectedBlobSize, expectPutContent);
            }, false).start();
        } else {
            readCompleteLatch.countDown();
        }
    };
    // create GetBlobOperation
    final Map<Integer, GetOperation> correlationIdToGetOperation = new HashMap<>();
    final RequestRegistrationCallback<GetOperation> requestRegistrationCallback = new RequestRegistrationCallback<>(correlationIdToGetOperation);
    NonBlockingRouter.currentOperationsCount.incrementAndGet();
    GetBlobOperation op = new GetBlobOperation(routerConfig, routerMetrics, mockClusterMap, responseHandler, blobId, options, callback, routerCallback, blobIdFactory, null, null, null, time, false, null);
    requestRegistrationCallback.setRequestsToSend(new ArrayList<>());
    // Wait operation to complete
    while (!op.isOperationComplete()) {
        op.poll(requestRegistrationCallback);
        List<ResponseInfo> responses = sendAndWaitForResponses(requestRegistrationCallback.getRequestsToSend());
        for (ResponseInfo responseInfo : responses) {
            GetResponse getResponse = RouterUtils.extractResponseAndNotifyResponseHandler(responseHandler, routerMetrics, responseInfo, stream -> GetResponse.readFrom(stream, mockClusterMap), response -> {
                ServerErrorCode serverError = response.getError();
                if (serverError == ServerErrorCode.No_Error) {
                    serverError = response.getPartitionResponseInfoList().get(0).getErrorCode();
                }
                return serverError;
            });
            op.handleResponse(responseInfo, getResponse);
            responseInfo.release();
        }
    }
    readCompleteLatch.await();
    Assert.assertTrue("Operation should be complete at this time", op.isOperationComplete());
    // Ensure that a ChannelClosed exception is not set when the ReadableStreamChannel is closed correctly.
    Assert.assertNull("Callback operation exception should be null", op.getOperationException());
    if (options.getBlobOptions.getOperationType() != GetBlobOptions.OperationType.BlobInfo && !options.getBlobOptions.isRawMode() && !options.getChunkIdsOnly) {
        int sizeWritten = expectedBlobSize;
        if (options.getBlobOptions.getRange() != null) {
            ByteRange range = options.getBlobOptions.getRange().toResolvedByteRange(expectedBlobSize, options.getBlobOptions.resolveRangeOnEmptyBlob());
            sizeWritten = (int) range.getRangeSize();
        }
        Assert.assertEquals("Size read must equal size written", sizeWritten, readCompleteResult.get());
    }
}

Example 20 with ByteBufferAsyncWritableChannel

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

the class GetBlobOperationTest method getAndAssertSuccess.

/**
 * Construct GetBlob operations with appropriate callbacks, then poll those operations until they complete,
 * and ensure that the whole blob data is read out and the contents match.
 * @param getChunksBeforeRead {@code true} if all chunks should be cached by the router before reading from the
 *                            stream.
 * @param initiateReadBeforeChunkGet Whether readInto() should be initiated immediately before data chunks are
 *                                   fetched by the router to simulate chunk arrival delay.
 * @param expectedLifeVersion the expected lifeVersion from get operation.
 */
private void getAndAssertSuccess(final boolean getChunksBeforeRead, final boolean initiateReadBeforeChunkGet, short expectedLifeVersion) throws Exception {
    final CountDownLatch readCompleteLatch = new CountDownLatch(1);
    final AtomicReference<Throwable> readCompleteThrowable = new AtomicReference<>(null);
    final AtomicLong readCompleteResult = new AtomicLong(0);
    final AtomicReference<Exception> operationException = new AtomicReference<>(null);
    final int numChunks = ((blobSize + maxChunkSize - 1) / maxChunkSize) + (blobSize > maxChunkSize ? 1 : 0);
    mockNetworkClient.resetProcessedResponseCount();
    Callback<GetBlobResultInternal> callback = (result, exception) -> {
        if (exception != null) {
            operationException.set(exception);
            readCompleteLatch.countDown();
        } else {
            try {
                if (options.getChunkIdsOnly) {
                    Assert.assertNull("Unexpected blob result when getChunkIdsOnly", result.getBlobResult);
                    if (blobSize > maxChunkSize) {
                        // CompositeBlob
                        Assert.assertNotNull("CompositeBlob should return a list of blob ids when getting chunk ids", result.storeKeys);
                        Assert.assertEquals(result.storeKeys.size(), (blobSize + maxChunkSize - 1) / maxChunkSize);
                    } else {
                        // SimpleBlob
                        Assert.assertNull("Unexpected list of blob id when getChunkIdsOnly is true on a simple blob", result.storeKeys);
                    }
                    readCompleteLatch.countDown();
                    return;
                }
                BlobInfo blobInfo;
                switch(options.getBlobOptions.getOperationType()) {
                    case All:
                        if (!options.getBlobOptions.isRawMode()) {
                            blobInfo = result.getBlobResult.getBlobInfo();
                            Assert.assertTrue("Blob properties must be the same", RouterTestHelpers.arePersistedFieldsEquivalent(blobProperties, blobInfo.getBlobProperties()));
                            Assert.assertEquals("Blob size should in received blobProperties should be the same as actual", blobSize, blobInfo.getBlobProperties().getBlobSize());
                            Assert.assertArrayEquals("User metadata must be the same", userMetadata, blobInfo.getUserMetadata());
                            Assert.assertEquals("LifeVersion mismatch", expectedLifeVersion, blobInfo.getLifeVersion());
                        }
                        break;
                    case Data:
                        Assert.assertNull("Unexpected blob info in operation result", result.getBlobResult.getBlobInfo());
                        break;
                    case BlobInfo:
                        blobInfo = result.getBlobResult.getBlobInfo();
                        Assert.assertTrue("Blob properties must be the same", RouterTestHelpers.arePersistedFieldsEquivalent(blobProperties, blobInfo.getBlobProperties()));
                        Assert.assertEquals("Blob size should in received blobProperties should be the same as actual", blobSize, blobInfo.getBlobProperties().getBlobSize());
                        Assert.assertNull("Unexpected blob data in operation result", result.getBlobResult.getBlobDataChannel());
                        Assert.assertEquals("LifeVersion mismatch", expectedLifeVersion, blobInfo.getLifeVersion());
                }
            } catch (Throwable e) {
                readCompleteThrowable.set(e);
            }
            if (options.getBlobOptions.getOperationType() != GetBlobOptions.OperationType.BlobInfo) {
                final ByteBufferAsyncWritableChannel asyncWritableChannel = new ByteBufferAsyncWritableChannel();
                final Future<Long> preSetReadIntoFuture = initiateReadBeforeChunkGet ? result.getBlobResult.getBlobDataChannel().readInto(asyncWritableChannel, null) : null;
                Utils.newThread(() -> {
                    if (getChunksBeforeRead) {
                        // wait for all chunks (data + metadata) to be received
                        while (mockNetworkClient.getProcessedResponseCount() < numChunks * routerConfig.routerGetRequestParallelism) {
                            Thread.yield();
                        }
                    }
                    Future<Long> readIntoFuture = initiateReadBeforeChunkGet ? preSetReadIntoFuture : result.getBlobResult.getBlobDataChannel().readInto(asyncWritableChannel, null);
                    assertBlobReadSuccess(options.getBlobOptions, readIntoFuture, asyncWritableChannel, result.getBlobResult.getBlobDataChannel(), readCompleteLatch, readCompleteResult, readCompleteThrowable);
                }, false).start();
            } else {
                readCompleteLatch.countDown();
            }
        }
    };
    GetBlobOperation op = createOperationAndComplete(callback);
    readCompleteLatch.await();
    Assert.assertTrue("Operation should be complete at this time", op.isOperationComplete());
    if (operationException.get() != null) {
        throw operationException.get();
    }
    if (readCompleteThrowable.get() != null) {
        throw new IllegalStateException(readCompleteThrowable.get());
    }
    // Ensure that a ChannelClosed exception is not set when the ReadableStreamChannel is closed correctly.
    Assert.assertNull("Callback operation exception should be null", op.getOperationException());
    if (options.getBlobOptions.getOperationType() != GetBlobOptions.OperationType.BlobInfo && !options.getBlobOptions.isRawMode() && !options.getChunkIdsOnly) {
        int sizeWritten = blobSize;
        if (options.getBlobOptions.getRange() != null) {
            ByteRange range = options.getBlobOptions.getRange().toResolvedByteRange(blobSize, options.getBlobOptions.resolveRangeOnEmptyBlob());
            sizeWritten = (int) range.getRangeSize();
        }
        Assert.assertEquals("Size read must equal size written", sizeWritten, readCompleteResult.get());
    }
}