[Improve][Connector-V2] Improve Paimon source split enumerator

seatunnel-connectors-v2/connector-paimon/src/main/java/org/apache/seatunnel/connectors/seatunnel/paimon/source/PaimonSourceSplitEnumerator.java

@@ -19,15 +19,20 @@
 
 import org.apache.seatunnel.api.source.SourceSplitEnumerator;
 
+import org.apache.commons.collections.map.HashedMap;
 import org.apache.paimon.table.Table;
 import org.apache.paimon.table.source.Split;
 
 import lombok.extern.slf4j.Slf4j;
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 
@@ -39,36 +44,51 @@ public class PaimonSourceSplitEnumerator
     /** Source split enumerator context */
     private final Context<PaimonSourceSplit> context;
 
-    /** The splits that has assigned */
-    private final Set<PaimonSourceSplit> assignedSplit;
+    private Map<Integer, List<PaimonSourceSplit>> pendingSplit;
 
-    /** The splits that have not assigned */
-    private Set<PaimonSourceSplit> pendingSplit;
+    private volatile boolean shouldEnumerate;
+
+    private final Object stateLock = new Object();
 
     /** The table that wants to read */
     private final Table table;
 
     public PaimonSourceSplitEnumerator(Context<PaimonSourceSplit> context, Table table) {
-        this.context = context;
-        this.table = table;
-        this.assignedSplit = new HashSet<>();
+        this(context, table, null);
     }
 
     public PaimonSourceSplitEnumerator(
             Context<PaimonSourceSplit> context, Table table, PaimonSourceState sourceState) {
         this.context = context;
         this.table = table;
-        this.assignedSplit = sourceState.getAssignedSplits();
+        this.pendingSplit = new HashMap<>();
+        this.shouldEnumerate = sourceState == null;
+        if (sourceState != null) {
+            this.shouldEnumerate = sourceState.isShouldEnumerate();
+            this.pendingSplit.putAll(sourceState.getPendingSplits());
+        }
     }
 
     @Override
     public void open() {
-        this.pendingSplit = new HashSet<>();
+        this.pendingSplit = new HashedMap();
     }
 
     @Override
     public void run() throws Exception {
-        // do nothing
+        Set<Integer> readers = context.registeredReaders();
+        if (shouldEnumerate) {
+            Set<PaimonSourceSplit> newSplits = getTableSplits();
+            synchronized (stateLock) {
+                addPendingSplit(newSplits);
+                shouldEnumerate = false;
+            }
+
+            assignSplit(readers);
+        }
+        log.debug(
+                "No more splits to assign." + " Sending NoMoreSplitsEvent to reader {}.", readers);
+        readers.forEach(context::signalNoMoreSplits);
     }
 
     @Override
@@ -79,8 +99,8 @@ public void close() throws IOException {
     @Override
     public void addSplitsBack(List<PaimonSourceSplit> splits, int subtaskId) {
         if (!splits.isEmpty()) {
-            pendingSplit.addAll(splits);
-            assignSplit(subtaskId);
+            addPendingSplit(splits);
+            assignSplit(Collections.singletonList(subtaskId));
         }
     }
 
@@ -91,13 +111,17 @@ public int currentUnassignedSplitSize() {
 
     @Override
     public void registerReader(int subtaskId) {
-        pendingSplit = getTableSplits();
-        assignSplit(subtaskId);
+        log.debug("Register reader {} to PaimonSourceSplitEnumerator.", subtaskId);
+        if (!pendingSplit.isEmpty()) {
+            assignSplit(Collections.singletonList(subtaskId));
+        }
     }
 
     @Override
     public PaimonSourceState snapshotState(long checkpointId) throws Exception {
-        return new PaimonSourceState(assignedSplit);
+        synchronized (stateLock) {
+            return new PaimonSourceState(pendingSplit, shouldEnumerate);
+        }
     }
 
     @Override
@@ -110,36 +134,41 @@ public void handleSplitRequest(int subtaskId) {
         // do nothing
     }
 
+    private void addPendingSplit(Collection<PaimonSourceSplit> splits) {
+        int readerCount = context.currentParallelism();
+        for (PaimonSourceSplit split : splits) {
+            int ownerReader = getSplitOwner(split.splitId(), readerCount);
+            log.info("Assigning {} to {} reader.", split.getSplit().toString(), ownerReader);
+            pendingSplit.computeIfAbsent(ownerReader, r -> new ArrayList<>()).add(split);
+        }
+    }
+
     /** Assign split by reader task id */
-    private void assignSplit(int taskId) {
-        ArrayList<PaimonSourceSplit> currentTaskSplits = new ArrayList<>();
-        if (context.currentParallelism() == 1) {
-            // if parallelism == 1, we should assign all the splits to reader
-            currentTaskSplits.addAll(pendingSplit);
-        } else {
-            // if parallelism > 1, according to hashCode of split's id to determine whether to
-            // allocate the current task
-            for (PaimonSourceSplit fileSourceSplit : pendingSplit) {
-                final int splitOwner =
-                        getSplitOwner(fileSourceSplit.splitId(), context.currentParallelism());
-                if (splitOwner == taskId) {
-                    currentTaskSplits.add(fileSourceSplit);
+    private void assignSplit(Collection<Integer> readers) {
+
+        log.debug("Assign pendingSplits to readers {}", readers);
+
+        for (int reader : readers) {
+            List<PaimonSourceSplit> assignmentForReader = pendingSplit.remove(reader);
+            if (assignmentForReader != null && !assignmentForReader.isEmpty()) {
+                log.info(
+                        "Assign splits {} to reader {}",
+                        assignmentForReader.stream()
+                                .map(p -> p.getSplit().toString())
+                                .collect(Collectors.joining(",")),
+                        reader);
+                try {
+                    context.assignSplit(reader, assignmentForReader);
+                } catch (Exception e) {
+                    log.error(
+                            "Failed to assign splits {} to reader {}",
+                            assignmentForReader,
+                            reader,
+                            e);
+                    pendingSplit.put(reader, assignmentForReader);
                 }
             }
         }
-        // assign splits
-        context.assignSplit(taskId, currentTaskSplits);
-        // save the state of assigned splits
-        assignedSplit.addAll(currentTaskSplits);
-        // remove the assigned splits from pending splits
-        currentTaskSplits.forEach(split -> pendingSplit.remove(split));
-        log.info(
-                "SubTask {} is assigned to [{}]",
-                taskId,
-                currentTaskSplits.stream()
-                        .map(PaimonSourceSplit::splitId)
-                        .collect(Collectors.joining(",")));
-        context.signalNoMoreSplits(taskId);
     }
 
     /** Get all splits of table */

seatunnel-connectors-v2/connector-paimon/src/main/java/org/apache/seatunnel/connectors/seatunnel/paimon/source/PaimonSourceState.java

@@ -17,21 +17,19 @@
 
 package org.apache.seatunnel.connectors.seatunnel.paimon.source;
 
+import lombok.AllArgsConstructor;
+import lombok.Getter;
+
 import java.io.Serializable;
-import java.util.Set;
+import java.util.List;
+import java.util.Map;
 
+@AllArgsConstructor
+@Getter
 /** Paimon connector source state, saves the splits has assigned to readers. */
 public class PaimonSourceState implements Serializable {
 
     private static final long serialVersionUID = 1L;
-
-    private final Set<PaimonSourceSplit> assignedSplits;
-
-    public PaimonSourceState(Set<PaimonSourceSplit> assignedSplits) {
-        this.assignedSplits = assignedSplits;
-    }
-
-    public Set<PaimonSourceSplit> getAssignedSplits() {
-        return assignedSplits;
-    }
+    private final Map<Integer, List<PaimonSourceSplit>> pendingSplits;
+    private boolean shouldEnumerate;
 }