[feature][connector][fake] Support mutil splits for fake source connector (#2974)

* feature Support more than splits and parallelism for fake connector

close #2961

* test parallelism

* Not practical average

* fix a bug

* remove split rowNum

* [Feature][Connector-V2] Improve class attribute

* fine-tuning

* Remove useless blank line.

* support STREAMING

* [feature][connector][fake] Support mutil-splits for fake source

* [connector][fake] Make sure the bounded flow is correct

* Update seatunnel-connectors-v2/connector-fake/src/main/java/org/apache/seatunnel/connectors/seatunnel/fake/source/FakeSourceReader.java

Co-authored-by: hailin0 <hailin088@gmail.com>

Co-authored-by: TyrantLucifer <tyrantlucifer@apache.org>
Co-authored-by: Zongwen Li <zongwen@apache.org>
Co-authored-by: hailin0 <hailin088@gmail.com>

Author: 4 people

seatunnel-connectors-v2/connector-fake/src/main/java/org/apache/seatunnel/connectors/seatunnel/fake/source/FakeConfig.java renamed to seatunnel-connectors-v2/connector-fake/src/main/java/org/apache/seatunnel/connectors/seatunnel/fake/config/FakeConfig.java

@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-package org.apache.seatunnel.connectors.seatunnel.fake.source;
+package org.apache.seatunnel.connectors.seatunnel.fake.config;
 
 import org.apache.seatunnel.shade.com.typesafe.config.Config;
 
@@ -27,11 +27,11 @@
 @Builder
 @Getter
 public class FakeConfig implements Serializable {
-    private static final String ROW_NUM = "row.num";
-    private static final String MAP_SIZE = "map.size";
-    private static final String ARRAY_SIZE = "array.size";
-    private static final String BYTES_LENGTH = "bytes.length";
-    private static final String STRING_LENGTH = "string.length";
+    public static final String ROW_NUM = "row.num";
+    public static final String MAP_SIZE = "map.size";
+    public static final String ARRAY_SIZE = "array.size";
+    public static final String BYTES_LENGTH = "bytes.length";
+    public static final String STRING_LENGTH = "string.length";
     private static final int DEFAULT_ROW_NUM = 5;
     private static final int DEFAULT_MAP_SIZE = 5;
     private static final int DEFAULT_ARRAY_SIZE = 5;

seatunnel-connectors-v2/connector-fake/src/main/java/org/apache/seatunnel/connectors/seatunnel/fake/source/FakeDataGenerator.java

@@ -25,6 +25,7 @@
 import org.apache.seatunnel.api.table.type.SeaTunnelRow;
 import org.apache.seatunnel.api.table.type.SeaTunnelRowType;
 import org.apache.seatunnel.connectors.seatunnel.common.schema.SeaTunnelSchema;
+import org.apache.seatunnel.connectors.seatunnel.fake.config.FakeConfig;
 
 import org.apache.commons.lang3.RandomStringUtils;
 import org.apache.commons.lang3.RandomUtils;

seatunnel-connectors-v2/connector-fake/src/main/java/org/apache/seatunnel/connectors/seatunnel/fake/source/FakeSource.java

@@ -20,22 +20,23 @@
 import org.apache.seatunnel.api.common.JobContext;
 import org.apache.seatunnel.api.source.Boundedness;
 import org.apache.seatunnel.api.source.SeaTunnelSource;
+import org.apache.seatunnel.api.source.SourceReader;
+import org.apache.seatunnel.api.source.SourceSplitEnumerator;
 import org.apache.seatunnel.api.table.type.SeaTunnelRow;
 import org.apache.seatunnel.api.table.type.SeaTunnelRowType;
 import org.apache.seatunnel.common.constants.JobMode;
 import org.apache.seatunnel.connectors.seatunnel.common.schema.SeaTunnelSchema;
-import org.apache.seatunnel.connectors.seatunnel.common.source.AbstractSingleSplitReader;
-import org.apache.seatunnel.connectors.seatunnel.common.source.AbstractSingleSplitSource;
-import org.apache.seatunnel.connectors.seatunnel.common.source.SingleSplitReaderContext;
+import org.apache.seatunnel.connectors.seatunnel.fake.config.FakeConfig;
 
 import org.apache.seatunnel.shade.com.typesafe.config.Config;
 
 import com.google.auto.service.AutoService;
 
+import java.io.Serializable;
+
 @AutoService(SeaTunnelSource.class)
-public class FakeSource extends AbstractSingleSplitSource<SeaTunnelRow> {
+public class FakeSource implements SeaTunnelSource<SeaTunnelRow, FakeSourceSplit, Serializable> {
 
-    private Config pluginConfig;
     private JobContext jobContext;
     private SeaTunnelSchema schema;
     private FakeConfig fakeConfig;
@@ -51,7 +52,17 @@ public SeaTunnelRowType getProducedType() {
     }
 
     @Override
-    public AbstractSingleSplitReader<SeaTunnelRow> createReader(SingleSplitReaderContext readerContext) throws Exception {
+    public SourceSplitEnumerator<FakeSourceSplit, Serializable> createEnumerator(SourceSplitEnumerator.Context<FakeSourceSplit> enumeratorContext) throws Exception {
+        return new FakeSourceSplitEnumerator(enumeratorContext);
+    }
+
+    @Override
+    public SourceSplitEnumerator<FakeSourceSplit, Serializable> restoreEnumerator(SourceSplitEnumerator.Context<FakeSourceSplit> enumeratorContext, Serializable checkpointState) throws Exception {
+        return new FakeSourceSplitEnumerator(enumeratorContext);
+    }
+
+    @Override
+    public SourceReader<SeaTunnelRow, FakeSourceSplit> createReader(SourceReader.Context readerContext) throws Exception {
         return new FakeSourceReader(readerContext, new FakeDataGenerator(schema, fakeConfig));
     }
 
@@ -62,7 +73,6 @@ public String getPluginName() {
 
     @Override
     public void prepare(Config pluginConfig) {
-        this.pluginConfig = pluginConfig;
         assert pluginConfig.hasPath(FakeDataGenerator.SCHEMA);
         this.schema = SeaTunnelSchema.buildWithConfig(pluginConfig.getConfig(FakeDataGenerator.SCHEMA));
         this.fakeConfig = FakeConfig.buildWithConfig(pluginConfig);

seatunnel-connectors-v2/connector-fake/src/main/java/org/apache/seatunnel/connectors/seatunnel/fake/source/FakeSourceReader.java

@@ -19,24 +19,27 @@
 
 import org.apache.seatunnel.api.source.Boundedness;
 import org.apache.seatunnel.api.source.Collector;
+import org.apache.seatunnel.api.source.SourceReader;
 import org.apache.seatunnel.api.table.type.SeaTunnelRow;
-import org.apache.seatunnel.connectors.seatunnel.common.source.AbstractSingleSplitReader;
-import org.apache.seatunnel.connectors.seatunnel.common.source.SingleSplitReaderContext;
 
 import lombok.extern.slf4j.Slf4j;
 
+import java.util.ArrayList;
+import java.util.Deque;
+import java.util.LinkedList;
 import java.util.List;
 
 @Slf4j
-public class FakeSourceReader extends AbstractSingleSplitReader<SeaTunnelRow> {
-
-    private final SingleSplitReaderContext context;
+public class FakeSourceReader implements SourceReader<SeaTunnelRow, FakeSourceSplit> {
 
+    private final SourceReader.Context context;
+    private final Deque<FakeSourceSplit> splits = new LinkedList<>();
     private final FakeDataGenerator fakeDataGenerator;
+    boolean noMoreSplit;
 
-    public FakeSourceReader(SingleSplitReaderContext context, FakeDataGenerator randomData) {
+    public FakeSourceReader(SourceReader.Context context, FakeDataGenerator fakeDataGenerator) {
         this.context = context;
-        this.fakeDataGenerator = randomData;
+        this.fakeDataGenerator = fakeDataGenerator;
     }
 
     @Override
@@ -52,16 +55,46 @@ public void close() {
     @Override
     @SuppressWarnings("magicnumber")
     public void pollNext(Collector<SeaTunnelRow> output) throws InterruptedException {
-        // Generate a random number of rows to emit.
-        List<SeaTunnelRow> seaTunnelRows = fakeDataGenerator.generateFakedRows();
-        for (SeaTunnelRow seaTunnelRow : seaTunnelRows) {
-            output.collect(seaTunnelRow);
-        }
-        if (Boundedness.BOUNDED.equals(context.getBoundedness())) {
-            // signal to the source that we have reached the end of the data.
-            log.info("Closed the bounded fake source");
-            context.signalNoMoreElement();
+        synchronized (output.getCheckpointLock()) {
+            FakeSourceSplit split = splits.poll();
+            if (null != split) {
+                // Generate a random number of rows to emit.
+                List<SeaTunnelRow> seaTunnelRows = fakeDataGenerator.generateFakedRows();
+                for (SeaTunnelRow seaTunnelRow : seaTunnelRows) {
+                    output.collect(seaTunnelRow);
+                }
+            } else {
+                if (noMoreSplit && Boundedness.BOUNDED.equals(context.getBoundedness())) {
+                    // signal to the source that we have reached the end of the data.
+                    log.info("Closed the bounded fake source");
+                    context.signalNoMoreElement();
+                }
+                if (!noMoreSplit) {
+                    log.info("wait split!");
+                }
+            }
+
         }
         Thread.sleep(1000L);
     }
+
+    @Override
+    public List<FakeSourceSplit> snapshotState(long checkpointId) throws Exception {
+        return new ArrayList<>(splits);
+    }
+
+    @Override
+    public void addSplits(List<FakeSourceSplit> splits) {
+        this.splits.addAll(splits);
+    }
+
+    @Override
+    public void handleNoMoreSplits() {
+        noMoreSplit = true;
+    }
+
+    @Override
+    public void notifyCheckpointComplete(long checkpointId) throws Exception {
+
+    }
 }

seatunnel-connectors-v2/connector-fake/src/main/java/org/apache/seatunnel/connectors/seatunnel/fake/source/FakeSourceSplit.java

@@ -0,0 +1,34 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.seatunnel.connectors.seatunnel.fake.source;
+
+import org.apache.seatunnel.api.source.SourceSplit;
+
+import lombok.AllArgsConstructor;
+import lombok.Data;
+
+@Data
+@AllArgsConstructor
+public class FakeSourceSplit implements SourceSplit {
+    private int splitId;
+
+    @Override
+    public String splitId() {
+        return String.valueOf(splitId);
+    }
+}

seatunnel-connectors-v2/connector-fake/src/main/java/org/apache/seatunnel/connectors/seatunnel/fake/source/FakeSourceSplitEnumerator.java

@@ -0,0 +1,122 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.seatunnel.connectors.seatunnel.fake.source;
+
+import org.apache.seatunnel.api.source.SourceSplitEnumerator;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+public class FakeSourceSplitEnumerator implements SourceSplitEnumerator<FakeSourceSplit, Serializable> {
+
+    private static final Logger LOG = LoggerFactory.getLogger(FakeSourceSplitEnumerator.class);
+    private final SourceSplitEnumerator.Context<FakeSourceSplit> enumeratorContext;
+    private final Map<Integer, Set<FakeSourceSplit>> pendingSplits;
+
+    public FakeSourceSplitEnumerator(SourceSplitEnumerator.Context<FakeSourceSplit> enumeratorContext) {
+        this.enumeratorContext = enumeratorContext;
+        this.pendingSplits = new HashMap<>();
+    }
+
+    @Override
+    public void open() {
+        // No connection needs to be opened
+    }
+
+    @Override
+    public void run() throws Exception {
+        discoverySplits();
+        assignPendingSplits();
+    }
+
+    @Override
+    public void close() throws IOException {
+        // nothing
+    }
+
+    @Override
+    public void addSplitsBack(List<FakeSourceSplit> splits, int subtaskId) {
+
+    }
+
+    @Override
+    public int currentUnassignedSplitSize() {
+        return 0;
+    }
+
+    @Override
+    public void handleSplitRequest(int subtaskId) {
+
+    }
+
+    @Override
+    public void registerReader(int subtaskId) {
+        // nothing
+    }
+
+    @Override
+    public Serializable snapshotState(long checkpointId) throws Exception {
+        return null;
+    }
+
+    @Override
+    public void notifyCheckpointComplete(long checkpointId) throws Exception {
+
+    }
+
+    private void discoverySplits() {
+        List<FakeSourceSplit> allSplit = new ArrayList<>();
+        LOG.info("Starting to calculate splits.");
+        int numReaders = enumeratorContext.currentParallelism();
+        for (int i = 0; i < numReaders; i++) {
+            allSplit.add(new FakeSourceSplit(i));
+        }
+        for (FakeSourceSplit split : allSplit) {
+            int ownerReader = split.getSplitId() % numReaders;
+            pendingSplits.computeIfAbsent(ownerReader, r -> new HashSet<>())
+                .add(split);
+        }
+        LOG.debug("Assigned {} to {} readers.", allSplit, numReaders);
+        LOG.info("Calculated splits successfully, the size of splits is {}.", allSplit.size());
+    }
+
+    private void assignPendingSplits() {
+        // Check if there's any pending splits for given readers
+        for (int pendingReader : enumeratorContext.registeredReaders()) {
+            // Remove pending assignment for the reader
+            final Set<FakeSourceSplit> pendingAssignmentForReader =
+                pendingSplits.remove(pendingReader);
+
+            if (pendingAssignmentForReader != null && !pendingAssignmentForReader.isEmpty()) {
+                // Assign pending splits to reader
+                LOG.info("Assigning splits to readers {}", pendingAssignmentForReader);
+                enumeratorContext.assignSplit(pendingReader, new ArrayList<>(pendingAssignmentForReader));
+                enumeratorContext.signalNoMoreSplits(pendingReader);
+            }
+        }
+    }
+}

seatunnel-connectors-v2/connector-fake/src/main/java/org/apache/seatunnel/connectors/seatunnel/fake/state/FakeSourceState.java

@@ -0,0 +1,23 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.seatunnel.connectors.seatunnel.fake.state;
+
+import java.io.Serializable;
+
+public class FakeSourceState implements Serializable {
+}

seatunnel-connectors-v2/connector-fake/src/test/java/org.apache.seatunnel.connectors.seatunnel.fake.source/FakeDataGeneratorTest.java

@@ -20,6 +20,7 @@
 import org.apache.seatunnel.api.table.type.SeaTunnelRow;
 import org.apache.seatunnel.api.table.type.SeaTunnelRowType;
 import org.apache.seatunnel.connectors.seatunnel.common.schema.SeaTunnelSchema;
+import org.apache.seatunnel.connectors.seatunnel.fake.config.FakeConfig;
 
 import org.apache.seatunnel.shade.com.typesafe.config.Config;
 import org.apache.seatunnel.shade.com.typesafe.config.ConfigFactory;

seatunnel-examples/seatunnel-flink-connector-v2-example/src/main/resources/examples/fake_to_console.conf

@@ -20,8 +20,8 @@
 
 env {
   # You can set flink configuration here
-  execution.parallelism = 1
-  #job.mode = "BATCH"
+  execution.parallelism = 2
+  job.mode = "STREAMING"
   #execution.checkpoint.interval = 10000
   #execution.checkpoint.data-uri = "hdfs://localhost:9000/checkpoint"
 }