/** Asserts that the message back matches the block response protocols. */
private void assertValid(
final DataServerMessage msg,
final ByteBuffer expectedData,
final long blockId,
final long offset,
final long length) {
Assert.assertEquals(expectedData, msg.getReadOnlyData());
Assert.assertEquals(blockId, msg.getBlockId());
Assert.assertEquals(offset, msg.getOffset());
Assert.assertEquals(length, msg.getLength());
}
/**
* Create a new socket to the data port and send a block request. The returned value is the
* response from the server.
*/
private DataServerMessage request(final BlockInfo block, final long offset, final long length)
throws IOException, TException {
DataServerMessage sendMsg =
DataServerMessage.createBlockRequestMessage(block.blockId, offset, length);
SocketChannel socketChannel =
SocketChannel.open(
new InetSocketAddress(
block.getLocations().get(0).getWorkerAddress().getHost(),
block.getLocations().get(0).getWorkerAddress().getDataPort()));
try {
while (!sendMsg.finishSending()) {
sendMsg.send(socketChannel);
}
DataServerMessage recvMsg =
DataServerMessage.createBlockResponseMessage(false, block.blockId, offset, length, null);
while (!recvMsg.isMessageReady()) {
int numRead = recvMsg.recv(socketChannel);
if (numRead == -1) {
break;
}
}
return recvMsg;
} finally {
socketChannel.close();
}
}
private void write(SelectionKey key) {
SocketChannel socketChannel = (SocketChannel) key.channel();
DataServerMessage sendMessage = mSendingData.get(socketChannel);
boolean closeChannel = false;
try {
sendMessage.send(socketChannel);
} catch (IOException e) {
closeChannel = true;
LOG.error(e.getMessage());
}
if (sendMessage.finishSending() || closeChannel) {
try {
key.channel().close();
} catch (IOException e) {
LOG.error(e.getMessage());
}
key.cancel();
mReceivingData.remove(socketChannel);
mSendingData.remove(socketChannel);
sendMessage.close();
mBlocksLocker.unlock(Math.abs(sendMessage.getBlockId()), sendMessage.getLockId());
}
}
private void read(SelectionKey key) throws IOException {
SocketChannel socketChannel = (SocketChannel) key.channel();
DataServerMessage tMessage;
if (mReceivingData.containsKey(socketChannel)) {
tMessage = mReceivingData.get(socketChannel);
} else {
tMessage = DataServerMessage.createBlockRequestMessage();
mReceivingData.put(socketChannel, tMessage);
}
// Attempt to read off the channel
int numRead;
try {
numRead = tMessage.recv(socketChannel);
} catch (IOException e) {
// The remote forcibly closed the connection, cancel the selection key and close the channel.
key.cancel();
socketChannel.close();
mReceivingData.remove(socketChannel);
mSendingData.remove(socketChannel);
return;
}
if (numRead == -1) {
// Remote entity shut the socket down cleanly. Do the same from our end and cancel the
// channel.
key.channel().close();
key.cancel();
mReceivingData.remove(socketChannel);
mSendingData.remove(socketChannel);
return;
}
if (tMessage.isMessageReady()) {
if (tMessage.getBlockId() <= 0) {
LOG.error("Invalid block id " + tMessage.getBlockId());
return;
}
key.interestOps(SelectionKey.OP_WRITE);
LOG.info("Get request for " + tMessage.getBlockId());
int lockId = mBlocksLocker.lock(tMessage.getBlockId());
DataServerMessage tResponseMessage =
DataServerMessage.createBlockResponseMessage(
true, tMessage.getBlockId(), tMessage.getOffset(), tMessage.getLength());
tResponseMessage.setLockId(lockId);
mSendingData.put(socketChannel, tResponseMessage);
}
}
/** Asserts that the message back matches the block response protocols for the error case. */
private void assertError(final DataServerMessage msg, final long blockId) {
Assert.assertEquals(blockId, msg.getBlockId());
Assert.assertEquals(0, msg.getLength());
Assert.assertNotEquals(msg.getStatus().getId(), RPCResponse.Status.SUCCESS.getId());
}
