[FLINK-38733] Add new SplitterEnumerator on JdbcSource

Author: eskabetxe

flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/source/JdbcSource.java

@@ -67,9 +67,7 @@ public class JdbcSource<OUT>
                 Source<OUT, JdbcSourceSplit, JdbcSourceEnumeratorState>,
                 ResultTypeQueryable<OUT> {
 
-    private final Boundedness boundedness;
     private final TypeInformation<OUT> typeInformation;
-    private final @Nullable ContinuousUnBoundingSettings continuousUnBoundingSettings;
 
     private final Configuration configuration;
     private final SplitterEnumerator splitterEnumerator;
@@ -90,11 +88,11 @@ public class JdbcSource<OUT>
         this(
                 configuration,
                 connectionProvider,
-                new JdbcSqlSplitterEnumerator(sqlSplitEnumeratorProvider),
+                new JdbcSqlSplitterEnumerator(
+                        sqlSplitEnumeratorProvider, continuousUnBoundingSettings),
                 resultExtractor,
                 typeInformation,
-                deliveryGuarantee,
-                continuousUnBoundingSettings);
+                deliveryGuarantee);
     }
 
     JdbcSource(
@@ -103,25 +101,19 @@ public class JdbcSource<OUT>
             SplitterEnumerator splitterEnumerator,
             ResultExtractor<OUT> resultExtractor,
             TypeInformation<OUT> typeInformation,
-            @Nullable DeliveryGuarantee deliveryGuarantee,
-            @Nullable ContinuousUnBoundingSettings continuousUnBoundingSettings) {
+            @Nullable DeliveryGuarantee deliveryGuarantee) {
         this.configuration = Preconditions.checkNotNull(configuration);
         this.connectionProvider = Preconditions.checkNotNull(connectionProvider);
         this.splitterEnumerator = Preconditions.checkNotNull(splitterEnumerator);
         this.resultExtractor = Preconditions.checkNotNull(resultExtractor);
         this.deliveryGuarantee =
                 Objects.isNull(deliveryGuarantee) ? DeliveryGuarantee.NONE : deliveryGuarantee;
         this.typeInformation = Preconditions.checkNotNull(typeInformation);
-        this.continuousUnBoundingSettings = continuousUnBoundingSettings;
-        this.boundedness =
-                Objects.isNull(continuousUnBoundingSettings)
-                        ? Boundedness.BOUNDED
-                        : Boundedness.CONTINUOUS_UNBOUNDED;
     }
 
     @Override
     public Boundedness getBoundedness() {
-        return boundedness;
+        return splitterEnumerator.getBoundedness();
     }
 
     @Override
@@ -144,11 +136,7 @@ public SourceReader<OUT, JdbcSourceSplit> createReader(SourceReaderContext reade
     public SplitEnumerator<JdbcSourceSplit, JdbcSourceEnumeratorState> createEnumerator(
             SplitEnumeratorContext<JdbcSourceSplit> enumContext) throws Exception {
         return new JdbcSourceEnumerator(
-                enumContext,
-                splitterEnumerator,
-                connectionProvider,
-                continuousUnBoundingSettings,
-                new ArrayList<>());
+                enumContext, splitterEnumerator, connectionProvider, new ArrayList<>());
     }
 
     @Override
@@ -162,7 +150,6 @@ public SplitEnumerator<JdbcSourceSplit, JdbcSourceEnumeratorState> restoreEnumer
                 enumContext,
                 splitterEnumerator.restoreState(optionalUserDefinedSplitEnumeratorState),
                 connectionProvider,
-                continuousUnBoundingSettings,
                 checkpoint.getRemainingSplits());
     }
 
@@ -223,14 +210,12 @@ public boolean equals(Object o) {
             return false;
         }
         JdbcSource<?> that = (JdbcSource<?>) o;
-        return boundedness == that.boundedness
-                && Objects.equals(typeInformation, that.typeInformation)
+        return Objects.equals(typeInformation, that.typeInformation)
                 && Objects.equals(configuration, that.configuration)
                 && Objects.equals(splitterEnumerator, that.splitterEnumerator)
                 && Objects.equals(connectionProvider, that.connectionProvider)
                 && Objects.equals(resultExtractor, that.resultExtractor)
-                && deliveryGuarantee == that.deliveryGuarantee
-                && Objects.equals(continuousUnBoundingSettings, that.continuousUnBoundingSettings);
+                && deliveryGuarantee == that.deliveryGuarantee;
     }
 
     @Override
@@ -243,6 +228,6 @@ public LineageVertex getLineageVertex() {
         LineageDataset dataset =
                 LineageUtils.datasetOf(
                         nameOpt.orElse(""), namespace, Arrays.asList(defaultTypeDatasetFacet));
-        return LineageUtils.sourceLineageVertexOf(boundedness, Collections.singleton(dataset));
+        return LineageUtils.sourceLineageVertexOf(getBoundedness(), Collections.singleton(dataset));
     }
 }

flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/source/JdbcSourceBuilder.java

@@ -332,7 +332,11 @@ public JdbcSource<OUT> build() {
             }
 
             this.splitterEnumerator =
-                    getSplitter(sql, jdbcParameterValuesProvider, optionalSqlSplitEnumeratorState);
+                    getSplitter(
+                            sql,
+                            jdbcParameterValuesProvider,
+                            optionalSqlSplitEnumeratorState,
+                            continuousUnBoundingSettings);
         }
 
         return new JdbcSource<>(
@@ -341,20 +345,20 @@ public JdbcSource<OUT> build() {
                 splitterEnumerator,
                 resultExtractor,
                 typeInformation,
-                deliveryGuarantee,
-                continuousUnBoundingSettings);
+                deliveryGuarantee);
     }
 
     private SplitterEnumerator getSplitter(
             String sqlTemplate,
             JdbcParameterValuesProvider parameterProvider,
-            Serializable userDefinedState) {
+            Serializable userDefinedState,
+            @Nullable ContinuousUnBoundingSettings continuousSettings) {
         JdbcSqlSplitEnumeratorBase.Provider<?> provider =
                 new SqlTemplateSplitEnumerator.TemplateSqlSplitEnumeratorProvider()
                         .setOptionalSqlSplitEnumeratorState(userDefinedState)
                         .setSqlTemplate(sqlTemplate)
                         .setParameterValuesProvider(parameterProvider);
 
-        return new JdbcSqlSplitterEnumerator(provider);
+        return new JdbcSqlSplitterEnumerator(provider, continuousSettings);
     }
 }

flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/source/config/ContinuousUnBoundingSettings.java

@@ -31,6 +31,7 @@
  * continuous discovery and streaming mode.
  */
 @PublicEvolving
+@Deprecated
 public class ContinuousUnBoundingSettings implements Serializable {
 
     private static final long serialVersionUID = 1L;

flink-connector-jdbc-core/src/main/java/org/apache/flink/connector/jdbc/core/datastream/source/enumerator/JdbcSourceEnumerator.java

@@ -22,7 +22,6 @@
 import org.apache.flink.api.connector.source.SourceEvent;
 import org.apache.flink.api.connector.source.SplitEnumerator;
 import org.apache.flink.api.connector.source.SplitEnumeratorContext;
-import org.apache.flink.connector.jdbc.core.datastream.source.config.ContinuousUnBoundingSettings;
 import org.apache.flink.connector.jdbc.core.datastream.source.enumerator.splitter.SplitterEnumerator;
 import org.apache.flink.connector.jdbc.core.datastream.source.split.JdbcSourceSplit;
 import org.apache.flink.connector.jdbc.datasource.connections.JdbcConnectionProvider;
@@ -37,56 +36,37 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Iterator;
-import java.util.LinkedHashMap;
 import java.util.List;
-import java.util.Map;
-import java.util.Objects;
 import java.util.Optional;
+import java.util.concurrent.atomic.AtomicInteger;
 
 /** JDBC source enumerator. */
 public class JdbcSourceEnumerator
         implements SplitEnumerator<JdbcSourceSplit, JdbcSourceEnumeratorState> {
     private static final Logger LOG = LoggerFactory.getLogger(JdbcSourceEnumerator.class);
 
     private final SplitEnumeratorContext<JdbcSourceSplit> context;
-    private final Boundedness boundedness;
-    private final LinkedHashMap<Integer, String> readersAwaitingSplit;
     private final List<JdbcSourceSplit> unassigned;
     private final SplitterEnumerator splitterEnumerator;
-    private final @Nullable ContinuousUnBoundingSettings continuousUnBoundingSettings;
     private final JdbcConnectionProvider connectionProvider;
+    private final List<Integer> readersWaitingForSplits = new ArrayList<>();
+    private final AtomicInteger asyncCallsPending = new AtomicInteger(0);
 
     public JdbcSourceEnumerator(
             SplitEnumeratorContext<JdbcSourceSplit> context,
             SplitterEnumerator splitterEnumerator,
             JdbcConnectionProvider connectionProvider,
-            ContinuousUnBoundingSettings continuousUnBoundingSettings,
             List<JdbcSourceSplit> unassigned) {
         this.context = Preconditions.checkNotNull(context);
         this.splitterEnumerator = Preconditions.checkNotNull(splitterEnumerator);
-        this.continuousUnBoundingSettings = continuousUnBoundingSettings;
-        this.boundedness =
-                Objects.isNull(continuousUnBoundingSettings)
-                        ? Boundedness.BOUNDED
-                        : Boundedness.CONTINUOUS_UNBOUNDED;
         this.unassigned = Preconditions.checkNotNull(unassigned);
-        this.readersAwaitingSplit = new LinkedHashMap<>();
         this.connectionProvider = connectionProvider;
     }
 
     @Override
     public void start() {
         splitterEnumerator.start(connectionProvider);
-        if (boundedness == Boundedness.CONTINUOUS_UNBOUNDED
-                && Objects.nonNull(continuousUnBoundingSettings)) {
-            context.callAsync(
-                    () -> splitterEnumerator.enumerateSplits(() -> 1024 - unassigned.size() > 0),
-                    this::processNewSplits,
-                    continuousUnBoundingSettings.getInitialDiscoveryDelay().toMillis(),
-                    continuousUnBoundingSettings.getDiscoveryInterval().toMillis());
-        } else {
-            unassigned.addAll(splitterEnumerator.enumerateSplits());
-        }
+        preDiscoverSplits();
     }
 
     @Override
@@ -101,11 +81,18 @@ public void addReader(int subtaskId) {
 
     @Override
     public void handleSplitRequest(int subtask, @Nullable String hostname) {
-        if (boundedness == Boundedness.BOUNDED) {
-            assignSplitsForBounded(subtask, hostname);
+        if (!context.registeredReaders().containsKey(subtask)) {
+            LOG.warn("Ignoring split request from unregistered reader {}", subtask);
+            return;
+        }
+        final Optional<JdbcSourceSplit> nextSplit = getNextSplit();
+        if (nextSplit.isPresent()) {
+            context.assignSplit(nextSplit.get(), subtask);
+            LOG.info("Assigned split to subtask {} : {}", subtask, nextSplit.get());
+            preDiscoverSplits();
         } else {
-            readersAwaitingSplit.put(subtask, hostname);
-            assignSplitsForUnbounded();
+            readersWaitingForSplits.add(subtask);
+            preDiscoverSplits();
         }
     }
 
@@ -116,18 +103,12 @@ public void handleSourceEvent(int subtaskId, SourceEvent sourceEvent) {
 
     @Override
     public void addSplitsBack(List<JdbcSourceSplit> splits, int subtaskId) {
-        LOG.debug("File Source Enumerator adds splits back: {}", splits);
+        LOG.debug("Source Enumerator adds splits back: {}", splits);
         unassigned.addAll(splits);
-        if (boundedness == Boundedness.BOUNDED) {
-            context.registeredReaders()
-                    .keySet()
-                    .forEach(subTask -> assignSplitsForBounded(subTask, null));
-        } else if (boundedness == Boundedness.CONTINUOUS_UNBOUNDED) {
-            if (context.registeredReaders().containsKey(subtaskId)) {
-                readersAwaitingSplit.put(subtaskId, null);
-            }
-            assignSplitsForUnbounded();
+        if (context.registeredReaders().containsKey(subtaskId)) {
+            readersWaitingForSplits.add(subtaskId);
         }
+        preDiscoverSplits();
     }
 
     @Override
@@ -150,58 +131,65 @@ private Optional<JdbcSourceSplit> getNextSplit() {
         return Optional.of(next);
     }
 
-    private void assignSplitsForBounded(int subtask, @Nullable String hostname) {
-        if (!context.registeredReaders().containsKey(subtask)) {
-            return;
-        }
-        if (LOG.isInfoEnabled()) {
-            final String hostInfo =
-                    hostname == null ? "(no host locality info)" : "(on host '" + hostname + "')";
-            LOG.info("Subtask {} {} is requesting a Jdbc source split", subtask, hostInfo);
-        }
-        final Optional<JdbcSourceSplit> nextSplit = getNextSplit();
-        if (nextSplit.isPresent()) {
-            final JdbcSourceSplit split = nextSplit.get();
-            context.assignSplit(split, subtask);
-            LOG.info("Assigned split to subtask {} : {}", subtask, split);
-        } else {
-            context.signalNoMoreSplits(subtask);
-            LOG.info("No more splits available for subtask {}", subtask);
+    private void preDiscoverSplits() {
+        int targetParallelism = context.currentParallelism();
+        while (asyncCallsPending.get() < targetParallelism
+                && !splitterEnumerator.isAllSplitsFinished()) {
+            asyncCallsPending.incrementAndGet();
+            context.callAsync(() -> splitterEnumerator.enumerateSplits(), this::onSplitsDiscovered);
         }
+
+        signalNoMoreSplitsIfDone();
     }
 
-    private void processNewSplits(List<JdbcSourceSplit> splits, Throwable error) {
+    private void onSplitsDiscovered(List<JdbcSourceSplit> splits, Throwable error) {
+        asyncCallsPending.decrementAndGet();
         if (error != null) {
-            LOG.error("Failed to enumerate sql splits.", error);
+            LOG.error("Failed to discover splits.", error);
+            preDiscoverSplits();
             return;
         }
-        this.unassigned.addAll(splits);
 
-        assignSplitsForUnbounded();
+        if (splits != null && !splits.isEmpty()) {
+            assignOrBuffer(splits);
+            preDiscoverSplits();
+        } else if (!splitterEnumerator.isAllSplitsFinished()) {
+            preDiscoverSplits();
+        } else {
+            signalNoMoreSplitsIfDone();
+        }
     }
 
-    private void assignSplitsForUnbounded() {
-        final Iterator<Map.Entry<Integer, String>> awaitingReader =
-                readersAwaitingSplit.entrySet().iterator();
-
-        while (awaitingReader.hasNext()) {
-            final Map.Entry<Integer, String> nextAwaiting = awaitingReader.next();
-
-            // if the reader that requested another split has failed in the meantime, remove
-            // it from the list of waiting readers
-            if (!context.registeredReaders().containsKey(nextAwaiting.getKey())) {
-                awaitingReader.remove();
-                continue;
+    private void assignOrBuffer(List<JdbcSourceSplit> splits) {
+        for (JdbcSourceSplit split : splits) {
+            if (!readersWaitingForSplits.isEmpty()) {
+                int subtaskId = readersWaitingForSplits.remove(0);
+                if (context.registeredReaders().containsKey(subtaskId)) {
+                    LOG.info(
+                            "Assigning discovered split {} to waiting subtask {}",
+                            split,
+                            subtaskId);
+                    context.assignSplit(split, subtaskId);
+                } else {
+                    unassigned.add(split);
+                }
+            } else {
+                unassigned.add(split);
             }
+        }
+    }
 
-            final int awaitingSubtask = nextAwaiting.getKey();
-            final Optional<JdbcSourceSplit> nextSplit = getNextSplit();
-            if (nextSplit.isPresent()) {
-                context.assignSplit(nextSplit.get(), awaitingSubtask);
-                awaitingReader.remove();
-            } else {
-                break;
+    private void signalNoMoreSplitsIfDone() {
+        if (splitterEnumerator.getBoundedness() == Boundedness.BOUNDED
+                && asyncCallsPending.get() == 0
+                && splitterEnumerator.isAllSplitsFinished()
+                && unassigned.isEmpty()
+                && !readersWaitingForSplits.isEmpty()) {
+            for (int subtaskId : readersWaitingForSplits) {
+                context.signalNoMoreSplits(subtaskId);
+                LOG.info("No more splits available for subtask {}", subtaskId);
             }
+            readersWaitingForSplits.clear();
         }
     }
 }