void addInputChannel(RemoteInputChannel listener) {
checkState(!channelError.get(), "There has been an error in the channel.");
if (!inputChannels.containsKey(listener.getInputChannelId())) {
inputChannels.put(listener.getInputChannelId(), listener);
}
}
private boolean decodeMsg(Object msg) throws Throwable {
final Class<?> msgClazz = msg.getClass();
// ---- Buffer --------------------------------------------------------
if (msgClazz == NettyMessage.BufferResponse.class) {
NettyMessage.BufferResponse bufferOrEvent = (NettyMessage.BufferResponse) msg;
RemoteInputChannel inputChannel = inputChannels.get(bufferOrEvent.receiverId);
if (inputChannel == null) {
bufferOrEvent.releaseBuffer();
cancelRequestFor(bufferOrEvent.receiverId);
return true;
}
return decodeBufferOrEvent(inputChannel, bufferOrEvent);
}
// ---- Error ---------------------------------------------------------
else if (msgClazz == NettyMessage.ErrorResponse.class) {
NettyMessage.ErrorResponse error = (NettyMessage.ErrorResponse) msg;
SocketAddress remoteAddr = ctx.channel().remoteAddress();
if (error.isFatalError()) {
notifyAllChannelsOfErrorAndClose(
new RemoteTransportException(
"Fatal error at remote task manager '" + remoteAddr + "'.",
remoteAddr,
error.cause));
} else {
RemoteInputChannel inputChannel = inputChannels.get(error.receiverId);
if (inputChannel != null) {
if (error.cause.getClass() == PartitionNotFoundException.class) {
inputChannel.onFailedPartitionRequest();
} else {
inputChannel.onError(
new RemoteTransportException(
"Error at remote task manager '" + remoteAddr + "'.", remoteAddr, error.cause));
}
}
}
} else {
throw new IllegalStateException("Received unknown message from producer: " + msg.getClass());
}
return true;
}
/**
* Continues the decoding of a staged buffer after a buffer has become available again.
*
* <p>This task is executed by the network I/O thread.
*/
@Override
public void run() {
boolean success = false;
Buffer buffer = null;
try {
if ((buffer = availableBuffer.getAndSet(null)) == null) {
throw new IllegalStateException("Running buffer availability task w/o a buffer.");
}
buffer.setSize(stagedBufferResponse.getSize());
stagedBufferResponse.getNettyBuffer().readBytes(buffer.getNioBuffer());
stagedBufferResponse.releaseBuffer();
RemoteInputChannel inputChannel = inputChannels.get(stagedBufferResponse.receiverId);
if (inputChannel != null) {
inputChannel.onBuffer(buffer, stagedBufferResponse.sequenceNumber);
success = true;
} else {
cancelRequestFor(stagedBufferResponse.receiverId);
}
stagedBufferResponse = null;
if (stagedMessages.isEmpty()) {
ctx.channel().config().setAutoRead(true);
ctx.channel().read();
} else {
ctx.channel().eventLoop().execute(stagedMessagesHandler);
}
} catch (Throwable t) {
notifyAllChannelsOfErrorAndClose(t);
} finally {
if (!success) {
if (buffer != null) {
buffer.recycle();
}
}
}
}
private void notifyAllChannelsOfErrorAndClose(Throwable cause) {
if (channelError.compareAndSet(false, true)) {
try {
for (RemoteInputChannel inputChannel : inputChannels.values()) {
inputChannel.onError(cause);
}
} catch (Throwable t) {
// We can only swallow the Exception at this point. :(
LOG.warn(
"An Exception was thrown during error notification of a " + "remote input channel.", t);
} finally {
inputChannels.clear();
if (ctx != null) {
ctx.close();
}
}
}
}
void removeInputChannel(RemoteInputChannel listener) {
inputChannels.remove(listener.getInputChannelId());
}
private boolean decodeBufferOrEvent(
RemoteInputChannel inputChannel, NettyMessage.BufferResponse bufferOrEvent) throws Throwable {
boolean releaseNettyBuffer = true;
try {
if (bufferOrEvent.isBuffer()) {
// ---- Buffer ------------------------------------------------
// Early return for empty buffers. Otherwise Netty's readBytes() throws an
// IndexOutOfBoundsException.
if (bufferOrEvent.getSize() == 0) {
inputChannel.onEmptyBuffer(bufferOrEvent.sequenceNumber);
return true;
}
BufferProvider bufferProvider = inputChannel.getBufferProvider();
if (bufferProvider == null) {
cancelRequestFor(bufferOrEvent.receiverId);
return false; // receiver has been cancelled/failed
}
while (true) {
Buffer buffer = bufferProvider.requestBuffer();
if (buffer != null) {
buffer.setSize(bufferOrEvent.getSize());
bufferOrEvent.getNettyBuffer().readBytes(buffer.getNioBuffer());
inputChannel.onBuffer(buffer, bufferOrEvent.sequenceNumber);
return true;
} else if (bufferListener.waitForBuffer(bufferProvider, bufferOrEvent)) {
releaseNettyBuffer = false;
return false;
} else if (bufferProvider.isDestroyed()) {
return false;
}
}
} else {
// ---- Event -------------------------------------------------
// TODO We can just keep the serialized data in the Netty buffer and release it later at the
// reader
byte[] byteArray = new byte[bufferOrEvent.getSize()];
bufferOrEvent.getNettyBuffer().readBytes(byteArray);
Buffer buffer =
new Buffer(new MemorySegment(byteArray), FreeingBufferRecycler.INSTANCE, false);
inputChannel.onBuffer(buffer, bufferOrEvent.sequenceNumber);
return true;
}
} finally {
if (releaseNettyBuffer) {
bufferOrEvent.releaseBuffer();
}
}
}