/**
* Copyright 2016 LinkedIn Corp. All rights reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
*/
package com.github.ambry.commons;
import com.github.ambry.router.AsyncWritableChannel;
import com.github.ambry.router.Callback;
import com.github.ambry.router.FutureResult;
import com.github.ambry.utils.Utils;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.ClosedChannelException;
import java.util.Random;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import org.junit.Test;
import static org.junit.Assert.*;
/**
* Tests functionality of {@link ByteBufferReadableStreamChannel}.
*/
public class ByteBufferReadableStreamChannelTest {
/**
* Tests the common case read operations i.e
* 1. Create {@link ByteBufferReadableStreamChannel} with random bytes.
* 2. Calls the different read operations of {@link ByteBufferReadableStreamChannel} and checks that the data read
* matches the data used to create the {@link ByteBufferReadableStreamChannel}.
* @throws Exception
*/
@Test
public void commonCaseTest() throws Exception {
ByteBuffer content = ByteBuffer.wrap(fillRandomBytes(new byte[1024]));
ByteBufferReadableStreamChannel readableStreamChannel = new ByteBufferReadableStreamChannel(content);
assertTrue("ByteBufferReadableStreamChannel is not open", readableStreamChannel.isOpen());
assertEquals("Size returned by ByteBufferReadableStreamChannel did not match source array size", content.capacity(),
readableStreamChannel.getSize());
ByteBufferAsyncWritableChannel writeChannel = new ByteBufferAsyncWritableChannel();
ReadIntoCallback callback = new ReadIntoCallback();
Future<Long> future = readableStreamChannel.readInto(writeChannel, callback);
ByteBuffer contentWrapper = ByteBuffer.wrap(content.array());
while (contentWrapper.hasRemaining()) {
ByteBuffer recvdContent = writeChannel.getNextChunk();
assertNotNull("Written content lesser than original content", recvdContent);
while (recvdContent.hasRemaining()) {
assertTrue("Written content is more than original content", contentWrapper.hasRemaining());
assertEquals("Unexpected byte", contentWrapper.get(), recvdContent.get());
}
writeChannel.resolveOldestChunk(null);
}
assertNull("There should have been no more data in the channel", writeChannel.getNextChunk(0));
writeChannel.close();
callback.awaitCallback();
if (callback.exception != null) {
throw callback.exception;
}
long futureBytesRead = future.get();
assertEquals("Total bytes written does not match (callback)", content.limit(), callback.bytesRead);
assertEquals("Total bytes written does not match (future)", content.limit(), futureBytesRead);
}
/**
* Tests that the right exceptions are thrown when reading into {@link AsyncWritableChannel} fails.
* @throws Exception
*/
@Test
public void readIntoAWCFailureTest() throws Exception {
String errMsg = "@@ExpectedExceptionMessage@@";
byte[] in = fillRandomBytes(new byte[1]);
// Bad AWC.
ByteBufferReadableStreamChannel readableStreamChannel = new ByteBufferReadableStreamChannel(ByteBuffer.wrap(in));
ReadIntoCallback callback = new ReadIntoCallback();
try {
readableStreamChannel.readInto(new BadAsyncWritableChannel(new IOException(errMsg)), callback).get();
fail("Should have failed because BadAsyncWritableChannel would have thrown exception");
} catch (ExecutionException e) {
Exception exception = (Exception) Utils.getRootCause(e);
assertEquals("Exception message does not match expected (future)", errMsg, exception.getMessage());
callback.awaitCallback();
assertEquals("Exception message does not match expected (callback)", errMsg, callback.exception.getMessage());
}
// Reading more than once.
readableStreamChannel = new ByteBufferReadableStreamChannel(ByteBuffer.wrap(in));
ByteBufferAsyncWritableChannel writeChannel = new ByteBufferAsyncWritableChannel();
readableStreamChannel.readInto(writeChannel, null);
try {
readableStreamChannel.readInto(writeChannel, null);
fail("Should have failed because readInto cannot be called more than once");
} catch (IllegalStateException e) {
// expected. Nothing to do.
}
// Read after close.
readableStreamChannel = new ByteBufferReadableStreamChannel(ByteBuffer.wrap(in));
readableStreamChannel.close();
writeChannel = new ByteBufferAsyncWritableChannel();
callback = new ReadIntoCallback();
try {
readableStreamChannel.readInto(writeChannel, callback).get();
fail("ByteBufferReadableStreamChannel has been closed, so read should have thrown ClosedChannelException");
} catch (ExecutionException e) {
Exception exception = (Exception) Utils.getRootCause(e);
assertTrue("Exception is not ClosedChannelException", exception instanceof ClosedChannelException);
callback.awaitCallback();
assertEquals("Exceptions of callback and future differ", exception.getMessage(), callback.exception.getMessage());
}
}
/**
* Tests behavior of read operations on some corner cases.
* <p/>
* Corner case list:
* 1. Blob size is 0.
* @throws Exception
*/
@Test
public void readAndWriteCornerCasesTest() throws Exception {
// 0 sized blob.
ByteBufferReadableStreamChannel readableStreamChannel = new ByteBufferReadableStreamChannel(ByteBuffer.allocate(0));
assertTrue("ByteBufferReadableStreamChannel is not open", readableStreamChannel.isOpen());
assertEquals("Size returned by ByteBufferReadableStreamChannel is not 0", 0, readableStreamChannel.getSize());
ByteBufferAsyncWritableChannel writeChannel = new ByteBufferAsyncWritableChannel();
ReadIntoCallback callback = new ReadIntoCallback();
Future<Long> future = readableStreamChannel.readInto(writeChannel, callback);
ByteBuffer chunk = writeChannel.getNextChunk(0);
while (chunk != null) {
writeChannel.resolveOldestChunk(null);
chunk = writeChannel.getNextChunk(0);
}
callback.awaitCallback();
assertEquals("There should have no bytes to read (future)", 0, future.get().longValue());
assertEquals("There should have no bytes to read (callback)", 0, callback.bytesRead);
if (callback.exception != null) {
throw callback.exception;
}
writeChannel.close();
readableStreamChannel.close();
}
/**
* Tests that no exceptions are thrown on repeating idempotent operations. Does <b><i>not</i></b> currently test that
* state changes are idempotent.
* @throws IOException
*/
@Test
public void idempotentOperationsTest() throws IOException {
byte[] in = fillRandomBytes(new byte[1]);
ByteBufferReadableStreamChannel byteBufferReadableStreamChannel =
new ByteBufferReadableStreamChannel(ByteBuffer.wrap(in));
assertTrue("ByteBufferReadableStreamChannel is not open", byteBufferReadableStreamChannel.isOpen());
byteBufferReadableStreamChannel.close();
assertFalse("ByteBufferReadableStreamChannel is not closed", byteBufferReadableStreamChannel.isOpen());
// should not throw exception.
byteBufferReadableStreamChannel.close();
assertFalse("ByteBufferReadableStreamChannel is not closed", byteBufferReadableStreamChannel.isOpen());
}
// helpers
// general
/**
* Fills random bytes into {@code in}.
* @param in the byte array that needs to be filled with random bytes.
* @return {@code in} filled with random bytes.
*/
private byte[] fillRandomBytes(byte[] in) {
new Random().nextBytes(in);
return in;
}
}
/**
* Callback for read operations on {@link ByteBufferReadableStreamChannel}.
*/
class ReadIntoCallback implements Callback<Long> {
public volatile long bytesRead;
public volatile Exception exception;
private final AtomicBoolean callbackInvoked = new AtomicBoolean(false);
private final CountDownLatch latch = new CountDownLatch(1);
@Override
public void onCompletion(Long result, Exception exception) {
if (callbackInvoked.compareAndSet(false, true)) {
bytesRead = result;
this.exception = exception;
latch.countDown();
} else {
this.exception = new IllegalStateException("Callback invoked more than once");
}
}
/**
* Waits for the callback to arrive for a limited amount of time.
* @throws InterruptedException
* @throws TimeoutException
*/
void awaitCallback() throws InterruptedException, TimeoutException {
if (!latch.await(1, TimeUnit.SECONDS)) {
throw new TimeoutException("Waiting too long for callback to arrive");
}
}
}
/**
* A {@link AsyncWritableChannel} that throws a custom exception (provided at construction time) on a call to
* {@link #write(ByteBuffer, Callback)}.
*/
class BadAsyncWritableChannel implements AsyncWritableChannel {
private final Exception exceptionToThrow;
private final AtomicBoolean isOpen = new AtomicBoolean(true);
/**
* Creates an instance of BadAsyncWritableChannel that throws {@code exceptionToThrow} on write.
* @param exceptionToThrow the {@link Exception} to throw on write.
*/
public BadAsyncWritableChannel(Exception exceptionToThrow) {
this.exceptionToThrow = exceptionToThrow;
}
@Override
public Future<Long> write(ByteBuffer src, Callback<Long> callback) {
if (exceptionToThrow instanceof RuntimeException) {
throw (RuntimeException) exceptionToThrow;
} else {
return markFutureInvokeCallback(callback, 0, exceptionToThrow);
}
}
@Override
public boolean isOpen() {
return isOpen.get();
}
@Override
public void close() throws IOException {
isOpen.set(false);
}
/**
* Creates and marks a future as done and invoked the callback with paramaters {@code totalBytesWritten} and
* {@code Exception}.
* @param callback the {@link Callback} to invoke.
* @param totalBytesWritten the number of bytes successfully written.
* @param exception the {@link Exception} that occurred if any.
* @return the {@link Future} that will contain the result of the operation.
*/
private Future<Long> markFutureInvokeCallback(Callback<Long> callback, long totalBytesWritten, Exception exception) {
FutureResult<Long> futureResult = new FutureResult<Long>();
futureResult.done(totalBytesWritten, exception);
if (callback != null) {
callback.onCompletion(totalBytesWritten, exception);
}
return futureResult;
}
}