[FLINK-20194] Change SourceReaderBase.onSplitFinished() to take a map…

… of SplitId -> SplitState.
bbejeck · Nov 24, 2020 · 2e2f637 · 2e2f637
1 parent 1a9212c
commit 2e2f637
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 12 deletions.
diff --git a/...or-base/src/main/java/org/apache/flink/connector/base/source/reader/SourceReaderBase.java b/...or-base/src/main/java/org/apache/flink/connector/base/source/reader/SourceReaderBase.java
@@ -36,7 +36,6 @@
 import javax.annotation.Nullable;
 
 import java.util.ArrayList;
-import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -172,11 +171,12 @@ private void finishCurrentFetch(final RecordsWithSplitIds<E> fetch, final Reader
  final Set<String> finishedSplits = fetch.finishedSplits();
  if (!finishedSplits.isEmpty()) {
  LOG.info("Finished reading split(s) {}", finishedSplits);
+ Map<String, SplitStateT> stateOfFinishedSplits = new HashMap<>();
  for (String finishedSplitId : finishedSplits) {
- splitStates.remove(finishedSplitId);
+ stateOfFinishedSplits.put(finishedSplitId, splitStates.remove(finishedSplitId).state);
  output.releaseOutputForSplit(finishedSplitId);
  }
- onSplitFinished(finishedSplits);
+ onSplitFinished(stateOfFinishedSplits);
  }
 
  fetch.recycle();
@@ -251,7 +251,7 @@ public int getNumberOfCurrentlyAssignedSplits() {
  /**
  * Handles the finished splits to clean the state if needed.
  */
- protected abstract void onSplitFinished(Collection<String> finishedSplitIds);
+ protected abstract void onSplitFinished(Map<String, SplitStateT> finishedSplitIds);
 
  /**
  * When new splits are added to the reader. The initialize the state of the new splits.

diff --git a/...ase/src/test/java/org/apache/flink/connector/base/source/reader/SourceReaderBaseTest.java b/...ase/src/test/java/org/apache/flink/connector/base/source/reader/SourceReaderBaseTest.java
@@ -44,9 +44,9 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Supplier;
 
@@ -296,7 +296,7 @@ private Configuration getConfig() {
  public void notifyCheckpointComplete(long checkpointId) {}
 
  @Override
- protected void onSplitFinished(Collection<String> finishedSplitIds) {}
+ protected void onSplitFinished(Map<String, TestingSourceSplit> finishedSplitIds) {}
 
  @Override
  protected TestingSourceSplit initializedState(TestingSourceSplit split) {

diff --git a/...e/src/test/java/org/apache/flink/connector/base/source/reader/mocks/MockSourceReader.java b/...e/src/test/java/org/apache/flink/connector/base/source/reader/mocks/MockSourceReader.java
@@ -27,7 +27,7 @@
 import org.apache.flink.connector.base.source.reader.splitreader.SplitReader;
 import org.apache.flink.connector.base.source.reader.synchronization.FutureCompletingBlockingQueue;
 
-import java.util.Collection;
+import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Supplier;
 
@@ -52,7 +52,7 @@ public MockSourceReader(FutureCompletingBlockingQueue<RecordsWithSplitIds<int[]>
  }
 
  @Override
- protected void onSplitFinished(Collection<String> finishedSplitIds) {
+ protected void onSplitFinished(Map<String, AtomicInteger> finishedSplitIds) {
 
  }
 

diff --git a/...nector-files/src/main/java/org/apache/flink/connector/file/src/impl/FileSourceReader.java b/...nector-files/src/main/java/org/apache/flink/connector/file/src/impl/FileSourceReader.java
@@ -28,7 +28,7 @@
 import org.apache.flink.connector.file.src.reader.BulkFormat;
 import org.apache.flink.connector.file.src.util.RecordAndPosition;
 
-import java.util.Collection;
+import java.util.Map;
 
 /**
  * A {@link SourceReader} that read records from {@link FileSourceSplit}.
@@ -54,7 +54,7 @@ public void start() {
  }
 
  @Override
- protected void onSplitFinished(Collection<String> finishedSplitIds) {
+ protected void onSplitFinished(Map<String, FileSourceSplitState<SplitT>> finishedSplitIds) {
  context.sendSplitRequest();
  }
 

diff --git a/...kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/KafkaSourceReader.java b/...kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/KafkaSourceReader.java
@@ -35,7 +35,6 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -67,7 +66,7 @@ public KafkaSourceReader(
  }
 
  @Override
- protected void onSplitFinished(Collection<String> finishedSplitIds) {
+ protected void onSplitFinished(Map<String, KafkaPartitionSplitState> finishedSplitIds) {
 
  }