You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by dw...@apache.org on 2021/09/24 06:43:48 UTC
[flink] 01/03: [hotfix] Make it possible to force MockSourceReader
to wait for initial splits assignment
This is an automated email from the ASF dual-hosted git repository.
dwysakowicz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit cfabea76dbdad9282c6827326874a5ba4d20f1ed
Author: Dawid Wysakowicz <dw...@apache.org>
AuthorDate: Tue Sep 14 13:27:35 2021 +0200
[hotfix] Make it possible to force MockSourceReader to wait for initial splits assignment
---
.../api/connector/source/mocks/MockSource.java | 97 +++++++++++++++++++++-
.../connector/source/mocks/MockSourceReader.java | 42 +++++++++-
2 files changed, 133 insertions(+), 6 deletions(-)
diff --git a/flink-core/src/test/java/org/apache/flink/api/connector/source/mocks/MockSource.java b/flink-core/src/test/java/org/apache/flink/api/connector/source/mocks/MockSource.java
index 8d85963..a5b202a 100644
--- a/flink-core/src/test/java/org/apache/flink/api/connector/source/mocks/MockSource.java
+++ b/flink-core/src/test/java/org/apache/flink/api/connector/source/mocks/MockSource.java
@@ -19,6 +19,7 @@
package org.apache.flink.api.connector.source.mocks;
import org.apache.flink.api.connector.source.Boundedness;
+import org.apache.flink.api.connector.source.ReaderOutput;
import org.apache.flink.api.connector.source.Source;
import org.apache.flink.api.connector.source.SourceReader;
import org.apache.flink.api.connector.source.SourceReaderContext;
@@ -38,7 +39,7 @@ public class MockSource implements Source<Integer, MockSourceSplit, Set<MockSour
private final Boundedness boundedness;
private final int numSplits;
- private final boolean readerWaitingForMoreSplits;
+ private final MockSourceReader.WaitingForSplits readerWaitingForMoreSplits;
private final boolean readerMarkIdleOnNoSplits;
protected List<MockSourceReader> createdReaders;
@@ -51,10 +52,24 @@ public class MockSource implements Source<Integer, MockSourceSplit, Set<MockSour
int numSplits,
boolean readerWaitingForMoreSplits,
boolean readerMarkIdleOnNoSplits) {
+ this(
+ boundedness,
+ numSplits,
+ readerWaitingForMoreSplits
+ ? MockSourceReader.WaitingForSplits.WAIT_UNTIL_ALL_SPLITS_ASSIGNED
+ : MockSourceReader.WaitingForSplits.DO_NOT_WAIT_FOR_SPLITS,
+ readerMarkIdleOnNoSplits);
+ }
+
+ private MockSource(
+ Boundedness boundedness,
+ int numSplits,
+ MockSourceReader.WaitingForSplits readerWaitingForSplitsBehaviour,
+ boolean readerMarkIdleOnNoSplits) {
this.boundedness = boundedness;
this.numSplits = numSplits;
this.createdReaders = new ArrayList<>();
- this.readerWaitingForMoreSplits = readerWaitingForMoreSplits;
+ this.readerWaitingForMoreSplits = readerWaitingForSplitsBehaviour;
this.readerMarkIdleOnNoSplits = readerMarkIdleOnNoSplits;
}
@@ -94,6 +109,84 @@ public class MockSource implements Source<Integer, MockSourceSplit, Set<MockSour
return new MockSplitEnumeratorCheckpointSerializer();
}
+ public static Builder continuous(int numSplits) {
+ return new Builder(Boundedness.CONTINUOUS_UNBOUNDED, numSplits);
+ }
+
+ public static Builder bounded(int numSplits) {
+ return new Builder(Boundedness.BOUNDED, numSplits);
+ }
+
+ /** A builder for {@link MockSource}. */
+ public static class Builder {
+ private final Boundedness boundedness;
+ private final int numSplits;
+ private MockSourceReader.WaitingForSplits readerWaitingForSplitsBehaviour =
+ MockSourceReader.WaitingForSplits.WAIT_FOR_INITIAL;
+ private boolean readerMarkIdleOnNoSplits = true;
+
+ private Builder(Boundedness boundedness, int numSplits) {
+ this.boundedness = boundedness;
+ this.numSplits = numSplits;
+ }
+
+ /**
+ * Instructs the {@link MockSourceReader} to not finish if there has been no splits
+ * assignment messages yet.
+ *
+ * @see #waitUntilAllSplitsAssigned()
+ * @see #doNotWaitForSplitsAssignment()
+ */
+ public Builder waitOnlyForInitialSplits() {
+ this.readerWaitingForSplitsBehaviour =
+ MockSourceReader.WaitingForSplits.WAIT_FOR_INITIAL;
+ return this;
+ }
+
+ /**
+ * Instructs the {@link MockSourceReader} to finish only if the {@link MockSplitEnumerator}
+ * assigned all splits and there will be no more splits distributed.
+ *
+ * @see #waitUntilAllSplitsAssigned()
+ * @see #doNotWaitForSplitsAssignment()
+ */
+ public Builder waitUntilAllSplitsAssigned() {
+ this.readerWaitingForSplitsBehaviour =
+ MockSourceReader.WaitingForSplits.WAIT_UNTIL_ALL_SPLITS_ASSIGNED;
+ return this;
+ }
+
+ /**
+ * Instructs the {@link MockSourceReader} to finish irrespective if there can be splits
+ * assigned in the future or not.
+ *
+ * @see #waitUntilAllSplitsAssigned()
+ * @see #doNotWaitForSplitsAssignment()
+ */
+ public Builder doNotWaitForSplitsAssignment() {
+ this.readerWaitingForSplitsBehaviour =
+ MockSourceReader.WaitingForSplits.DO_NOT_WAIT_FOR_SPLITS;
+ return this;
+ }
+
+ /**
+ * If enabled the {@link MockSourceReader} will mark the {@link ReaderOutput#markIdle()
+ * idle} if it has no splits assigned.
+ */
+ public Builder markReaderIdleOnNoSplits(boolean enable) {
+ this.readerMarkIdleOnNoSplits = enable;
+ return this;
+ }
+
+ public MockSource build() {
+ return new MockSource(
+ boundedness,
+ numSplits,
+ readerWaitingForSplitsBehaviour,
+ readerMarkIdleOnNoSplits);
+ }
+ }
+
// --------------- methods for testing -------------
public List<MockSourceReader> getCreatedReaders() {
diff --git a/flink-core/src/test/java/org/apache/flink/api/connector/source/mocks/MockSourceReader.java b/flink-core/src/test/java/org/apache/flink/api/connector/source/mocks/MockSourceReader.java
index f871ea8..2cc084b 100644
--- a/flink-core/src/test/java/org/apache/flink/api/connector/source/mocks/MockSourceReader.java
+++ b/flink-core/src/test/java/org/apache/flink/api/connector/source/mocks/MockSourceReader.java
@@ -40,7 +40,20 @@ public class MockSourceReader implements SourceReader<Integer, MockSourceSplit>
private int currentSplitIndex = 0;
private boolean started;
private int timesClosed;
- private boolean waitingForMoreSplits;
+ private final WaitingForSplits waitingForSplitsBehaviour;
+ private SplitsAssignmentState splitsAssignmentState = SplitsAssignmentState.NO_SPLITS_ASSIGNED;
+
+ enum WaitingForSplits {
+ WAIT_FOR_INITIAL,
+ WAIT_UNTIL_ALL_SPLITS_ASSIGNED,
+ DO_NOT_WAIT_FOR_SPLITS
+ }
+
+ private enum SplitsAssignmentState {
+ NO_SPLITS_ASSIGNED,
+ INITIAL_SPLITS_ASSIGNED,
+ NO_MORE_SPLITS
+ }
@GuardedBy("this")
private CompletableFuture<Void> availableFuture;
@@ -50,10 +63,19 @@ public class MockSourceReader implements SourceReader<Integer, MockSourceSplit>
}
public MockSourceReader(boolean waitingForMoreSplits, boolean markIdleOnNoSplits) {
+ this(
+ waitingForMoreSplits
+ ? WaitingForSplits.WAIT_UNTIL_ALL_SPLITS_ASSIGNED
+ : WaitingForSplits.DO_NOT_WAIT_FOR_SPLITS,
+ markIdleOnNoSplits);
+ }
+
+ public MockSourceReader(
+ WaitingForSplits waitingForSplitsBehaviour, boolean markIdleOnNoSplits) {
this.started = false;
this.timesClosed = 0;
this.availableFuture = CompletableFuture.completedFuture(null);
- this.waitingForMoreSplits = waitingForMoreSplits;
+ this.waitingForSplitsBehaviour = waitingForSplitsBehaviour;
this.markIdleOnNoSplits = markIdleOnNoSplits;
}
@@ -64,7 +86,16 @@ public class MockSourceReader implements SourceReader<Integer, MockSourceSplit>
@Override
public InputStatus pollNext(ReaderOutput<Integer> sourceOutput) throws Exception {
- boolean finished = !waitingForMoreSplits;
+
+ if (waitingForSplitsBehaviour == WaitingForSplits.WAIT_FOR_INITIAL
+ && splitsAssignmentState == SplitsAssignmentState.NO_SPLITS_ASSIGNED) {
+ markUnavailable();
+ return InputStatus.NOTHING_AVAILABLE;
+ }
+
+ boolean finished =
+ splitsAssignmentState == SplitsAssignmentState.NO_MORE_SPLITS
+ || waitingForSplitsBehaviour == WaitingForSplits.DO_NOT_WAIT_FOR_SPLITS;
currentSplitIndex = 0;
// Find first splits with available records.
while (currentSplitIndex < assignedSplits.size()
@@ -101,13 +132,16 @@ public class MockSourceReader implements SourceReader<Integer, MockSourceSplit>
@Override
public void addSplits(List<MockSourceSplit> splits) {
+ if (splitsAssignmentState == SplitsAssignmentState.NO_SPLITS_ASSIGNED) {
+ splitsAssignmentState = SplitsAssignmentState.INITIAL_SPLITS_ASSIGNED;
+ }
assignedSplits.addAll(splits);
markAvailable();
}
@Override
public void notifyNoMoreSplits() {
- waitingForMoreSplits = false;
+ splitsAssignmentState = SplitsAssignmentState.NO_MORE_SPLITS;
markAvailable();
}