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

Example 1 with Callback

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

the class GetBlobOperationTest method testEarlyReadableStreamChannelClose.

/**
 * Test that the operation is completed and an exception with the error code {@link RouterErrorCode#ChannelClosed} is
 * set when the {@link ReadableStreamChannel} is closed before all chunks are read for a specific blob size and
 * number of chunks to read.
 * @param numChunksInBlob the number of chunks in the blob to put/get.
 * @param numChunksToRead the number of chunks to read from the {@link AsyncWritableChannel} before closing the
 *                        {@link ReadableStreamChannel}.
 * @throws Exception
 */
private void testEarlyReadableStreamChannelClose(int numChunksInBlob, final int numChunksToRead) throws Exception {
    final AtomicReference<Exception> callbackException = new AtomicReference<>();
    final AtomicReference<Future<Long>> readIntoFuture = new AtomicReference<>();
    final CountDownLatch readCompleteLatch = new CountDownLatch(1);
    Callback<GetBlobResultInternal> callback = new Callback<GetBlobResultInternal>() {

        @Override
        public void onCompletion(final GetBlobResultInternal result, Exception exception) {
            if (exception != null) {
                callbackException.set(exception);
                readCompleteLatch.countDown();
            } else {
                final ByteBufferAsyncWritableChannel writableChannel = new ByteBufferAsyncWritableChannel();
                readIntoFuture.set(result.getBlobResult.getBlobDataChannel().readInto(writableChannel, null));
                Utils.newThread(new Runnable() {

                    @Override
                    public void run() {
                        try {
                            int chunksLeftToRead = numChunksToRead;
                            while (chunksLeftToRead > 0) {
                                writableChannel.getNextChunk();
                                writableChannel.resolveOldestChunk(null);
                                chunksLeftToRead--;
                            }
                            result.getBlobResult.getBlobDataChannel().close();
                            while (writableChannel.getNextChunk(100) != null) {
                                writableChannel.resolveOldestChunk(null);
                            }
                        } catch (Exception e) {
                            callbackException.set(e);
                        } finally {
                            readCompleteLatch.countDown();
                        }
                    }
                }, false).start();
            }
        }
    };
    blobSize = numChunksInBlob * maxChunkSize;
    doPut();
    GetBlobOperation op = createOperationAndComplete(callback);
    Assert.assertTrue("Timeout waiting for read to complete", readCompleteLatch.await(2, TimeUnit.SECONDS));
    if (callbackException.get() != null) {
        throw callbackException.get();
    }
    try {
        readIntoFuture.get().get();
        Assert.fail("Expected ExecutionException");
    } catch (ExecutionException e) {
        Assert.assertTrue("Unexpected type for exception: " + e.getCause(), e.getCause() instanceof RouterException);
        Assert.assertEquals("Unexpected RouterErrorCode", RouterErrorCode.ChannelClosed, ((RouterException) e.getCause()).getErrorCode());
    }
    Exception operationException = op.getOperationException();
    Assert.assertTrue("Unexpected type for exception: " + operationException, operationException instanceof RouterException);
    Assert.assertEquals("Unexpected RouterErrorCode", RouterErrorCode.ChannelClosed, ((RouterException) operationException).getErrorCode());
}

Example 2 with Callback

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

the class GetBlobOperationTest method testDataChunkError.

/**
 * Test that an operation is completed with a specified {@link RouterErrorCode} when all gets on data chunks
 * in a multi-part blob return a specified {@link ServerErrorCode}
 * @param serverErrorCode The error code to be returned when fetching data chunks.
 * @param expectedErrorCode The operation's expected error code.
 * @throws Exception
 */
private void testDataChunkError(ServerErrorCode serverErrorCode, final RouterErrorCode expectedErrorCode) throws Exception {
    blobSize = maxChunkSize * 2 + 1;
    doPut();
    final CountDownLatch readCompleteLatch = new CountDownLatch(1);
    final AtomicReference<Exception> readCompleteException = new AtomicReference<>(null);
    final ByteBufferAsyncWritableChannel asyncWritableChannel = new ByteBufferAsyncWritableChannel();
    mockServerLayout.getMockServers().forEach(mockServer -> mockServer.setGetErrorOnDataBlobOnly(true));
    RouterTestHelpers.testWithErrorCodes(Collections.singletonMap(serverErrorCode, 9), mockServerLayout, expectedErrorCode, new ErrorCodeChecker() {

        @Override
        public void testAndAssert(RouterErrorCode expectedError) throws Exception {
            Callback<GetBlobResultInternal> callback = new Callback<GetBlobResultInternal>() {

                @Override
                public void onCompletion(final GetBlobResultInternal result, final Exception exception) {
                    if (exception != null) {
                        asyncWritableChannel.close();
                        readCompleteLatch.countDown();
                    } else {
                        Utils.newThread(new Runnable() {

                            @Override
                            public void run() {
                                try {
                                    result.getBlobResult.getBlobDataChannel().readInto(asyncWritableChannel, new Callback<Long>() {

                                        @Override
                                        public void onCompletion(Long result, Exception exception) {
                                            asyncWritableChannel.close();
                                        }
                                    });
                                    asyncWritableChannel.getNextChunk();
                                } catch (Exception e) {
                                    readCompleteException.set(e);
                                } finally {
                                    readCompleteLatch.countDown();
                                }
                            }
                        }, false).start();
                    }
                }
            };
            GetBlobOperation op = createOperationAndComplete(callback);
            Assert.assertTrue(readCompleteLatch.await(2, TimeUnit.SECONDS));
            Assert.assertTrue("Operation should be complete at this time", op.isOperationComplete());
            if (readCompleteException.get() != null) {
                throw readCompleteException.get();
            }
            Assert.assertFalse("AsyncWriteableChannel should have been closed.", asyncWritableChannel.isOpen());
            assertFailureAndCheckErrorCode(op, expectedError);
        }
    });
}

Example 3 with Callback

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

the class MockRouterCallback method testInstantiation.

/**
 * Test {@link GetBlobInfoOperation} instantiation and validate the get methods.
 */
@Test
public void testInstantiation() {
    BlobId blobId = new BlobId(routerConfig.routerBlobidCurrentVersion, BlobId.BlobIdType.NATIVE, ClusterMap.UNKNOWN_DATACENTER_ID, Utils.getRandomShort(random), Utils.getRandomShort(random), mockClusterMap.getWritablePartitionIds(MockClusterMap.DEFAULT_PARTITION_CLASS).get(0), false, BlobId.BlobDataType.DATACHUNK);
    Callback<GetBlobResultInternal> getOperationCallback = (result, exception) -> {
    // no op.
    };
    // test a good case
    GetBlobInfoOperation op = new GetBlobInfoOperation(routerConfig, routerMetrics, mockClusterMap, responseHandler, blobId, options, getOperationCallback, routerCallback, kms, cryptoService, cryptoJobHandler, time, false, quotaChargeCallback);
    Assert.assertEquals("Callback must match", getOperationCallback, op.getCallback());
    Assert.assertEquals("Blob ids must match", blobId.getID(), op.getBlobIdStr());
    // test the case where the tracker type is bad
    Properties properties = getNonBlockingRouterProperties(true);
    properties.setProperty("router.get.operation.tracker.type", "NonExistentTracker");
    RouterConfig badConfig = new RouterConfig(new VerifiableProperties(properties));
    try {
        new GetBlobInfoOperation(badConfig, routerMetrics, mockClusterMap, responseHandler, blobId, options, getOperationCallback, routerCallback, kms, cryptoService, cryptoJobHandler, time, false, quotaChargeCallback);
        Assert.fail("Instantiation of GetBlobInfoOperation with an invalid tracker type must fail");
    } catch (IllegalArgumentException e) {
    // expected. Nothing to do.
    }
}

Example 4 with Callback

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

the class AmbryServerSecurityServiceTest method validateConnectionTest.

/**
 * Tests for {@link AmbryServerSecurityService#validateConnection(SSLSession, Callback)}
 * @throws Exception
 */
@Test
public void validateConnectionTest() throws Exception {
    // sslSession is null
    TestUtils.assertException(IllegalArgumentException.class, () -> serverSecurityService.validateConnection(null).get(), null);
    // success case
    SSLSession sslSession = Mockito.mock(SSLSession.class);
    serverSecurityService.validateConnection(sslSession, (r, e) -> {
        Assert.assertNull("result not null", r);
        Assert.assertNull("exception not null", e);
    });
    // service is closed
    serverSecurityService.close();
    ThrowingConsumer<ExecutionException> errorAction = e -> {
        Assert.assertTrue("Exception should have been an instance of RestServiceException", e.getCause() instanceof RestServiceException);
        RestServiceException re = (RestServiceException) e.getCause();
        Assert.assertEquals("Unexpected RestServerErrorCode (Future)", RestServiceErrorCode.ServiceUnavailable, re.getErrorCode());
    };
    TestUtils.assertException(ExecutionException.class, () -> serverSecurityService.validateConnection(sslSession).get(), errorAction);
}

Example 5 with Callback

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

the class AsyncWritableChannel method write.

/**
 * Write data in {@code src} to the channel and the {@code callback} will be invoked once the write succeeds or fails.
 * This method is the counterpart for {@link ByteBuf}. It shares the same guarantee as the {@link #write(ByteBuffer, Callback)}.
 * Whoever implements this interface, shouldn't release the {@code src}. If releasing is expected after finishing writing
 * to channel, please release this {@link ByteBuf} in the callback method.
 * @param src The data taht needs to be written to the channel.
 * @param callback The {@link Callback} that will be invoked once the write succeeds/fails. This can be null.
 * @return a {@link Future} that will eventually contain the result of the write operation (the number of bytes
 *         written).
 */
default Future<Long> write(ByteBuf src, Callback<Long> callback) {
    if (src == null) {
        throw new IllegalArgumentException("Source ByteBuf cannot be null");
    }
    int numBuffers = src.nioBufferCount();
    FutureResult<Long> futureResult = new FutureResult<>();
    Callback<Long> singleBufferCallback = (result, exception) -> {
        if (result != 0) {
            src.readerIndex(src.readerIndex() + (int) result.longValue());
        }
        futureResult.done(result, exception);
        if (callback != null) {
            callback.onCompletion(result, exception);
        }
    };
    if (numBuffers < 1) {
        ByteBuffer byteBuffer = ByteBuffer.allocate(src.readableBytes());
        src.getBytes(src.readerIndex(), byteBuffer);
        write(byteBuffer, singleBufferCallback);
    } else if (numBuffers == 1) {
        write(src.nioBuffer(), singleBufferCallback);
    } else {
        ByteBuffer[] buffers = src.nioBuffers();
        AtomicLong size = new AtomicLong(0);
        AtomicInteger index = new AtomicInteger(0);
        AtomicBoolean callbackInvoked = new AtomicBoolean(false);
        Callback<Long> cb = (result, exception) -> {
            index.addAndGet(1);
            size.addAndGet(result);
            if (result != 0) {
                src.readerIndex(src.readerIndex() + (int) result.longValue());
            }
            if ((exception != null || index.get() == buffers.length) && callbackInvoked.compareAndSet(false, true)) {
                futureResult.done(size.get(), exception);
                if (callback != null) {
                    callback.onCompletion(size.get(), exception);
                }
            }
        };
        for (int i = 0; i < buffers.length; i++) {
            write(buffers[i], cb);
        }
    }
    return futureResult;
}