@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());
}
}
