/** Sets up all dependencies before running a test. */
@Before
public void before() {
mLineageMasterClient = PowerMockito.mock(LineageMasterClient.class);
mLineageContext = PowerMockito.mock(LineageContext.class);
Mockito.when(mLineageContext.acquireMasterClient()).thenReturn(mLineageMasterClient);
Whitebox.setInternalState(LineageContext.class, "INSTANCE", mLineageContext);
mAlluxioLineageFileSystem = LineageFileSystem.get();
Whitebox.setInternalState(mAlluxioLineageFileSystem, "mLineageContext", mLineageContext);
FileSystemContext fileSystemContext = PowerMockito.mock(FileSystemContext.class);
FileSystemMasterClient fileSystemMasterClient = PowerMockito.mock(FileSystemMasterClient.class);
Mockito.when(fileSystemContext.acquireMasterClient()).thenReturn(fileSystemMasterClient);
Whitebox.setInternalState(FileSystemContext.class, "INSTANCE", fileSystemContext);
Whitebox.setInternalState(mAlluxioLineageFileSystem, "mContext", fileSystemContext);
}
@Override
public void close() throws IOException {
if (mContext != FileSystemContext.INSTANCE) {
mContext.close();
}
super.close();
}
/** Sets the file system and context. */
private void updateFileSystemAndContext() {
Subject subject = getHadoopSubject();
if (subject != null) {
mContext = FileSystemContext.create(subject);
mFileSystem = FileSystem.Factory.get(mContext);
} else {
mContext = FileSystemContext.INSTANCE;
mFileSystem = FileSystem.Factory.get();
}
}
/**
* Creates an instance of {@link NettyPacketReader}. If this is used to read a block remotely, it
* requires the block to be locked beforehand and the lock ID is passed to this class.
*
* @param context the file system context
* @param address the netty data server network address
* @param id the block ID or UFS file ID
* @param offset the offset
* @param len the length to read
* @param lockId the lock ID
* @param sessionId the session ID
* @param type the request type (block or UFS file)
* @throws IOException if it fails to acquire a netty channel
*/
private NettyPacketReader(
FileSystemContext context,
InetSocketAddress address,
long id,
long offset,
long len,
long lockId,
long sessionId,
Protocol.RequestType type)
throws IOException {
Preconditions.checkArgument(offset >= 0 && len > 0);
mContext = context;
mAddress = address;
mId = id;
mStart = offset;
mPosToRead = offset;
mBytesToRead = len;
mRequestType = type;
mChannel = context.acquireNettyChannel(address);
ChannelPipeline pipeline = mChannel.pipeline();
if (!(pipeline.last() instanceof RPCMessageDecoder)) {
throw new RuntimeException(
String.format(
"Channel pipeline has unexpected handlers %s.",
pipeline.last().getClass().getCanonicalName()));
}
mChannel.pipeline().addLast(new PacketReadHandler());
Protocol.ReadRequest readRequest =
Protocol.ReadRequest.newBuilder()
.setId(id)
.setOffset(offset)
.setLength(len)
.setLockId(lockId)
.setSessionId(sessionId)
.setType(type)
.build();
mChannel
.writeAndFlush(new RPCProtoMessage(readRequest))
.addListener(ChannelFutureListener.CLOSE_ON_FAILURE);
}
@Override
public void close() {
try {
if (mDone) {
return;
}
if (!mChannel.isOpen()) {
return;
}
try {
if (!CANCEL_ENABLED) {
mChannel.close().sync();
return;
}
if (remaining() > 0) {
Protocol.ReadRequest cancelRequest =
Protocol.ReadRequest.newBuilder()
.setId(mId)
.setCancel(true)
.setType(mRequestType)
.build();
mChannel
.writeAndFlush(new RPCProtoMessage(cancelRequest))
.addListener(ChannelFutureListener.CLOSE_ON_FAILURE);
}
} catch (InterruptedException e) {
mChannel.close();
throw Throwables.propagate(e);
}
while (true) {
try {
ByteBuf buf = readPacket();
// A null packet indicates the end of the stream.
if (buf == null) {
return;
}
buf.release();
} catch (IOException e) {
LOG.warn(
"Failed to close the NettyBlockReader (block: {}, address: {}).", mId, mAddress, e);
try {
mChannel.close().sync();
} catch (InterruptedException ee) {
throw Throwables.propagate(ee);
}
return;
}
}
} finally {
if (mChannel.isOpen()) {
Preconditions.checkState(mChannel.pipeline().last() instanceof PacketReadHandler);
mChannel.pipeline().removeLast();
// Make sure "autoread" is on before realsing the channel.
resume();
}
mContext.releaseNettyChannel(mAddress, mChannel);
mClosed = true;
}
}
