This is an automated email from the ASF dual-hosted git repository.
davidzollo pushed a commit to branch dev
in repository https://gitbox.apache.org/repos/asf/seatunnel.git
The following commit(s) were added to refs/heads/dev by this push:
new 1fe15cbb84 [Improve][Connector-V2] Improve source split round-robin
assignment for InfluxDB IoTDB Kudu and TDengine (#11573)
1fe15cbb84 is described below
commit 1fe15cbb847d7dc0a000ea035deec7950c0d8bba
Author: JeremyXin <[email protected]>
AuthorDate: Tue Jul 28 09:29:14 2026 +0800
[Improve][Connector-V2] Improve source split round-robin assignment for
InfluxDB IoTDB Kudu and TDengine (#11573)
---
.../source/InfluxDBSourceSplitEnumerator.java | 38 ++++-
.../influxdb/state/InfluxDBSourceState.java | 1 +
.../source/InfluxDBSourceSplitEnumeratorTest.java | 158 ++++++++++++++++++
.../source/IoTDBv2SourceSplitEnumerator.java | 38 ++++-
.../iotdbv2/state/IoTDBv2SourceState.java | 1 +
.../source/IoTDBv2SourceSplitEnumeratorTest.java | 158 ++++++++++++++++++
.../iotdb/source/IoTDBSourceSplitEnumerator.java | 36 +++-
.../seatunnel/iotdb/state/IoTDBSourceState.java | 1 +
.../source/IoTDBSourceSplitEnumeratorTest.java | 153 +++++++++++++++++
.../kudu/source/KuduSourceSplitEnumerator.java | 23 ++-
.../seatunnel/kudu/state/KuduSourceState.java | 1 +
.../kudu/source/KuduSourceSplitEnumeratorTest.java | 185 +++++++++++++++++++++
.../source/TDengineSourceSplitEnumerator.java | 35 +++-
.../tdengine/state/TDengineSourceState.java | 1 +
.../source/TDengineSourceSplitEnumeratorTest.java | 164 ++++++++++++++++++
15 files changed, 959 insertions(+), 34 deletions(-)
diff --git
a/seatunnel-connectors-v2/connector-influxdb/src/main/java/org/apache/seatunnel/connectors/seatunnel/influxdb/source/InfluxDBSourceSplitEnumerator.java
b/seatunnel-connectors-v2/connector-influxdb/src/main/java/org/apache/seatunnel/connectors/seatunnel/influxdb/source/InfluxDBSourceSplitEnumerator.java
index 694765a9fa..5e06afda18 100644
---
a/seatunnel-connectors-v2/connector-influxdb/src/main/java/org/apache/seatunnel/connectors/seatunnel/influxdb/source/InfluxDBSourceSplitEnumerator.java
+++
b/seatunnel-connectors-v2/connector-influxdb/src/main/java/org/apache/seatunnel/connectors/seatunnel/influxdb/source/InfluxDBSourceSplitEnumerator.java
@@ -31,11 +31,14 @@ import lombok.extern.slf4j.Slf4j;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
+import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
@Slf4j
public class InfluxDBSourceSplitEnumerator
@@ -44,6 +47,7 @@ public class InfluxDBSourceSplitEnumerator
private final Context<InfluxDBSourceSplit> context;
private final Map<Integer, List<InfluxDBSourceSplit>> pendingSplit;
private final Object stateLock = new Object();
+ private final AtomicInteger assignCount = new AtomicInteger(0);
private volatile boolean shouldEnumerate;
public InfluxDBSourceSplitEnumerator(
@@ -62,6 +66,7 @@ public class InfluxDBSourceSplitEnumerator
if (sourceState != null) {
this.shouldEnumerate = sourceState.isShouldEnumerate();
this.pendingSplit.putAll(sourceState.getPendingSplit());
+ this.assignCount.set(sourceState.getAssignCount());
}
}
@@ -85,11 +90,18 @@ public class InfluxDBSourceSplitEnumerator
}
@Override
- public void addSplitsBack(List splits, int subtaskId) {
+ public void addSplitsBack(List<InfluxDBSourceSplit> splits, int subtaskId)
{
log.debug("Add back splits {} to InfluxDBSourceSplitEnumerator.",
splits);
if (!splits.isEmpty()) {
- addPendingSplit(splits);
- assignSplit(Collections.singletonList(subtaskId));
+ addPendingSplit(splits, subtaskId);
+ if (context.registeredReaders().contains(subtaskId)) {
+ assignSplit(Collections.singletonList(subtaskId));
+ } else {
+ log.warn(
+ "Reader {} is not registered. Pending splits {} are
not assigned.",
+ subtaskId,
+ splits);
+ }
}
}
@@ -109,7 +121,7 @@ public class InfluxDBSourceSplitEnumerator
@Override
public InfluxDBSourceState snapshotState(long checkpointId) {
synchronized (stateLock) {
- return new InfluxDBSourceState(shouldEnumerate, pendingSplit);
+ return new InfluxDBSourceState(shouldEnumerate, pendingSplit,
assignCount.get());
}
}
@@ -181,13 +193,23 @@ public class InfluxDBSourceSplitEnumerator
private void addPendingSplit(Collection<InfluxDBSourceSplit> splits) {
int readerCount = context.currentParallelism();
- for (InfluxDBSourceSplit split : splits) {
- int ownerReader = getSplitOwner(split.splitId(), readerCount);
+
+ List<InfluxDBSourceSplit> sortedSplits =
+ splits.stream()
+
.sorted(Comparator.comparing(InfluxDBSourceSplit::splitId))
+ .collect(Collectors.toList());
+
+ for (InfluxDBSourceSplit split : sortedSplits) {
+ int ownerReader = getSplitOwner(assignCount.getAndIncrement(),
readerCount);
log.info("Assigning {} to {} reader.", split, ownerReader);
pendingSplit.computeIfAbsent(ownerReader, r -> new
ArrayList<>()).add(split);
}
}
+ private void addPendingSplit(Collection<InfluxDBSourceSplit> splits, int
ownerReader) {
+ pendingSplit.computeIfAbsent(ownerReader, r -> new
ArrayList<>()).addAll(splits);
+ }
+
private void assignSplit(Collection<Integer> readers) {
log.debug("Assign pendingSplits to readers {}", readers);
@@ -209,8 +231,8 @@ public class InfluxDBSourceSplitEnumerator
}
}
- private static int getSplitOwner(String tp, int numReaders) {
- return (tp.hashCode() & Integer.MAX_VALUE) % numReaders;
+ private static int getSplitOwner(int currentAssignCount, int numReaders) {
+ return currentAssignCount % numReaders;
}
@Override
diff --git
a/seatunnel-connectors-v2/connector-influxdb/src/main/java/org/apache/seatunnel/connectors/seatunnel/influxdb/state/InfluxDBSourceState.java
b/seatunnel-connectors-v2/connector-influxdb/src/main/java/org/apache/seatunnel/connectors/seatunnel/influxdb/state/InfluxDBSourceState.java
index 29b0dc2317..4a7d11abbf 100644
---
a/seatunnel-connectors-v2/connector-influxdb/src/main/java/org/apache/seatunnel/connectors/seatunnel/influxdb/state/InfluxDBSourceState.java
+++
b/seatunnel-connectors-v2/connector-influxdb/src/main/java/org/apache/seatunnel/connectors/seatunnel/influxdb/state/InfluxDBSourceState.java
@@ -33,4 +33,5 @@ public class InfluxDBSourceState implements Serializable {
private static final long serialVersionUID = 7132198105704653582L;
private boolean shouldEnumerate;
private Map<Integer, List<InfluxDBSourceSplit>> pendingSplit;
+ private int assignCount;
}
diff --git
a/seatunnel-connectors-v2/connector-influxdb/src/test/java/org/apache/seatunnel/connectors/seatunnel/influxdb/source/InfluxDBSourceSplitEnumeratorTest.java
b/seatunnel-connectors-v2/connector-influxdb/src/test/java/org/apache/seatunnel/connectors/seatunnel/influxdb/source/InfluxDBSourceSplitEnumeratorTest.java
new file mode 100644
index 0000000000..4b787d4bff
--- /dev/null
+++
b/seatunnel-connectors-v2/connector-influxdb/src/test/java/org/apache/seatunnel/connectors/seatunnel/influxdb/source/InfluxDBSourceSplitEnumeratorTest.java
@@ -0,0 +1,158 @@
+/*
+ * 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.influxdb.source;
+
+import org.apache.seatunnel.api.common.metrics.MetricsContext;
+import org.apache.seatunnel.api.event.EventListener;
+import org.apache.seatunnel.api.source.SourceEvent;
+import org.apache.seatunnel.api.source.SourceSplitEnumerator;
+import
org.apache.seatunnel.connectors.seatunnel.influxdb.state.InfluxDBSourceState;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.lang.reflect.Method;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+public class InfluxDBSourceSplitEnumeratorTest {
+
+ @Test
+ public void shouldBalanceSplitsEvenlyAcrossReaders() throws Exception {
+ TestingContext context = new TestingContext(4);
+ InfluxDBSourceSplitEnumerator enumerator =
+ new InfluxDBSourceSplitEnumerator(context, null, null);
+
+ invokeAddPendingSplit(enumerator, buildSplits(10));
+
+ InfluxDBSourceState state = enumerator.snapshotState(1L);
+ Map<Integer, List<InfluxDBSourceSplit>> pendingSplits =
state.getPendingSplit();
+
+ Assertions.assertEquals(4, pendingSplits.size());
+ Assertions.assertEquals(3, pendingSplits.get(0).size());
+ Assertions.assertEquals(3, pendingSplits.get(1).size());
+ Assertions.assertEquals(2, pendingSplits.get(2).size());
+ Assertions.assertEquals(2, pendingSplits.get(3).size());
+ }
+
+ @Test
+ public void shouldContinueRoundRobinAfterRestore() throws Exception {
+ TestingContext context = new TestingContext(3);
+ InfluxDBSourceState restoredState = new InfluxDBSourceState(true, new
HashMap<>(), 1);
+ InfluxDBSourceSplitEnumerator enumerator =
+ new InfluxDBSourceSplitEnumerator(context, restoredState,
null);
+
+ invokeAddPendingSplit(
+ enumerator, Collections.singletonList(new
InfluxDBSourceSplit("7", "")));
+
+ InfluxDBSourceState state = enumerator.snapshotState(1L);
+
+ Assertions.assertFalse(state.getPendingSplit().containsKey(0));
+ Assertions.assertEquals(1, state.getPendingSplit().get(1).size());
+ Assertions.assertFalse(state.getPendingSplit().containsKey(2));
+ }
+
+ @Test
+ public void shouldReassignReturnedSplitsToOriginalReader() throws
Exception {
+ TestingContext context = new TestingContext(3);
+ InfluxDBSourceSplitEnumerator enumerator =
+ new InfluxDBSourceSplitEnumerator(context, null, null);
+
+ enumerator.addSplitsBack(buildSplits(3), 2);
+
+ Assertions.assertEquals(0, context.getAssignedSplitCount(0));
+ Assertions.assertEquals(0, context.getAssignedSplitCount(1));
+ Assertions.assertEquals(3, context.getAssignedSplitCount(2));
+
+ InfluxDBSourceState state = enumerator.snapshotState(1L);
+ Assertions.assertTrue(state.getPendingSplit().isEmpty());
+ }
+
+ private void invokeAddPendingSplit(
+ InfluxDBSourceSplitEnumerator enumerator,
List<InfluxDBSourceSplit> splits)
+ throws Exception {
+ Method addPendingSplit =
+ InfluxDBSourceSplitEnumerator.class.getDeclaredMethod(
+ "addPendingSplit", java.util.Collection.class);
+ addPendingSplit.setAccessible(true);
+ addPendingSplit.invoke(enumerator, splits);
+ }
+
+ private List<InfluxDBSourceSplit> buildSplits(int size) {
+ List<InfluxDBSourceSplit> splits = new ArrayList<>();
+ for (int i = 0; i < size; i++) {
+ splits.add(new InfluxDBSourceSplit(String.valueOf(i), "query-" +
i));
+ }
+ return splits;
+ }
+
+ private static class TestingContext
+ implements SourceSplitEnumerator.Context<InfluxDBSourceSplit> {
+ private final int parallelism;
+ private final Set<Integer> readers;
+ private final Map<Integer, List<InfluxDBSourceSplit>> assignments =
new HashMap<>();
+
+ private TestingContext(int parallelism) {
+ this.parallelism = parallelism;
+ this.readers = new LinkedHashSet<>();
+ for (int i = 0; i < parallelism; i++) {
+ readers.add(i);
+ }
+ }
+
+ @Override
+ public int currentParallelism() {
+ return parallelism;
+ }
+
+ @Override
+ public Set<Integer> registeredReaders() {
+ return readers;
+ }
+
+ @Override
+ public void assignSplit(int subtaskId, List<InfluxDBSourceSplit>
splits) {
+ assignments.computeIfAbsent(subtaskId, ignored -> new
ArrayList<>()).addAll(splits);
+ }
+
+ @Override
+ public void signalNoMoreSplits(int subtask) {}
+
+ @Override
+ public void sendEventToSourceReader(int subtaskId, SourceEvent event)
{}
+
+ @Override
+ public MetricsContext getMetricsContext() {
+ return null;
+ }
+
+ @Override
+ public EventListener getEventListener() {
+ return null;
+ }
+
+ private int getAssignedSplitCount(int subtaskId) {
+ return assignments.getOrDefault(subtaskId,
Collections.emptyList()).size();
+ }
+ }
+}
diff --git
a/seatunnel-connectors-v2/connector-iotdb-v2/src/main/java/org/apache/seatunnel/connectors/seatunnel/iotdbv2/source/IoTDBv2SourceSplitEnumerator.java
b/seatunnel-connectors-v2/connector-iotdb-v2/src/main/java/org/apache/seatunnel/connectors/seatunnel/iotdbv2/source/IoTDBv2SourceSplitEnumerator.java
index 0da8de9bae..b3bc5f6ada 100644
---
a/seatunnel-connectors-v2/connector-iotdb-v2/src/main/java/org/apache/seatunnel/connectors/seatunnel/iotdbv2/source/IoTDBv2SourceSplitEnumerator.java
+++
b/seatunnel-connectors-v2/connector-iotdb-v2/src/main/java/org/apache/seatunnel/connectors/seatunnel/iotdbv2/source/IoTDBv2SourceSplitEnumerator.java
@@ -30,11 +30,14 @@ import lombok.extern.slf4j.Slf4j;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
+import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
import static
org.apache.seatunnel.connectors.seatunnel.iotdbv2.config.IoTDBv2SourceOptions.LOWER_BOUND;
import static
org.apache.seatunnel.connectors.seatunnel.iotdbv2.config.IoTDBv2SourceOptions.NUM_PARTITIONS;
@@ -59,6 +62,7 @@ public class IoTDBv2SourceSplitEnumerator
private final Context<IoTDBv2SourceSplit> context;
private final ReadonlyConfig conf;
private final Map<Integer, List<IoTDBv2SourceSplit>> pendingSplit;
+ private final AtomicInteger assignCount = new AtomicInteger(0);
private volatile boolean shouldEnumerate;
public IoTDBv2SourceSplitEnumerator(
@@ -77,6 +81,7 @@ public class IoTDBv2SourceSplitEnumerator
if (sourceState != null) {
this.shouldEnumerate = sourceState.isShouldEnumerate();
this.pendingSplit.putAll(sourceState.getPendingSplit());
+ this.assignCount.set(sourceState.getAssignCount());
}
}
@@ -183,8 +188,15 @@ public class IoTDBv2SourceSplitEnumerator
public void addSplitsBack(List<IoTDBv2SourceSplit> splits, int subtaskId) {
log.debug("Add back splits {} to IoTDBSourceSplitEnumerator.", splits);
if (!splits.isEmpty()) {
- addPendingSplit(splits);
- assignSplit(Collections.singletonList(subtaskId));
+ addPendingSplit(splits, subtaskId);
+ if (context.registeredReaders().contains(subtaskId)) {
+ assignSplit(Collections.singletonList(subtaskId));
+ } else {
+ log.warn(
+ "Reader {} is not registered. Pending splits {} are
not assigned.",
+ subtaskId,
+ splits);
+ }
}
}
@@ -204,14 +216,26 @@ public class IoTDBv2SourceSplitEnumerator
private void addPendingSplit(Collection<IoTDBv2SourceSplit> splits) {
synchronized (stateLock) {
int readerCount = context.currentParallelism();
- for (IoTDBv2SourceSplit split : splits) {
- int ownerReader = getSplitOwner(split.splitId(), readerCount);
+
+ List<IoTDBv2SourceSplit> sortedSplits =
+ splits.stream()
+
.sorted(Comparator.comparing(IoTDBv2SourceSplit::splitId))
+ .collect(Collectors.toList());
+
+ for (IoTDBv2SourceSplit split : sortedSplits) {
+ int ownerReader = getSplitOwner(assignCount.getAndIncrement(),
readerCount);
log.info("Assigning {} to {} reader.", split, ownerReader);
pendingSplit.computeIfAbsent(ownerReader, r -> new
ArrayList<>()).add(split);
}
}
}
+ private void addPendingSplit(Collection<IoTDBv2SourceSplit> splits, int
ownerReader) {
+ synchronized (stateLock) {
+ pendingSplit.computeIfAbsent(ownerReader, r -> new
ArrayList<>()).addAll(splits);
+ }
+ }
+
private void assignSplit(Collection<Integer> readers) {
log.debug("Assign pendingSplits to readers {}", readers);
@@ -238,12 +262,12 @@ public class IoTDBv2SourceSplitEnumerator
@Override
public IoTDBv2SourceState snapshotState(long checkpointId) throws
Exception {
synchronized (stateLock) {
- return new IoTDBv2SourceState(shouldEnumerate, pendingSplit);
+ return new IoTDBv2SourceState(shouldEnumerate, pendingSplit,
assignCount.get());
}
}
- private static int getSplitOwner(String tp, int numReaders) {
- return (tp.hashCode() & Integer.MAX_VALUE) % numReaders;
+ private static int getSplitOwner(int currentAssignCount, int numReaders) {
+ return currentAssignCount % numReaders;
}
@Override
diff --git
a/seatunnel-connectors-v2/connector-iotdb-v2/src/main/java/org/apache/seatunnel/connectors/seatunnel/iotdbv2/state/IoTDBv2SourceState.java
b/seatunnel-connectors-v2/connector-iotdb-v2/src/main/java/org/apache/seatunnel/connectors/seatunnel/iotdbv2/state/IoTDBv2SourceState.java
index 9ddb15725c..149e0556e4 100644
---
a/seatunnel-connectors-v2/connector-iotdb-v2/src/main/java/org/apache/seatunnel/connectors/seatunnel/iotdbv2/state/IoTDBv2SourceState.java
+++
b/seatunnel-connectors-v2/connector-iotdb-v2/src/main/java/org/apache/seatunnel/connectors/seatunnel/iotdbv2/state/IoTDBv2SourceState.java
@@ -33,4 +33,5 @@ public class IoTDBv2SourceState implements Serializable {
private static final long serialVersionUID = 7142773921678153583L;
private boolean shouldEnumerate;
private Map<Integer, List<IoTDBv2SourceSplit>> pendingSplit;
+ private int assignCount;
}
diff --git
a/seatunnel-connectors-v2/connector-iotdb-v2/src/test/java/org/apache/seatunnel/connectors/seatunnel/iotdbv2/source/IoTDBv2SourceSplitEnumeratorTest.java
b/seatunnel-connectors-v2/connector-iotdb-v2/src/test/java/org/apache/seatunnel/connectors/seatunnel/iotdbv2/source/IoTDBv2SourceSplitEnumeratorTest.java
new file mode 100644
index 0000000000..226bb893b8
--- /dev/null
+++
b/seatunnel-connectors-v2/connector-iotdb-v2/src/test/java/org/apache/seatunnel/connectors/seatunnel/iotdbv2/source/IoTDBv2SourceSplitEnumeratorTest.java
@@ -0,0 +1,158 @@
+/*
+ * 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.iotdbv2.source;
+
+import org.apache.seatunnel.api.common.metrics.MetricsContext;
+import org.apache.seatunnel.api.event.EventListener;
+import org.apache.seatunnel.api.source.SourceEvent;
+import org.apache.seatunnel.api.source.SourceSplitEnumerator;
+import
org.apache.seatunnel.connectors.seatunnel.iotdbv2.state.IoTDBv2SourceState;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.lang.reflect.Method;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+public class IoTDBv2SourceSplitEnumeratorTest {
+
+ @Test
+ public void shouldBalanceSplitsEvenlyAcrossReaders() throws Exception {
+ TestingContext context = new TestingContext(4);
+ IoTDBv2SourceSplitEnumerator enumerator =
+ new IoTDBv2SourceSplitEnumerator(context, null, null);
+
+ invokeAddPendingSplit(enumerator, buildSplits(10));
+
+ IoTDBv2SourceState state = enumerator.snapshotState(1L);
+ Map<Integer, List<IoTDBv2SourceSplit>> pendingSplits =
state.getPendingSplit();
+
+ Assertions.assertEquals(4, pendingSplits.size());
+ Assertions.assertEquals(3, pendingSplits.get(0).size());
+ Assertions.assertEquals(3, pendingSplits.get(1).size());
+ Assertions.assertEquals(2, pendingSplits.get(2).size());
+ Assertions.assertEquals(2, pendingSplits.get(3).size());
+ }
+
+ @Test
+ public void shouldContinueRoundRobinAfterRestore() throws Exception {
+ TestingContext context = new TestingContext(3);
+ IoTDBv2SourceState restoredState = new IoTDBv2SourceState(true, new
HashMap<>(), 1);
+ IoTDBv2SourceSplitEnumerator enumerator =
+ new IoTDBv2SourceSplitEnumerator(context, null, restoredState);
+
+ invokeAddPendingSplit(
+ enumerator, Collections.singletonList(new
IoTDBv2SourceSplit("7", "")));
+
+ IoTDBv2SourceState state = enumerator.snapshotState(1L);
+
+ Assertions.assertFalse(state.getPendingSplit().containsKey(0));
+ Assertions.assertEquals(1, state.getPendingSplit().get(1).size());
+ Assertions.assertFalse(state.getPendingSplit().containsKey(2));
+ }
+
+ @Test
+ public void shouldReassignReturnedSplitsToOriginalReader() throws
Exception {
+ TestingContext context = new TestingContext(3);
+ IoTDBv2SourceSplitEnumerator enumerator =
+ new IoTDBv2SourceSplitEnumerator(context, null, null);
+
+ enumerator.addSplitsBack(buildSplits(3), 2);
+
+ Assertions.assertEquals(0, context.getAssignedSplitCount(0));
+ Assertions.assertEquals(0, context.getAssignedSplitCount(1));
+ Assertions.assertEquals(3, context.getAssignedSplitCount(2));
+
+ IoTDBv2SourceState state = enumerator.snapshotState(1L);
+ Assertions.assertTrue(state.getPendingSplit().isEmpty());
+ }
+
+ private void invokeAddPendingSplit(
+ IoTDBv2SourceSplitEnumerator enumerator, List<IoTDBv2SourceSplit>
splits)
+ throws Exception {
+ Method addPendingSplit =
+ IoTDBv2SourceSplitEnumerator.class.getDeclaredMethod(
+ "addPendingSplit", java.util.Collection.class);
+ addPendingSplit.setAccessible(true);
+ addPendingSplit.invoke(enumerator, splits);
+ }
+
+ private List<IoTDBv2SourceSplit> buildSplits(int size) {
+ List<IoTDBv2SourceSplit> splits = new ArrayList<>();
+ for (int i = 0; i < size; i++) {
+ splits.add(new IoTDBv2SourceSplit(String.valueOf(i), "query-" +
i));
+ }
+ return splits;
+ }
+
+ private static class TestingContext
+ implements SourceSplitEnumerator.Context<IoTDBv2SourceSplit> {
+ private final int parallelism;
+ private final Set<Integer> readers;
+ private final Map<Integer, List<IoTDBv2SourceSplit>> assignments = new
HashMap<>();
+
+ private TestingContext(int parallelism) {
+ this.parallelism = parallelism;
+ this.readers = new LinkedHashSet<>();
+ for (int i = 0; i < parallelism; i++) {
+ readers.add(i);
+ }
+ }
+
+ @Override
+ public int currentParallelism() {
+ return parallelism;
+ }
+
+ @Override
+ public Set<Integer> registeredReaders() {
+ return readers;
+ }
+
+ @Override
+ public void assignSplit(int subtaskId, List<IoTDBv2SourceSplit>
splits) {
+ assignments.computeIfAbsent(subtaskId, ignored -> new
ArrayList<>()).addAll(splits);
+ }
+
+ @Override
+ public void signalNoMoreSplits(int subtask) {}
+
+ @Override
+ public void sendEventToSourceReader(int subtaskId, SourceEvent event)
{}
+
+ @Override
+ public MetricsContext getMetricsContext() {
+ return null;
+ }
+
+ @Override
+ public EventListener getEventListener() {
+ return null;
+ }
+
+ private int getAssignedSplitCount(int subtaskId) {
+ return assignments.getOrDefault(subtaskId,
Collections.emptyList()).size();
+ }
+ }
+}
diff --git
a/seatunnel-connectors-v2/connector-iotdb/src/main/java/org/apache/seatunnel/connectors/seatunnel/iotdb/source/IoTDBSourceSplitEnumerator.java
b/seatunnel-connectors-v2/connector-iotdb/src/main/java/org/apache/seatunnel/connectors/seatunnel/iotdb/source/IoTDBSourceSplitEnumerator.java
index a004819591..d9bdd0e221 100644
---
a/seatunnel-connectors-v2/connector-iotdb/src/main/java/org/apache/seatunnel/connectors/seatunnel/iotdb/source/IoTDBSourceSplitEnumerator.java
+++
b/seatunnel-connectors-v2/connector-iotdb/src/main/java/org/apache/seatunnel/connectors/seatunnel/iotdb/source/IoTDBSourceSplitEnumerator.java
@@ -30,11 +30,14 @@ import lombok.extern.slf4j.Slf4j;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
+import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
import static
org.apache.iotdb.tsfile.common.constant.QueryConstant.RESERVED_TIME;
import static
org.apache.seatunnel.connectors.seatunnel.iotdb.config.IoTDBSourceOptions.LOWER_BOUND;
@@ -59,6 +62,7 @@ public class IoTDBSourceSplitEnumerator
private final Context<IoTDBSourceSplit> context;
private final ReadonlyConfig conf;
private final Map<Integer, List<IoTDBSourceSplit>> pendingSplit;
+ private final AtomicInteger assignCount = new AtomicInteger(0);
private volatile boolean shouldEnumerate;
public IoTDBSourceSplitEnumerator(
@@ -77,6 +81,7 @@ public class IoTDBSourceSplitEnumerator
if (sourceState != null) {
this.shouldEnumerate = sourceState.isShouldEnumerate();
this.pendingSplit.putAll(sourceState.getPendingSplit());
+ this.assignCount.set(sourceState.getAssignCount());
}
}
@@ -183,8 +188,15 @@ public class IoTDBSourceSplitEnumerator
public void addSplitsBack(List<IoTDBSourceSplit> splits, int subtaskId) {
log.debug("Add back splits {} to IoTDBSourceSplitEnumerator.", splits);
if (!splits.isEmpty()) {
- addPendingSplit(splits);
- assignSplit(Collections.singletonList(subtaskId));
+ addPendingSplit(splits, subtaskId);
+ if (context.registeredReaders().contains(subtaskId)) {
+ assignSplit(Collections.singletonList(subtaskId));
+ } else {
+ log.warn(
+ "Reader {} is not registered. Pending splits {} are
not assigned.",
+ subtaskId,
+ splits);
+ }
}
}
@@ -203,13 +215,23 @@ public class IoTDBSourceSplitEnumerator
private void addPendingSplit(Collection<IoTDBSourceSplit> splits) {
int readerCount = context.currentParallelism();
- for (IoTDBSourceSplit split : splits) {
- int ownerReader = getSplitOwner(split.splitId(), readerCount);
+
+ List<IoTDBSourceSplit> sortedSplits =
+ splits.stream()
+
.sorted(Comparator.comparing(IoTDBSourceSplit::splitId))
+ .collect(Collectors.toList());
+
+ for (IoTDBSourceSplit split : sortedSplits) {
+ int ownerReader = getSplitOwner(assignCount.getAndIncrement(),
readerCount);
log.info("Assigning {} to {} reader.", split, ownerReader);
pendingSplit.computeIfAbsent(ownerReader, r -> new
ArrayList<>()).add(split);
}
}
+ private void addPendingSplit(Collection<IoTDBSourceSplit> splits, int
ownerReader) {
+ pendingSplit.computeIfAbsent(ownerReader, r -> new
ArrayList<>()).addAll(splits);
+ }
+
private void assignSplit(Collection<Integer> readers) {
log.debug("Assign pendingSplits to readers {}", readers);
@@ -234,12 +256,12 @@ public class IoTDBSourceSplitEnumerator
@Override
public IoTDBSourceState snapshotState(long checkpointId) throws Exception {
synchronized (stateLock) {
- return new IoTDBSourceState(shouldEnumerate, pendingSplit);
+ return new IoTDBSourceState(shouldEnumerate, pendingSplit,
assignCount.get());
}
}
- private static int getSplitOwner(String tp, int numReaders) {
- return (tp.hashCode() & Integer.MAX_VALUE) % numReaders;
+ private static int getSplitOwner(int currentAssignCount, int numReaders) {
+ return currentAssignCount % numReaders;
}
@Override
diff --git
a/seatunnel-connectors-v2/connector-iotdb/src/main/java/org/apache/seatunnel/connectors/seatunnel/iotdb/state/IoTDBSourceState.java
b/seatunnel-connectors-v2/connector-iotdb/src/main/java/org/apache/seatunnel/connectors/seatunnel/iotdb/state/IoTDBSourceState.java
index 1d8fdc39bd..b2d5eba147 100644
---
a/seatunnel-connectors-v2/connector-iotdb/src/main/java/org/apache/seatunnel/connectors/seatunnel/iotdb/state/IoTDBSourceState.java
+++
b/seatunnel-connectors-v2/connector-iotdb/src/main/java/org/apache/seatunnel/connectors/seatunnel/iotdb/state/IoTDBSourceState.java
@@ -33,4 +33,5 @@ public class IoTDBSourceState implements Serializable {
private static final long serialVersionUID = 7142773921678153583L;
private boolean shouldEnumerate;
private Map<Integer, List<IoTDBSourceSplit>> pendingSplit;
+ private int assignCount;
}
diff --git
a/seatunnel-connectors-v2/connector-iotdb/src/test/java/org/apache/seatunnel/connectors/seatunnel/iotdb/source/IoTDBSourceSplitEnumeratorTest.java
b/seatunnel-connectors-v2/connector-iotdb/src/test/java/org/apache/seatunnel/connectors/seatunnel/iotdb/source/IoTDBSourceSplitEnumeratorTest.java
new file mode 100644
index 0000000000..9647d9235e
--- /dev/null
+++
b/seatunnel-connectors-v2/connector-iotdb/src/test/java/org/apache/seatunnel/connectors/seatunnel/iotdb/source/IoTDBSourceSplitEnumeratorTest.java
@@ -0,0 +1,153 @@
+/*
+ * 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.iotdb.source;
+
+import org.apache.seatunnel.api.common.metrics.MetricsContext;
+import org.apache.seatunnel.api.event.EventListener;
+import org.apache.seatunnel.api.source.SourceEvent;
+import org.apache.seatunnel.api.source.SourceSplitEnumerator;
+import org.apache.seatunnel.connectors.seatunnel.iotdb.state.IoTDBSourceState;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.lang.reflect.Method;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+public class IoTDBSourceSplitEnumeratorTest {
+
+ @Test
+ public void shouldBalanceSplitsEvenlyAcrossReaders() throws Exception {
+ TestingContext context = new TestingContext(4);
+ IoTDBSourceSplitEnumerator enumerator = new
IoTDBSourceSplitEnumerator(context, null, null);
+
+ invokeAddPendingSplit(enumerator, buildSplits(10));
+
+ IoTDBSourceState state = enumerator.snapshotState(1L);
+ Map<Integer, List<IoTDBSourceSplit>> pendingSplits =
state.getPendingSplit();
+
+ Assertions.assertEquals(4, pendingSplits.size());
+ Assertions.assertEquals(3, pendingSplits.get(0).size());
+ Assertions.assertEquals(3, pendingSplits.get(1).size());
+ Assertions.assertEquals(2, pendingSplits.get(2).size());
+ Assertions.assertEquals(2, pendingSplits.get(3).size());
+ }
+
+ @Test
+ public void shouldContinueRoundRobinAfterRestore() throws Exception {
+ TestingContext context = new TestingContext(3);
+ IoTDBSourceState restoredState = new IoTDBSourceState(true, new
HashMap<>(), 1);
+ IoTDBSourceSplitEnumerator enumerator =
+ new IoTDBSourceSplitEnumerator(context, null, restoredState);
+
+ invokeAddPendingSplit(enumerator, Collections.singletonList(new
IoTDBSourceSplit("7", "")));
+
+ IoTDBSourceState state = enumerator.snapshotState(1L);
+
+ Assertions.assertFalse(state.getPendingSplit().containsKey(0));
+ Assertions.assertEquals(1, state.getPendingSplit().get(1).size());
+ Assertions.assertFalse(state.getPendingSplit().containsKey(2));
+ }
+
+ @Test
+ public void shouldReassignReturnedSplitsToOriginalReader() throws
Exception {
+ TestingContext context = new TestingContext(3);
+ IoTDBSourceSplitEnumerator enumerator = new
IoTDBSourceSplitEnumerator(context, null, null);
+
+ enumerator.addSplitsBack(buildSplits(3), 2);
+
+ Assertions.assertEquals(0, context.getAssignedSplitCount(0));
+ Assertions.assertEquals(0, context.getAssignedSplitCount(1));
+ Assertions.assertEquals(3, context.getAssignedSplitCount(2));
+
+ IoTDBSourceState state = enumerator.snapshotState(1L);
+ Assertions.assertTrue(state.getPendingSplit().isEmpty());
+ }
+
+ private void invokeAddPendingSplit(
+ IoTDBSourceSplitEnumerator enumerator, List<IoTDBSourceSplit>
splits) throws Exception {
+ Method addPendingSplit =
+ IoTDBSourceSplitEnumerator.class.getDeclaredMethod(
+ "addPendingSplit", java.util.Collection.class);
+ addPendingSplit.setAccessible(true);
+ addPendingSplit.invoke(enumerator, splits);
+ }
+
+ private List<IoTDBSourceSplit> buildSplits(int size) {
+ List<IoTDBSourceSplit> splits = new ArrayList<>();
+ for (int i = 0; i < size; i++) {
+ splits.add(new IoTDBSourceSplit(String.valueOf(i), "query-" + i));
+ }
+ return splits;
+ }
+
+ private static class TestingContext implements
SourceSplitEnumerator.Context<IoTDBSourceSplit> {
+ private final int parallelism;
+ private final Set<Integer> readers;
+ private final Map<Integer, List<IoTDBSourceSplit>> assignments = new
HashMap<>();
+
+ private TestingContext(int parallelism) {
+ this.parallelism = parallelism;
+ this.readers = new LinkedHashSet<>();
+ for (int i = 0; i < parallelism; i++) {
+ readers.add(i);
+ }
+ }
+
+ @Override
+ public int currentParallelism() {
+ return parallelism;
+ }
+
+ @Override
+ public Set<Integer> registeredReaders() {
+ return readers;
+ }
+
+ @Override
+ public void assignSplit(int subtaskId, List<IoTDBSourceSplit> splits) {
+ assignments.computeIfAbsent(subtaskId, ignored -> new
ArrayList<>()).addAll(splits);
+ }
+
+ @Override
+ public void signalNoMoreSplits(int subtask) {}
+
+ @Override
+ public void sendEventToSourceReader(int subtaskId, SourceEvent event)
{}
+
+ @Override
+ public MetricsContext getMetricsContext() {
+ return null;
+ }
+
+ @Override
+ public EventListener getEventListener() {
+ return null;
+ }
+
+ private int getAssignedSplitCount(int subtaskId) {
+ return assignments.getOrDefault(subtaskId,
Collections.emptyList()).size();
+ }
+ }
+}
diff --git
a/seatunnel-connectors-v2/connector-kudu/src/main/java/org/apache/seatunnel/connectors/seatunnel/kudu/source/KuduSourceSplitEnumerator.java
b/seatunnel-connectors-v2/connector-kudu/src/main/java/org/apache/seatunnel/connectors/seatunnel/kudu/source/KuduSourceSplitEnumerator.java
index 469332da1c..6121e005a4 100644
---
a/seatunnel-connectors-v2/connector-kudu/src/main/java/org/apache/seatunnel/connectors/seatunnel/kudu/source/KuduSourceSplitEnumerator.java
+++
b/seatunnel-connectors-v2/connector-kudu/src/main/java/org/apache/seatunnel/connectors/seatunnel/kudu/source/KuduSourceSplitEnumerator.java
@@ -33,11 +33,13 @@ import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
+import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;
import java.util.stream.Collectors;
@@ -53,6 +55,7 @@ public class KuduSourceSplitEnumerator
private final Map<Integer, List<KuduSourceSplit>> pendingSplits;
private final Map<TablePath, KuduSourceTableConfig> tables;
private final KuduInputFormat kuduInputFormat;
+ private final AtomicInteger assignCount = new AtomicInteger(0);
private final Object stateLock = new Object();
@@ -79,6 +82,7 @@ public class KuduSourceSplitEnumerator
} else {
this.pendingTables = new
ConcurrentLinkedQueue<>(checkpointState.getPendingTables());
this.pendingSplits = new
HashMap<>(checkpointState.getPendingSplits());
+ this.assignCount.set(checkpointState.getAssignCount());
}
}
@@ -160,8 +164,14 @@ public class KuduSourceSplitEnumerator
private void addPendingSplit(Collection<KuduSourceSplit> splits) {
int readerCount = enumeratorContext.currentParallelism();
- for (KuduSourceSplit split : splits) {
- int ownerReader = getSplitOwner(split.splitId(), readerCount);
+
+ List<KuduSourceSplit> sortedSplits =
+ splits.stream()
+ .sorted(Comparator.comparing(KuduSourceSplit::splitId))
+ .collect(Collectors.toList());
+
+ for (KuduSourceSplit split : sortedSplits) {
+ int ownerReader = getSplitOwner(assignCount.getAndIncrement(),
readerCount);
log.info("Assigning {} to {} reader.", split, ownerReader);
pendingSplits.computeIfAbsent(ownerReader, r -> new
ArrayList<>()).add(split);
}
@@ -171,8 +181,8 @@ public class KuduSourceSplitEnumerator
pendingSplits.computeIfAbsent(ownerReader, r -> new
ArrayList<>()).addAll(splits);
}
- private int getSplitOwner(String splitId, int numReaders) {
- return (splitId.hashCode() & Integer.MAX_VALUE) % numReaders;
+ private int getSplitOwner(int currentAssignCount, int numReaders) {
+ return currentAssignCount % numReaders;
}
@Override
@@ -198,7 +208,10 @@ public class KuduSourceSplitEnumerator
@Override
public KuduSourceState snapshotState(long checkpointId) throws Exception {
synchronized (stateLock) {
- return new KuduSourceState(new ArrayList(pendingTables), new
HashMap<>(pendingSplits));
+ return new KuduSourceState(
+ new ArrayList<>(pendingTables),
+ new HashMap<>(pendingSplits),
+ assignCount.get());
}
}
diff --git
a/seatunnel-connectors-v2/connector-kudu/src/main/java/org/apache/seatunnel/connectors/seatunnel/kudu/state/KuduSourceState.java
b/seatunnel-connectors-v2/connector-kudu/src/main/java/org/apache/seatunnel/connectors/seatunnel/kudu/state/KuduSourceState.java
index 95cd02dd90..0b76218aef 100644
---
a/seatunnel-connectors-v2/connector-kudu/src/main/java/org/apache/seatunnel/connectors/seatunnel/kudu/state/KuduSourceState.java
+++
b/seatunnel-connectors-v2/connector-kudu/src/main/java/org/apache/seatunnel/connectors/seatunnel/kudu/state/KuduSourceState.java
@@ -33,4 +33,5 @@ public class KuduSourceState implements Serializable {
private static final long serialVersionUID = -3141157457869831037L;
private List<TablePath> pendingTables;
private Map<Integer, List<KuduSourceSplit>> pendingSplits;
+ private int assignCount;
}
diff --git
a/seatunnel-connectors-v2/connector-kudu/src/test/java/org/apache/seatunnel/connectors/seatunnel/kudu/source/KuduSourceSplitEnumeratorTest.java
b/seatunnel-connectors-v2/connector-kudu/src/test/java/org/apache/seatunnel/connectors/seatunnel/kudu/source/KuduSourceSplitEnumeratorTest.java
new file mode 100644
index 0000000000..d992f037a3
--- /dev/null
+++
b/seatunnel-connectors-v2/connector-kudu/src/test/java/org/apache/seatunnel/connectors/seatunnel/kudu/source/KuduSourceSplitEnumeratorTest.java
@@ -0,0 +1,185 @@
+/*
+ * 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.kudu.source;
+
+import org.apache.seatunnel.api.common.metrics.MetricsContext;
+import org.apache.seatunnel.api.event.EventListener;
+import org.apache.seatunnel.api.source.SourceEvent;
+import org.apache.seatunnel.api.source.SourceSplitEnumerator;
+import org.apache.seatunnel.api.table.catalog.TablePath;
+import org.apache.seatunnel.connectors.seatunnel.kudu.config.KuduSourceConfig;
+import org.apache.seatunnel.connectors.seatunnel.kudu.state.KuduSourceState;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import sun.misc.Unsafe;
+
+import java.lang.reflect.Field;
+import java.lang.reflect.Method;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+public class KuduSourceSplitEnumeratorTest {
+
+ @Test
+ public void shouldBalanceSplitsEvenlyAcrossReaders() throws Exception {
+ TestingContext context = new TestingContext(4);
+ KuduSourceSplitEnumerator enumerator =
+ new KuduSourceSplitEnumerator(context, mockSourceConfig(),
null);
+
+ invokeAddPendingSplit(enumerator, buildSplits(10));
+
+ KuduSourceState state = enumerator.snapshotState(1L);
+ Map<Integer, List<KuduSourceSplit>> pendingSplits =
state.getPendingSplits();
+
+ Assertions.assertEquals(4, pendingSplits.size());
+ Assertions.assertEquals(3, pendingSplits.get(0).size());
+ Assertions.assertEquals(3, pendingSplits.get(1).size());
+ Assertions.assertEquals(2, pendingSplits.get(2).size());
+ Assertions.assertEquals(2, pendingSplits.get(3).size());
+ }
+
+ @Test
+ public void shouldContinueRoundRobinAfterRestore() throws Exception {
+ TestingContext context = new TestingContext(3);
+ KuduSourceState restoredState =
+ new KuduSourceState(Collections.emptyList(), new HashMap<>(),
1);
+ KuduSourceSplitEnumerator enumerator =
+ new KuduSourceSplitEnumerator(context, mockSourceConfig(),
restoredState);
+
+ invokeAddPendingSplit(enumerator,
Collections.singletonList(buildSplit(7)));
+
+ KuduSourceState state = enumerator.snapshotState(1L);
+
+ Assertions.assertFalse(state.getPendingSplits().containsKey(0));
+ Assertions.assertEquals(1, state.getPendingSplits().get(1).size());
+ Assertions.assertFalse(state.getPendingSplits().containsKey(2));
+ }
+
+ @Test
+ public void shouldReassignReturnedSplitsToOriginalReader() {
+ TestingContext context = new TestingContext(3);
+ KuduSourceSplitEnumerator enumerator =
+ new KuduSourceSplitEnumerator(context, mockSourceConfig(),
null);
+
+ enumerator.addSplitsBack(buildSplits(3), 2);
+
+ Assertions.assertEquals(0, context.getAssignedSplitCount(0));
+ Assertions.assertEquals(0, context.getAssignedSplitCount(1));
+ Assertions.assertEquals(3, context.getAssignedSplitCount(2));
+
+ Assertions.assertTrue(context.getAssignments().containsKey(2));
+ }
+
+ private void invokeAddPendingSplit(
+ KuduSourceSplitEnumerator enumerator, List<KuduSourceSplit>
splits) throws Exception {
+ Method addPendingSplit =
+ KuduSourceSplitEnumerator.class.getDeclaredMethod(
+ "addPendingSplit", java.util.Collection.class);
+ addPendingSplit.setAccessible(true);
+ addPendingSplit.invoke(enumerator, splits);
+ }
+
+ private List<KuduSourceSplit> buildSplits(int size) {
+ List<KuduSourceSplit> splits = new ArrayList<>();
+ for (int i = 0; i < size; i++) {
+ splits.add(buildSplit(i));
+ }
+ return splits;
+ }
+
+ private KuduSourceSplit buildSplit(int splitId) {
+ return new KuduSourceSplit(TablePath.of("db", null, "table"), splitId,
new byte[0]);
+ }
+
+ private KuduSourceConfig mockSourceConfig() {
+ try {
+ Field unsafeField = Unsafe.class.getDeclaredField("theUnsafe");
+ unsafeField.setAccessible(true);
+ Unsafe unsafe = (Unsafe) unsafeField.get(null);
+ KuduSourceConfig sourceConfig =
+ (KuduSourceConfig)
unsafe.allocateInstance(KuduSourceConfig.class);
+
+ Field tableConfigListField =
KuduSourceConfig.class.getDeclaredField("tableConfigList");
+ tableConfigListField.setAccessible(true);
+ tableConfigListField.set(sourceConfig, Collections.emptyList());
+ return sourceConfig;
+ } catch (ReflectiveOperationException e) {
+ throw new IllegalStateException("Failed to create test
KuduSourceConfig", e);
+ }
+ }
+
+ private static class TestingContext implements
SourceSplitEnumerator.Context<KuduSourceSplit> {
+ private final int parallelism;
+ private final Set<Integer> readers;
+ private final Map<Integer, List<KuduSourceSplit>> assignments = new
HashMap<>();
+
+ private TestingContext(int parallelism) {
+ this.parallelism = parallelism;
+ this.readers = new LinkedHashSet<>();
+ for (int i = 0; i < parallelism; i++) {
+ readers.add(i);
+ }
+ }
+
+ @Override
+ public int currentParallelism() {
+ return parallelism;
+ }
+
+ @Override
+ public Set<Integer> registeredReaders() {
+ return readers;
+ }
+
+ @Override
+ public void assignSplit(int subtaskId, List<KuduSourceSplit> splits) {
+ assignments.computeIfAbsent(subtaskId, ignored -> new
ArrayList<>()).addAll(splits);
+ }
+
+ @Override
+ public void signalNoMoreSplits(int subtask) {}
+
+ @Override
+ public void sendEventToSourceReader(int subtaskId, SourceEvent event)
{}
+
+ @Override
+ public MetricsContext getMetricsContext() {
+ return null;
+ }
+
+ @Override
+ public EventListener getEventListener() {
+ return null;
+ }
+
+ private int getAssignedSplitCount(int subtaskId) {
+ return assignments.getOrDefault(subtaskId,
Collections.emptyList()).size();
+ }
+
+ private Map<Integer, List<KuduSourceSplit>> getAssignments() {
+ return assignments;
+ }
+ }
+}
diff --git
a/seatunnel-connectors-v2/connector-tdengine/src/main/java/org/apache/seatunnel/connectors/seatunnel/tdengine/source/TDengineSourceSplitEnumerator.java
b/seatunnel-connectors-v2/connector-tdengine/src/main/java/org/apache/seatunnel/connectors/seatunnel/tdengine/source/TDengineSourceSplitEnumerator.java
index 911a9a6ec1..a95f6fbf85 100644
---
a/seatunnel-connectors-v2/connector-tdengine/src/main/java/org/apache/seatunnel/connectors/seatunnel/tdengine/source/TDengineSourceSplitEnumerator.java
+++
b/seatunnel-connectors-v2/connector-tdengine/src/main/java/org/apache/seatunnel/connectors/seatunnel/tdengine/source/TDengineSourceSplitEnumerator.java
@@ -29,10 +29,12 @@ import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
+import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
@Slf4j
@@ -45,6 +47,7 @@ public class TDengineSourceSplitEnumerator
private volatile boolean shouldEnumerate;
private final Object stateLock = new Object();
private final Map<Integer, List<TDengineSourceSplit>> pendingSplits = new
ConcurrentHashMap<>();
+ private final AtomicInteger assignCount = new AtomicInteger(0);
public TDengineSourceSplitEnumerator(
StableMetadata stableMetadata,
@@ -65,11 +68,12 @@ public class TDengineSourceSplitEnumerator
if (sourceState != null) {
this.shouldEnumerate = sourceState.isShouldEnumerate();
this.pendingSplits.putAll(sourceState.getPendingSplits());
+ this.assignCount.set(sourceState.getAssignCount());
}
}
- private static int getSplitOwner(String tp, int numReaders) {
- return (tp.hashCode() & Integer.MAX_VALUE) % numReaders;
+ private static int getSplitOwner(int currentAssignCount, int numReaders) {
+ return currentAssignCount % numReaders;
}
@Override
@@ -95,8 +99,14 @@ public class TDengineSourceSplitEnumerator
private void addPendingSplit(List<TDengineSourceSplit> newSplits) {
int readerCount = context.currentParallelism();
- for (TDengineSourceSplit split : newSplits) {
- int ownerReader = getSplitOwner(split.splitId(), readerCount);
+
+ List<TDengineSourceSplit> sortedSplits =
+ newSplits.stream()
+
.sorted(Comparator.comparing(TDengineSourceSplit::splitId))
+ .collect(Collectors.toList());
+
+ for (TDengineSourceSplit split : sortedSplits) {
+ int ownerReader = getSplitOwner(assignCount.getAndIncrement(),
readerCount);
pendingSplits.computeIfAbsent(ownerReader, r -> new
ArrayList<>()).add(split);
}
}
@@ -145,8 +155,15 @@ public class TDengineSourceSplitEnumerator
public void addSplitsBack(List<TDengineSourceSplit> splits, int subtaskId)
{
log.info("Add back splits {} to TDengineSourceSplitEnumerator.",
splits);
if (!splits.isEmpty()) {
- addPendingSplit(splits);
- assignSplit(Collections.singletonList(subtaskId));
+ addPendingSplit(splits, subtaskId);
+ if (context.registeredReaders().contains(subtaskId)) {
+ assignSplit(Collections.singletonList(subtaskId));
+ } else {
+ log.warn(
+ "Reader {} is not registered. Pending splits {} are
not assigned.",
+ subtaskId,
+ splits);
+ }
}
}
@@ -184,10 +201,14 @@ public class TDengineSourceSplitEnumerator
}
}
+ private void addPendingSplit(Collection<TDengineSourceSplit> splits, int
ownerReader) {
+ pendingSplits.computeIfAbsent(ownerReader, r -> new
ArrayList<>()).addAll(splits);
+ }
+
@Override
public TDengineSourceState snapshotState(long checkpointId) {
synchronized (stateLock) {
- return new TDengineSourceState(shouldEnumerate, pendingSplits);
+ return new TDengineSourceState(shouldEnumerate, pendingSplits,
assignCount.get());
}
}
diff --git
a/seatunnel-connectors-v2/connector-tdengine/src/main/java/org/apache/seatunnel/connectors/seatunnel/tdengine/state/TDengineSourceState.java
b/seatunnel-connectors-v2/connector-tdengine/src/main/java/org/apache/seatunnel/connectors/seatunnel/tdengine/state/TDengineSourceState.java
index 0bd300a8d3..27e1f72b6a 100644
---
a/seatunnel-connectors-v2/connector-tdengine/src/main/java/org/apache/seatunnel/connectors/seatunnel/tdengine/state/TDengineSourceState.java
+++
b/seatunnel-connectors-v2/connector-tdengine/src/main/java/org/apache/seatunnel/connectors/seatunnel/tdengine/state/TDengineSourceState.java
@@ -32,4 +32,5 @@ public class TDengineSourceState implements Serializable {
private static final long serialVersionUID = 6915087497958523069L;
private boolean shouldEnumerate;
private final Map<Integer, List<TDengineSourceSplit>> pendingSplits;
+ private int assignCount;
}
diff --git
a/seatunnel-connectors-v2/connector-tdengine/src/test/java/org/apache/seatunnel/connectors/seatunnel/tdengine/source/TDengineSourceSplitEnumeratorTest.java
b/seatunnel-connectors-v2/connector-tdengine/src/test/java/org/apache/seatunnel/connectors/seatunnel/tdengine/source/TDengineSourceSplitEnumeratorTest.java
new file mode 100644
index 0000000000..70893c3900
--- /dev/null
+++
b/seatunnel-connectors-v2/connector-tdengine/src/test/java/org/apache/seatunnel/connectors/seatunnel/tdengine/source/TDengineSourceSplitEnumeratorTest.java
@@ -0,0 +1,164 @@
+/*
+ * 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.tdengine.source;
+
+import org.apache.seatunnel.api.common.metrics.MetricsContext;
+import org.apache.seatunnel.api.event.EventListener;
+import org.apache.seatunnel.api.source.SourceEvent;
+import org.apache.seatunnel.api.source.SourceSplitEnumerator;
+import
org.apache.seatunnel.connectors.seatunnel.tdengine.state.TDengineSourceState;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.lang.reflect.Method;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+public class TDengineSourceSplitEnumeratorTest {
+
+ @Test
+ public void shouldBalanceSplitsEvenlyAcrossReaders() throws Exception {
+ TestingContext context = new TestingContext(4);
+ TDengineSourceSplitEnumerator enumerator = newEnumerator(context,
null);
+
+ invokeAddPendingSplit(enumerator, buildSplits(10));
+
+ TDengineSourceState state = enumerator.snapshotState(1L);
+ Map<Integer, List<TDengineSourceSplit>> pendingSplits =
state.getPendingSplits();
+
+ Assertions.assertEquals(4, pendingSplits.size());
+ Assertions.assertEquals(3, pendingSplits.get(0).size());
+ Assertions.assertEquals(3, pendingSplits.get(1).size());
+ Assertions.assertEquals(2, pendingSplits.get(2).size());
+ Assertions.assertEquals(2, pendingSplits.get(3).size());
+ }
+
+ @Test
+ public void shouldContinueRoundRobinAfterRestore() throws Exception {
+ TestingContext context = new TestingContext(3);
+ TDengineSourceState restoredState = new TDengineSourceState(true, new
HashMap<>(), 1);
+ TDengineSourceSplitEnumerator enumerator = newEnumerator(context,
restoredState);
+
+ invokeAddPendingSplit(
+ enumerator, Collections.singletonList(new
TDengineSourceSplit("7", "")));
+
+ TDengineSourceState state = enumerator.snapshotState(1L);
+
+ Assertions.assertFalse(state.getPendingSplits().containsKey(0));
+ Assertions.assertEquals(1, state.getPendingSplits().get(1).size());
+ Assertions.assertFalse(state.getPendingSplits().containsKey(2));
+ }
+
+ @Test
+ public void shouldReassignReturnedSplitsToOriginalReader() throws
Exception {
+ TestingContext context = new TestingContext(3);
+ TDengineSourceSplitEnumerator enumerator = newEnumerator(context,
null);
+
+ enumerator.addSplitsBack(buildSplits(3), 2);
+
+ Assertions.assertEquals(0, context.getAssignedSplitCount(0));
+ Assertions.assertEquals(0, context.getAssignedSplitCount(1));
+ Assertions.assertEquals(3, context.getAssignedSplitCount(2));
+
+ TDengineSourceState state = enumerator.snapshotState(1L);
+ Assertions.assertTrue(state.getPendingSplits().isEmpty());
+ }
+
+ private TDengineSourceSplitEnumerator newEnumerator(
+ TestingContext context, TDengineSourceState sourceState) {
+ return new TDengineSourceSplitEnumerator(
+ new StableMetadata(null, null, Collections.emptyList()),
+ null,
+ sourceState,
+ context);
+ }
+
+ private void invokeAddPendingSplit(
+ TDengineSourceSplitEnumerator enumerator,
List<TDengineSourceSplit> splits)
+ throws Exception {
+ Method addPendingSplit =
+ TDengineSourceSplitEnumerator.class.getDeclaredMethod(
+ "addPendingSplit", List.class);
+ addPendingSplit.setAccessible(true);
+ addPendingSplit.invoke(enumerator, splits);
+ }
+
+ private List<TDengineSourceSplit> buildSplits(int size) {
+ List<TDengineSourceSplit> splits = new ArrayList<>();
+ for (int i = 0; i < size; i++) {
+ splits.add(new TDengineSourceSplit(String.valueOf(i), "query-" +
i));
+ }
+ return splits;
+ }
+
+ private static class TestingContext
+ implements SourceSplitEnumerator.Context<TDengineSourceSplit> {
+ private final int parallelism;
+ private final Set<Integer> readers;
+ private final Map<Integer, List<TDengineSourceSplit>> assignments =
new HashMap<>();
+
+ private TestingContext(int parallelism) {
+ this.parallelism = parallelism;
+ this.readers = new LinkedHashSet<>();
+ for (int i = 0; i < parallelism; i++) {
+ readers.add(i);
+ }
+ }
+
+ @Override
+ public int currentParallelism() {
+ return parallelism;
+ }
+
+ @Override
+ public Set<Integer> registeredReaders() {
+ return readers;
+ }
+
+ @Override
+ public void assignSplit(int subtaskId, List<TDengineSourceSplit>
splits) {
+ assignments.computeIfAbsent(subtaskId, ignored -> new
ArrayList<>()).addAll(splits);
+ }
+
+ @Override
+ public void signalNoMoreSplits(int subtask) {}
+
+ @Override
+ public void sendEventToSourceReader(int subtaskId, SourceEvent event)
{}
+
+ @Override
+ public MetricsContext getMetricsContext() {
+ return null;
+ }
+
+ @Override
+ public EventListener getEventListener() {
+ return null;
+ }
+
+ private int getAssignedSplitCount(int subtaskId) {
+ return assignments.getOrDefault(subtaskId,
Collections.emptyList()).size();
+ }
+ }
+}
