[hotfix][network] Implement UnionInputGate#pollNextBufferOrEvent method

is · May 10, 2019 · cd13fdf · cd13fdf
1 parent dd1eb08
commit cd13fdf
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 20 deletions.
diff --git a/.../src/main/java/org/apache/flink/runtime/io/network/partition/consumer/UnionInputGate.java b/.../src/main/java/org/apache/flink/runtime/io/network/partition/consumer/UnionInputGate.java
@@ -159,14 +159,28 @@ public void requestPartitions() throws IOException, InterruptedException {
 
  @Override
  public Optional<BufferOrEvent> getNextBufferOrEvent() throws IOException, InterruptedException {
+ return getNextBufferOrEvent(true);
+ }
+
+ @Override
+ public Optional<BufferOrEvent> pollNextBufferOrEvent() throws IOException, InterruptedException {
+ return getNextBufferOrEvent(false);
+ }
+
+ private Optional<BufferOrEvent> getNextBufferOrEvent(boolean blocking) throws IOException, InterruptedException {
  if (inputGatesWithRemainingData.isEmpty()) {
  return Optional.empty();
  }
 
  // Make sure to request the partitions, if they have not been requested before.
  requestPartitions();
 
- InputGateWithData inputGateWithData = waitAndGetNextInputGate();
+ Optional<InputGateWithData> next = waitAndGetNextInputGate(blocking);
+ if (!next.isPresent()) {
+ return Optional.empty();
+ }
+
+ InputGateWithData inputGateWithData = next.get();
  InputGate inputGate = inputGateWithData.inputGate;
  BufferOrEvent bufferOrEvent = inputGateWithData.bufferOrEvent;
 
@@ -197,18 +211,17 @@ public Optional<BufferOrEvent> getNextBufferOrEvent() throws IOException, Interr
  return Optional.of(bufferOrEvent);
  }
 
- @Override
- public Optional<BufferOrEvent> pollNextBufferOrEvent() throws UnsupportedOperationException {
- throw new UnsupportedOperationException();
- }
-
- private InputGateWithData waitAndGetNextInputGate() throws IOException, InterruptedException {
+ private Optional<InputGateWithData> waitAndGetNextInputGate(boolean blocking) throws IOException, InterruptedException {
  while (true) {
  InputGate inputGate;
  boolean moreInputGatesAvailable;
  synchronized (inputGatesWithData) {
  while (inputGatesWithData.size() == 0) {
- inputGatesWithData.wait();
+ if (blocking) {
+ inputGatesWithData.wait();
+ } else {
+ return Optional.empty();
+ }
  }
  inputGate = inputGatesWithData.remove();
  enqueuedInputGatesWithData.remove(inputGate);
@@ -218,7 +231,7 @@ private InputGateWithData waitAndGetNextInputGate() throws IOException, Interrup
  // In case of inputGatesWithData being inaccurate do not block on an empty inputGate, but just poll the data.
  Optional<BufferOrEvent> bufferOrEvent = inputGate.pollNextBufferOrEvent();
  if (bufferOrEvent.isPresent()) {
- return new InputGateWithData(inputGate, bufferOrEvent.get(), moreInputGatesAvailable);
+ return Optional.of(new InputGateWithData(inputGate, bufferOrEvent.get(), moreInputGatesAvailable));
  }
  }
  }

diff --git a/...test/java/org/apache/flink/runtime/io/network/partition/consumer/SingleInputGateTest.java b/...test/java/org/apache/flink/runtime/io/network/partition/consumer/SingleInputGateTest.java
@@ -619,17 +619,7 @@ static void verifyBufferOrEvent(
  assertEquals(expectedChannelIndex, bufferOrEvent.get().getChannelIndex());
  assertEquals(expectedMoreAvailable, bufferOrEvent.get().moreAvailable());
  if (!expectedMoreAvailable) {
- try {
- assertFalse(inputGate.pollNextBufferOrEvent().isPresent());
- }
- catch (UnsupportedOperationException ex) {
- /**
- * {@link UnionInputGate#pollNextBufferOrEvent()} is unsupported at the moment.
- */
- if (!(inputGate instanceof UnionInputGate)) {
- throw ex;
- }
- }
+ assertFalse(inputGate.pollNextBufferOrEvent().isPresent());
  }
  }
 }