/**
* Adds channel handlers that perform encryption / decryption of data using SASL.
*
* @param channel The channel.
* @param backend The SASL backend.
* @param maxOutboundBlockSize Max size in bytes of outgoing encrypted blocks, to control memory
* usage.
*/
static void addToChannel(
Channel channel, SaslEncryptionBackend backend, int maxOutboundBlockSize) {
channel
.pipeline()
.addFirst(ENCRYPTION_HANDLER_NAME, new EncryptionHandler(backend, maxOutboundBlockSize))
.addFirst("saslDecryption", new DecryptionHandler(backend))
.addFirst("saslFrameDecoder", NettyUtils.createFrameDecoder());
}
public ExternalShuffleBlockResolver(TransportConf conf, File registeredExecutorFile)
throws IOException {
this(
conf,
registeredExecutorFile,
Executors.newSingleThreadExecutor(
// Add `spark` prefix because it will run in NM in Yarn mode.
NettyUtils.createThreadFactory("spark-shuffle-directory-cleaner")));
}
private void processStreamRequest(final StreamRequest req) {
final String client = NettyUtils.getRemoteAddress(channel);
ManagedBuffer buf;
try {
buf = streamManager.openStream(req.streamId);
} catch (Exception e) {
logger.error(
String.format("Error opening stream %s for request from %s", req.streamId, client), e);
respond(new StreamFailure(req.streamId, Throwables.getStackTraceAsString(e)));
return;
}
respond(new StreamResponse(req.streamId, buf.size(), buf));
}
private void processFetchRequest(final ChunkFetchRequest req) {
final String client = NettyUtils.getRemoteAddress(channel);
logger.trace("Received req from {} to fetch block {}", client, req.streamChunkId);
ManagedBuffer buf;
try {
streamManager.checkAuthorization(reverseClient, req.streamChunkId.streamId);
streamManager.registerChannel(channel, req.streamChunkId.streamId);
buf = streamManager.getChunk(req.streamChunkId.streamId, req.streamChunkId.chunkIndex);
} catch (Exception e) {
logger.error(
String.format("Error opening block %s for request from %s", req.streamChunkId, client),
e);
respond(new ChunkFetchFailure(req.streamChunkId, Throwables.getStackTraceAsString(e)));
return;
}
respond(new ChunkFetchSuccess(req.streamChunkId, buf));
}
/**
* Request to stream the data with the given stream ID from the remote end.
*
* @param streamId The stream to fetch.
* @param callback Object to call with the stream data.
*/
public void stream(final String streamId, final StreamCallback callback) {
final String serverAddr = NettyUtils.getRemoteAddress(channel);
final long startTime = System.currentTimeMillis();
logger.debug("Sending stream request for {} to {}", streamId, serverAddr);
// Need to synchronize here so that the callback is added to the queue and the RPC is
// written to the socket atomically, so that callbacks are called in the right order
// when responses arrive.
synchronized (this) {
handler.addStreamCallback(callback);
channel
.writeAndFlush(new StreamRequest(streamId))
.addListener(
new ChannelFutureListener() {
@Override
public void operationComplete(ChannelFuture future) throws Exception {
if (future.isSuccess()) {
long timeTaken = System.currentTimeMillis() - startTime;
logger.trace(
"Sending request for {} to {} took {} ms", streamId, serverAddr, timeTaken);
} else {
String errorMsg =
String.format(
"Failed to send request for %s to %s: %s",
streamId, serverAddr, future.cause());
logger.error(errorMsg, future.cause());
channel.close();
try {
callback.onFailure(streamId, new IOException(errorMsg, future.cause()));
} catch (Exception e) {
logger.error("Uncaught exception in RPC response callback handler!", e);
}
}
}
});
}
}
/**
* Sends an opaque message to the RpcHandler on the server-side. The callback will be invoked with
* the server's response or upon any failure.
*
* @param message The message to send.
* @param callback Callback to handle the RPC's reply.
* @return The RPC's id.
*/
public long sendRpc(ByteBuffer message, final RpcResponseCallback callback) {
final String serverAddr = NettyUtils.getRemoteAddress(channel);
final long startTime = System.currentTimeMillis();
logger.trace("Sending RPC to {}", serverAddr);
final long requestId = Math.abs(UUID.randomUUID().getLeastSignificantBits());
handler.addRpcRequest(requestId, callback);
channel
.writeAndFlush(new RpcRequest(requestId, new NioManagedBuffer(message)))
.addListener(
new ChannelFutureListener() {
@Override
public void operationComplete(ChannelFuture future) throws Exception {
if (future.isSuccess()) {
long timeTaken = System.currentTimeMillis() - startTime;
logger.trace(
"Sending request {} to {} took {} ms", requestId, serverAddr, timeTaken);
} else {
String errorMsg =
String.format(
"Failed to send RPC %s to %s: %s", requestId, serverAddr, future.cause());
logger.error(errorMsg, future.cause());
handler.removeRpcRequest(requestId);
channel.close();
try {
callback.onFailure(new IOException(errorMsg, future.cause()));
} catch (Exception e) {
logger.error("Uncaught exception in RPC response callback handler!", e);
}
}
}
});
return requestId;
}
/**
* Requests a single chunk from the remote side, from the pre-negotiated streamId.
*
* <p>Chunk indices go from 0 onwards. It is valid to request the same chunk multiple times,
* though some streams may not support this.
*
* <p>Multiple fetchChunk requests may be outstanding simultaneously, and the chunks are
* guaranteed to be returned in the same order that they were requested, assuming only a single
* TransportClient is used to fetch the chunks.
*
* @param streamId Identifier that refers to a stream in the remote StreamManager. This should be
* agreed upon by client and server beforehand.
* @param chunkIndex 0-based index of the chunk to fetch
* @param callback Callback invoked upon successful receipt of chunk, or upon any failure.
*/
public void fetchChunk(
long streamId, final int chunkIndex, final ChunkReceivedCallback callback) {
final String serverAddr = NettyUtils.getRemoteAddress(channel);
final long startTime = System.currentTimeMillis();
logger.debug("Sending fetch chunk request {} to {}", chunkIndex, serverAddr);
final StreamChunkId streamChunkId = new StreamChunkId(streamId, chunkIndex);
handler.addFetchRequest(streamChunkId, callback);
channel
.writeAndFlush(new ChunkFetchRequest(streamChunkId))
.addListener(
new ChannelFutureListener() {
@Override
public void operationComplete(ChannelFuture future) throws Exception {
if (future.isSuccess()) {
long timeTaken = System.currentTimeMillis() - startTime;
logger.trace(
"Sending request {} to {} took {} ms", streamChunkId, serverAddr, timeTaken);
} else {
String errorMsg =
String.format(
"Failed to send request %s to %s: %s",
streamChunkId, serverAddr, future.cause());
logger.error(errorMsg, future.cause());
handler.removeFetchRequest(streamChunkId);
channel.close();
try {
callback.onFailure(chunkIndex, new IOException(errorMsg, future.cause()));
} catch (Exception e) {
logger.error("Uncaught exception in RPC response callback handler!", e);
}
}
}
});
}
