public void clearBuffers() {
for (RecordSerializer<?> serializer : serializers) {
synchronized (serializer) {
try {
Buffer buffer = serializer.getCurrentBuffer();
if (buffer != null) {
buffer.recycle();
}
} finally {
serializer.clear();
}
}
}
}
@Override
public void requestFailed(Buffer buffer, IOException error) {
// Recycle the buffer and forward the error
buffer.recycle();
subpartitionView.notifyError(error);
}
/**
* Returns a successful buffer read request.
*
* <p>Note: This method is always called from the same I/O thread.
*/
private void returnBufferFromIOThread(Buffer buffer) {
final NotificationListener listener;
synchronized (lock) {
if (hasReachedEndOfFile || isReleased) {
buffer.recycle();
return;
}
returnedBuffers.add(buffer);
listener = registeredListener;
registeredListener = null;
// If this was the last buffer before we reached EOF, set the corresponding flag to
// ensure that further buffers are correctly recycled and eventually no further reads
// are triggered.
if (asyncFileReader.hasReachedEndOfFile()) {
hasReachedEndOfFile = true;
}
}
if (listener != null) {
listener.onNotification();
}
}
@Override
public BufferOrEvent getNextBufferOrEvent() throws IOException, InterruptedException {
if (hasReceivedAllEndOfPartitionEvents) {
return null;
}
if (isReleased) {
throw new IllegalStateException("Already released.");
}
requestPartitions();
InputChannel currentChannel = null;
while (currentChannel == null) {
currentChannel = inputChannelsWithData.poll(2, TimeUnit.SECONDS);
}
final Buffer buffer = currentChannel.getNextBuffer();
// Sanity check that notifications only happen when data is available
if (buffer == null) {
throw new IllegalStateException(
"Bug in input gate/channel logic: input gate got"
+ "notified by channel about available data, but none was available.");
}
if (buffer.isBuffer()) {
return new BufferOrEvent(buffer, currentChannel.getChannelIndex());
} else {
final AbstractEvent event = EventSerializer.fromBuffer(buffer, getClass().getClassLoader());
if (event.getClass() == EndOfPartitionEvent.class) {
channelsWithEndOfPartitionEvents.set(currentChannel.getChannelIndex());
if (channelsWithEndOfPartitionEvents.cardinality() == numberOfInputChannels) {
hasReceivedAllEndOfPartitionEvents = true;
}
currentChannel.notifySubpartitionConsumed();
currentChannel.releaseAllResources();
}
return new BufferOrEvent(event, currentChannel.getChannelIndex());
}
}
/**
* 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();
}
}
}
}
@Override
public void releaseAllResources() throws IOException {
try {
synchronized (lock) {
if (!isReleased) {
// Recycle all buffers. Buffers, which are in flight are recycled as soon as
// they return from the I/O thread.
Buffer buffer;
while ((buffer = returnedBuffers.poll()) != null) {
buffer.recycle();
}
isReleased = true;
}
}
} finally {
asyncFileReader.close();
}
}
/**
* Adds a buffer to the subpartition with the given index.
*
* <p>For PIPELINED results, this will trigger the deployment of consuming tasks after the first
* buffer has been added.
*/
public void add(Buffer buffer, int subpartitionIndex) throws IOException {
boolean success = false;
try {
checkInProduceState();
final ResultSubpartition subpartition = subpartitions[subpartitionIndex];
synchronized (subpartition) {
success = subpartition.add(buffer);
// Update statistics
totalNumberOfBuffers++;
totalNumberOfBytes += buffer.getSize();
}
} finally {
if (success) {
notifyPipelinedConsumers();
} else {
buffer.recycle();
}
}
}
// Called by the recycling thread (not network I/O thread)
@Override
public void onEvent(Buffer buffer) {
boolean success = false;
try {
if (buffer != null) {
if (availableBuffer.compareAndSet(null, buffer)) {
ctx.channel().eventLoop().execute(this);
success = true;
} else {
throw new IllegalStateException(
"Received a buffer notification, "
+ " but the previous one has not been handled yet.");
}
} else {
// The buffer pool has been destroyed
stagedBufferResponse = null;
if (stagedMessages.isEmpty()) {
ctx.channel().config().setAutoRead(true);
ctx.channel().read();
} else {
ctx.channel().eventLoop().execute(stagedMessagesHandler);
}
}
} catch (Throwable t) {
ctx.channel().eventLoop().execute(new AsyncErrorNotificationTask(t));
} finally {
if (!success) {
if (buffer != null) {
buffer.recycle();
}
}
}
}
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();
}
}
}