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 7b67bc8134 [Fix][Core] Fix enumerator re-registration synchronization
(#10830)
7b67bc8134 is described below
commit 7b67bc81341583f33c379c56587e92dbd5a15087
Author: Jast <[email protected]>
AuthorDate: Wed May 13 20:53:16 2026 +0800
[Fix][Core] Fix enumerator re-registration synchronization (#10830)
---
.../jdbc/source/JdbcSourceSplitEnumeratorTest.java | 84 +++++++
.../server/task/SourceSplitEnumeratorTask.java | 16 +-
.../server/task/SourceSplitEnumeratorTaskTest.java | 253 +++++++++++++++++++++
.../flink/source/FlinkSourceEnumerator.java | 36 ++-
.../flink/source/FlinkSourceEnumeratorTest.java | 91 ++++++++
5 files changed, 462 insertions(+), 18 deletions(-)
diff --git
a/seatunnel-connectors-v2/connector-jdbc/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/source/JdbcSourceSplitEnumeratorTest.java
b/seatunnel-connectors-v2/connector-jdbc/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/source/JdbcSourceSplitEnumeratorTest.java
index 4e90b41e70..9f0748d320 100644
---
a/seatunnel-connectors-v2/connector-jdbc/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/source/JdbcSourceSplitEnumeratorTest.java
+++
b/seatunnel-connectors-v2/connector-jdbc/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/source/JdbcSourceSplitEnumeratorTest.java
@@ -203,6 +203,86 @@ class JdbcSourceSplitEnumeratorTest {
Assertions.assertEquals(0, enumerator.currentUnassignedSplitSize());
}
+ @Test
+ void testRegisterReaderReassignsReturnedSplitsWithoutResignaling() throws
Exception {
+ int parallelism = 1;
+ TablePath tablePath = TablePath.of("db", "schema", "table");
+
+ Map<TablePath, JdbcSourceTable> tables = new HashMap<>();
+ tables.put(tablePath, createJdbcSourceTable(tablePath));
+
+ Set<Integer> registeredReaders = new
HashSet<>(Collections.singleton(0));
+ List<List<String>> assignedSplitBatches = new ArrayList<>();
+ AtomicInteger noMoreSplitsCallCount = new AtomicInteger();
+
+ SourceSplitEnumerator.Context<JdbcSourceSplit> context =
+ new SourceSplitEnumerator.Context<JdbcSourceSplit>() {
+ @Override
+ public int currentParallelism() {
+ return parallelism;
+ }
+
+ @Override
+ public Set<Integer> registeredReaders() {
+ return new HashSet<>(registeredReaders);
+ }
+
+ @Override
+ public void assignSplit(int subtaskId,
List<JdbcSourceSplit> splits) {
+ assignedSplitBatches.add(
+ splits.stream()
+ .map(JdbcSourceSplit::splitId)
+
.collect(java.util.stream.Collectors.toList()));
+ }
+
+ @Override
+ public void signalNoMoreSplits(int subtask) {
+ noMoreSplitsCallCount.incrementAndGet();
+ }
+
+ @Override
+ public void sendEventToSourceReader(int subtaskId,
SourceEvent event) {}
+
+ @Override
+ public MetricsContext getMetricsContext() {
+ return null;
+ }
+
+ @Override
+ public EventListener getEventListener() {
+ return null;
+ }
+ };
+
+ JdbcSourceConfig sourceConfig =
+ JdbcSourceConfig.builder()
+ .jdbcConnectionConfig(
+ JdbcConnectionConfig.builder()
+ .url("jdbc:generic://localhost:0/test")
+ .driverName("org.example.Driver")
+ .build())
+ .build();
+
+ JdbcSourceSplitEnumerator enumerator =
+ new JdbcSourceSplitEnumerator(context, sourceConfig, tables,
null);
+
+ enumerator.open();
+ enumerator.run();
+
+ registeredReaders.clear();
+
+ JdbcSourceSplit returnedSplit = createSplit(tablePath,
"returned-split");
+ enumerator.addSplitsBack(Collections.singletonList(returnedSplit), 0);
+
+ registeredReaders.add(0);
+ enumerator.registerReader(0);
+
+ Assertions.assertEquals(2, assignedSplitBatches.size());
+ Assertions.assertEquals(
+ Collections.singletonList("returned-split"),
assignedSplitBatches.get(1));
+ Assertions.assertEquals(1, noMoreSplitsCallCount.get());
+ }
+
private JdbcSourceTable createJdbcSourceTable(TablePath tablePath) {
TableIdentifier tableId = TableIdentifier.of("default", tablePath);
TableSchema tableSchema =
TableSchema.builder().columns(Collections.emptyList()).build();
@@ -211,4 +291,8 @@ class JdbcSourceSplitEnumeratorTest {
tableId, tableSchema, Collections.emptyMap(),
Collections.emptyList(), "");
return
JdbcSourceTable.builder().tablePath(tablePath).catalogTable(catalogTable).build();
}
+
+ private JdbcSourceSplit createSplit(TablePath tablePath, String splitId) {
+ return new JdbcSourceSplit(tablePath, splitId, "SELECT 1", "id", null,
0, 1);
+ }
}
diff --git
a/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/task/SourceSplitEnumeratorTask.java
b/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/task/SourceSplitEnumeratorTask.java
index b22ca022d6..18e1e8217d 100644
---
a/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/task/SourceSplitEnumeratorTask.java
+++
b/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/task/SourceSplitEnumeratorTask.java
@@ -224,15 +224,17 @@ public class SourceSplitEnumeratorTask<SplitT extends
SourceSplit> extends Coord
SourceSplitEnumerator<SplitT, Serializable> enumerator =
getEnumerator();
int readerIndex = readerId.getTaskIndex();
+ boolean needResignalNoMoreSplits;
this.addTaskMemberMapping(readerId, memberAddr);
- synchronized (this) {
+ synchronized (enumeratorContext) {
enumerator.registerReader(readerIndex);
- if (enumeratorContext.hasNoMoreSplitsSignaled(readerIndex)) {
- log.info(
- "Reader [{}] re-registered after failover.
Re-signaling NoMoreSplitsEvent.",
- readerIndex);
- enumeratorContext.signalNoMoreSplits(readerIndex);
- }
+ needResignalNoMoreSplits =
enumeratorContext.hasNoMoreSplitsSignaled(readerIndex);
+ }
+ if (needResignalNoMoreSplits) {
+ log.info(
+ "Reader [{}] re-registered after failover. Re-signaling
NoMoreSplitsEvent.",
+ readerIndex);
+ enumeratorContext.signalNoMoreSplits(readerIndex);
}
int taskSize = taskMemberMapping.size();
if (maxReaderSize == taskSize) {
diff --git
a/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/task/SourceSplitEnumeratorTaskTest.java
b/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/task/SourceSplitEnumeratorTaskTest.java
index b74c8d1dba..0c5cd6bd83 100644
---
a/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/task/SourceSplitEnumeratorTaskTest.java
+++
b/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/task/SourceSplitEnumeratorTaskTest.java
@@ -18,6 +18,8 @@
package org.apache.seatunnel.engine.server.task;
import org.apache.seatunnel.api.source.SeaTunnelSource;
+import org.apache.seatunnel.api.source.SourceEvent;
+import org.apache.seatunnel.api.source.SourceSplit;
import org.apache.seatunnel.api.source.SourceSplitEnumerator;
import org.apache.seatunnel.engine.common.utils.concurrent.CompletableFuture;
import org.apache.seatunnel.engine.core.dag.actions.SourceAction;
@@ -26,22 +28,42 @@ import
org.apache.seatunnel.engine.server.execution.TaskExecutionContext;
import org.apache.seatunnel.engine.server.execution.TaskGroupLocation;
import org.apache.seatunnel.engine.server.execution.TaskLocation;
import
org.apache.seatunnel.engine.server.task.context.SeaTunnelSplitEnumeratorContext;
+import
org.apache.seatunnel.engine.server.task.operation.source.AssignSplitOperation;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
import org.mockito.Mockito;
import com.hazelcast.cluster.Address;
+import com.hazelcast.spi.impl.operationservice.Operation;
import com.hazelcast.spi.impl.operationservice.impl.InvocationFuture;
+import java.io.Serializable;
+import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
public class SourceSplitEnumeratorTaskTest {
+ private static final class DummySplit implements SourceSplit {
+ private static final long serialVersionUID = 1L;
+
+ @Override
+ public String splitId() {
+ return "dummy";
+ }
+ }
+
@Test
void testOpenShouldBeforeReaderRegister() throws Exception {
@@ -160,4 +182,235 @@ public class SourceSplitEnumeratorTaskTest {
Mockito.verify(context, Mockito.times(2)).sendToMember(Mockito.any(),
Mockito.any());
}
+
+ @Test
+ void testReceivedReaderUsesEnumeratorContextLock() throws Exception {
+ SeaTunnelSource<?, DummySplit, Serializable> source =
Mockito.mock(SeaTunnelSource.class);
+ LockAwareEnumerator enumerator = new LockAwareEnumerator();
+
+ AtomicReference<SeaTunnelSplitEnumeratorContext<DummySplit>>
enumeratorContextRef =
+ new AtomicReference<>();
+ Mockito.when(source.createEnumerator(Mockito.any()))
+ .thenAnswer(
+ invocation -> {
+
enumeratorContextRef.set(invocation.getArgument(0));
+ return enumerator;
+ });
+
+ SourceAction<?, DummySplit, Serializable> action =
+ new SourceAction<>(1, "fake", source, new HashSet<>(),
Collections.emptySet());
+ SourceSplitEnumeratorTask<DummySplit> enumeratorTask =
+ new SourceSplitEnumeratorTask<>(
+ 1, new TaskLocation(new TaskGroupLocation(1, 1, 1), 1,
1), action);
+
+ TaskExecutionContext context =
Mockito.mock(TaskExecutionContext.class);
+ InvocationFuture<Object> future = Mockito.mock(InvocationFuture.class);
+
Mockito.when(context.getOrCreateMetricsContext(Mockito.any())).thenReturn(null);
+ Mockito.when(context.sendToMaster(Mockito.any())).thenReturn(future);
+ Mockito.when(context.sendToMember(Mockito.any(),
Mockito.any())).thenReturn(future);
+ Mockito.when(future.join()).thenReturn(null);
+ TaskExecutionService taskExecutionService =
Mockito.mock(TaskExecutionService.class);
+
Mockito.when(context.getTaskExecutionService()).thenReturn(taskExecutionService);
+
+ enumeratorTask.setTaskExecutionContext(context);
+ enumeratorTask.init();
+ enumeratorTask.restoreState(new ArrayList<>());
+
+ TaskLocation readerLocation = new TaskLocation(new
TaskGroupLocation(1, 1, 1), 1, 1);
+ Address address = Address.createUnresolvedAddress("localhost", 5701);
+
+ SeaTunnelSplitEnumeratorContext<DummySplit> enumeratorContext =
enumeratorContextRef.get();
+ Assertions.assertNotNull(enumeratorContext);
+
+ ExecutorService executorService = Executors.newSingleThreadExecutor();
+ try {
+ synchronized (enumeratorContext) {
+ Future<?> blockedRegistration =
+ executorService.submit(
+ () -> {
+ try {
+
enumeratorTask.receivedReader(readerLocation, address);
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ });
+ Assertions.assertFalse(
+ enumerator.registerReaderCalled.await(200,
TimeUnit.MILLISECONDS));
+ blockedRegistration.cancel(true);
+ }
+
+ Future<?> registrationAfterUnlock =
+ executorService.submit(
+ () -> {
+ try {
+
enumeratorTask.receivedReader(readerLocation, address);
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ });
+
+ Assertions.assertTrue(enumerator.registerReaderCalled.await(1,
TimeUnit.SECONDS));
+ registrationAfterUnlock.get(1, TimeUnit.SECONDS);
+ } finally {
+ executorService.shutdownNow();
+ }
+ }
+
+ @Test
+ void
testResignalNoMoreSplitsOutsideLockWithPreciseAssignSplitVerification() throws
Exception {
+ SeaTunnelSource<?, DummySplit, Serializable> source =
Mockito.mock(SeaTunnelSource.class);
+ SourceSplitEnumerator<DummySplit, Serializable> enumerator =
+ Mockito.mock(SourceSplitEnumerator.class);
+
+ AtomicReference<SeaTunnelSplitEnumeratorContext<DummySplit>>
enumeratorContextRef =
+ new AtomicReference<>();
+ Mockito.when(source.createEnumerator(Mockito.any()))
+ .thenAnswer(
+ invocation -> {
+
enumeratorContextRef.set(invocation.getArgument(0));
+ return enumerator;
+ });
+
+ CountDownLatch blockingJoinEntered = new CountDownLatch(1);
+ CountDownLatch allowBlockingJoin = new CountDownLatch(1);
+ CountDownLatch secondReaderRegistered = new CountDownLatch(1);
+
+ Mockito.doAnswer(
+ invocation -> {
+ if (invocation.getArgument(0, Integer.class) == 2)
{
+ secondReaderRegistered.countDown();
+ }
+ return null;
+ })
+ .when(enumerator)
+ .registerReader(Mockito.anyInt());
+
+ SourceAction<?, DummySplit, Serializable> action =
+ new SourceAction<>(1, "fake", source, new HashSet<>(),
Collections.emptySet());
+ SourceSplitEnumeratorTask<DummySplit> enumeratorTask =
+ new SourceSplitEnumeratorTask<>(
+ 1, new TaskLocation(new TaskGroupLocation(1, 1, 1), 1,
1), action);
+
+ TaskExecutionContext context =
Mockito.mock(TaskExecutionContext.class);
+ InvocationFuture<Object> future = Mockito.mock(InvocationFuture.class);
+
Mockito.when(context.getOrCreateMetricsContext(Mockito.any())).thenReturn(null);
+ Mockito.when(context.sendToMaster(Mockito.any())).thenReturn(future);
+ Mockito.when(context.sendToMember(Mockito.any(),
Mockito.any())).thenReturn(future);
+ Mockito.when(future.join())
+ .thenAnswer(
+ invocation -> {
+ blockingJoinEntered.countDown();
+ allowBlockingJoin.await(5, TimeUnit.SECONDS);
+ return null;
+ });
+ TaskExecutionService taskExecutionService =
Mockito.mock(TaskExecutionService.class);
+
Mockito.when(context.getTaskExecutionService()).thenReturn(taskExecutionService);
+
+ enumeratorTask.setTaskExecutionContext(context);
+ enumeratorTask.init();
+ enumeratorTask.restoreState(new ArrayList<>());
+
+ SeaTunnelSplitEnumeratorContext<DummySplit> enumeratorContext =
enumeratorContextRef.get();
+ Assertions.assertNotNull(enumeratorContext);
+ getNoMoreSplitsSignaledReaders(enumeratorContext).add(1);
+
+ TaskLocation firstReader = new TaskLocation(new TaskGroupLocation(1,
1, 1), 1, 1);
+ TaskLocation secondReader = new TaskLocation(new TaskGroupLocation(1,
1, 1), 1, 2);
+ Address firstAddress = Address.createUnresolvedAddress("localhost",
5701);
+ Address secondAddress = Address.createUnresolvedAddress("localhost",
5702);
+
+ ExecutorService executorService = Executors.newFixedThreadPool(2);
+ try {
+ Future<?> blockingFuture =
+ executorService.submit(
+ () -> {
+ try {
+ enumeratorTask.receivedReader(firstReader,
firstAddress);
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ });
+
+ Assertions.assertTrue(blockingJoinEntered.await(1,
TimeUnit.SECONDS));
+
+ Future<?> concurrentRegistration =
+ executorService.submit(
+ () -> {
+ try {
+
enumeratorTask.receivedReader(secondReader, secondAddress);
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ });
+
+ Assertions.assertTrue(secondReaderRegistered.await(1,
TimeUnit.SECONDS));
+
+ allowBlockingJoin.countDown();
+ blockingFuture.get(1, TimeUnit.SECONDS);
+ concurrentRegistration.get(1, TimeUnit.SECONDS);
+ } finally {
+ executorService.shutdownNow();
+ }
+
+ ArgumentCaptor<Operation> operationCaptor =
ArgumentCaptor.forClass(Operation.class);
+ Mockito.verify(context).sendToMember(operationCaptor.capture(),
Mockito.eq(firstAddress));
+ Operation operation = operationCaptor.getValue();
+ Assertions.assertInstanceOf(AssignSplitOperation.class, operation);
+ Assertions.assertEquals(firstReader, readField(operation, "taskID"));
+ Assertions.assertTrue(((java.util.List<?>) readField(operation,
"splits")).isEmpty());
+ }
+
+ @SuppressWarnings("unchecked")
+ private Set<Integer> getNoMoreSplitsSignaledReaders(
+ SeaTunnelSplitEnumeratorContext<DummySplit> enumeratorContext)
throws Exception {
+ return (Set<Integer>) readField(enumeratorContext,
"noMoreSplitsSignaledReaders");
+ }
+
+ private static Object readField(Object target, String fieldName) throws
Exception {
+ Field field = target.getClass().getDeclaredField(fieldName);
+ field.setAccessible(true);
+ return field.get(target);
+ }
+
+ private static final class LockAwareEnumerator
+ implements SourceSplitEnumerator<DummySplit, Serializable> {
+
+ private final CountDownLatch registerReaderCalled = new
CountDownLatch(1);
+
+ @Override
+ public void open() {}
+
+ @Override
+ public void run() {}
+
+ @Override
+ public void close() {}
+
+ @Override
+ public void addSplitsBack(java.util.List<DummySplit> splits, int
subtaskId) {}
+
+ @Override
+ public int currentUnassignedSplitSize() {
+ return 0;
+ }
+
+ @Override
+ public void handleSplitRequest(int subtaskId) {}
+
+ @Override
+ public void registerReader(int subtaskId) {
+ registerReaderCalled.countDown();
+ }
+
+ @Override
+ public void handleSourceEvent(int subtaskId, SourceEvent sourceEvent)
{}
+
+ @Override
+ public Serializable snapshotState(long checkpointId) {
+ return null;
+ }
+
+ @Override
+ public void notifyCheckpointComplete(long checkpointId) {}
+ }
}
diff --git
a/seatunnel-translation/seatunnel-translation-flink/seatunnel-translation-flink-common/src/main/java/org/apache/seatunnel/translation/flink/source/FlinkSourceEnumerator.java
b/seatunnel-translation/seatunnel-translation-flink/seatunnel-translation-flink-common/src/main/java/org/apache/seatunnel/translation/flink/source/FlinkSourceEnumerator.java
index bd3dba9c28..960536dcdb 100644
---
a/seatunnel-translation/seatunnel-translation-flink/seatunnel-translation-flink-common/src/main/java/org/apache/seatunnel/translation/flink/source/FlinkSourceEnumerator.java
+++
b/seatunnel-translation/seatunnel-translation-flink/seatunnel-translation-flink-common/src/main/java/org/apache/seatunnel/translation/flink/source/FlinkSourceEnumerator.java
@@ -32,9 +32,9 @@ import org.slf4j.LoggerFactory;
import javax.annotation.Nullable;
import java.io.IOException;
+import java.util.HashSet;
import java.util.List;
import java.util.Set;
-import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;
/**
@@ -58,10 +58,10 @@ public class FlinkSourceEnumerator<SplitT extends
SourceSplit, EnumStateT>
private final Set<Integer> noMoreSplitsSignaledReaders;
private final Object lock = new Object();
+ private final Set<Integer> registeredReaderIds = new HashSet<>();
- private AtomicBoolean isRun = new AtomicBoolean(false);
-
- private volatile int currentRegisterReaders = 0;
+ private volatile boolean isRun = false;
+ private volatile boolean isRunning = false;
public FlinkSourceEnumerator(
SourceSplitEnumerator<SplitT, EnumStateT> enumerator,
@@ -96,23 +96,37 @@ public class FlinkSourceEnumerator<SplitT extends
SourceSplit, EnumStateT>
@Override
public void addReader(int subtaskId) {
+ boolean needResignalNoMoreSplits;
+ boolean shouldRun = false;
synchronized (lock) {
sourceSplitEnumerator.registerReader(subtaskId);
- currentRegisterReaders++;
- if (noMoreSplitsSignaledReaders.contains(subtaskId)) {
- LOGGER.info(
- "Reader [{}] re-registered after failover.
Re-signaling NoMoreSplitsEvent.",
- subtaskId);
- enumeratorContext.signalNoMoreSplits(subtaskId);
+ registeredReaderIds.add(subtaskId);
+ needResignalNoMoreSplits =
noMoreSplitsSignaledReaders.contains(subtaskId);
+ if (!isRun && !isRunning && registeredReaderIds.size() ==
parallelism) {
+ shouldRun = true;
+ isRunning = true;
}
}
- if (currentRegisterReaders == parallelism && !isRun.getAndSet(true)) {
+ if (shouldRun) {
try {
sourceSplitEnumerator.run();
+ synchronized (lock) {
+ isRun = true;
+ isRunning = false;
+ }
} catch (Exception e) {
+ synchronized (lock) {
+ isRunning = false;
+ }
throw new RuntimeException(e);
}
}
+ if (needResignalNoMoreSplits) {
+ LOGGER.info(
+ "Reader [{}] re-registered after failover. Re-signaling
NoMoreSplitsEvent.",
+ subtaskId);
+ enumeratorContext.signalNoMoreSplits(subtaskId);
+ }
}
@Override
diff --git
a/seatunnel-translation/seatunnel-translation-flink/seatunnel-translation-flink-common/src/test/java/org/apache/seatunnel/translation/flink/source/FlinkSourceEnumeratorTest.java
b/seatunnel-translation/seatunnel-translation-flink/seatunnel-translation-flink-common/src/test/java/org/apache/seatunnel/translation/flink/source/FlinkSourceEnumeratorTest.java
index 949a66f838..83ca41af73 100644
---
a/seatunnel-translation/seatunnel-translation-flink/seatunnel-translation-flink-common/src/test/java/org/apache/seatunnel/translation/flink/source/FlinkSourceEnumeratorTest.java
+++
b/seatunnel-translation/seatunnel-translation-flink/seatunnel-translation-flink-common/src/test/java/org/apache/seatunnel/translation/flink/source/FlinkSourceEnumeratorTest.java
@@ -22,12 +22,18 @@ import
org.apache.seatunnel.api.source.SourceSplitEnumerator;
import org.apache.flink.api.connector.source.SplitEnumeratorContext;
+import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import org.mockito.Mockito;
import java.io.Serializable;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
class FlinkSourceEnumeratorTest {
@@ -59,4 +65,89 @@ class FlinkSourceEnumeratorTest {
Mockito.verify(enumeratorContext).signalNoMoreSplits(0);
}
+
+ @Test
+ void testDuplicateReaderRegistrationDoesNotStartEnumeratorEarly() throws
Exception {
+ SourceSplitEnumerator<DummySplit, Serializable> sourceSplitEnumerator =
+ Mockito.mock(SourceSplitEnumerator.class);
+ SplitEnumeratorContext<SplitWrapper<DummySplit>> enumeratorContext =
+ Mockito.mock(SplitEnumeratorContext.class);
+ Mockito.when(enumeratorContext.currentParallelism()).thenReturn(2);
+
+ FlinkSourceEnumerator<DummySplit, Serializable> enumerator =
+ new FlinkSourceEnumerator<>(
+ sourceSplitEnumerator, enumeratorContext,
ConcurrentHashMap.newKeySet());
+
+ enumerator.addReader(0);
+ enumerator.addReader(0);
+
+ Mockito.verify(sourceSplitEnumerator, Mockito.never()).run();
+
+ enumerator.addReader(1);
+
+ Mockito.verify(sourceSplitEnumerator).run();
+ }
+
+ @Test
+ void testRunFailureCanRetryOnReaderReregister() throws Exception {
+ SourceSplitEnumerator<DummySplit, Serializable> sourceSplitEnumerator =
+ Mockito.mock(SourceSplitEnumerator.class);
+ SplitEnumeratorContext<SplitWrapper<DummySplit>> enumeratorContext =
+ Mockito.mock(SplitEnumeratorContext.class);
+ Mockito.when(enumeratorContext.currentParallelism()).thenReturn(1);
+ Mockito.doThrow(new RuntimeException("run failed"))
+ .doNothing()
+ .when(sourceSplitEnumerator)
+ .run();
+
+ FlinkSourceEnumerator<DummySplit, Serializable> enumerator =
+ new FlinkSourceEnumerator<>(
+ sourceSplitEnumerator, enumeratorContext,
ConcurrentHashMap.newKeySet());
+
+ Assertions.assertThrows(RuntimeException.class, () ->
enumerator.addReader(0));
+
+ enumerator.addReader(0);
+
+ Mockito.verify(sourceSplitEnumerator, Mockito.times(2)).run();
+ }
+
+ @Test
+ void testSnapshotStateDoesNotWaitForBlockingRun() throws Exception {
+ SourceSplitEnumerator<DummySplit, Serializable> sourceSplitEnumerator =
+ Mockito.mock(SourceSplitEnumerator.class);
+ SplitEnumeratorContext<SplitWrapper<DummySplit>> enumeratorContext =
+ Mockito.mock(SplitEnumeratorContext.class);
+ Mockito.when(enumeratorContext.currentParallelism()).thenReturn(1);
+
Mockito.when(sourceSplitEnumerator.snapshotState(1L)).thenReturn("checkpoint");
+
+ CountDownLatch runEntered = new CountDownLatch(1);
+ CountDownLatch releaseRun = new CountDownLatch(1);
+ Mockito.doAnswer(
+ invocation -> {
+ runEntered.countDown();
+ releaseRun.await(5, TimeUnit.SECONDS);
+ return null;
+ })
+ .when(sourceSplitEnumerator)
+ .run();
+
+ FlinkSourceEnumerator<DummySplit, Serializable> enumerator =
+ new FlinkSourceEnumerator<>(
+ sourceSplitEnumerator, enumeratorContext,
ConcurrentHashMap.newKeySet());
+
+ ExecutorService executorService = Executors.newFixedThreadPool(2);
+ Future<?> addReaderFuture = executorService.submit(() ->
enumerator.addReader(0));
+ try {
+ Assertions.assertTrue(runEntered.await(1, TimeUnit.SECONDS));
+
+ Future<Serializable> snapshotFuture =
+ executorService.submit(() -> enumerator.snapshotState(1L));
+
+ Assertions.assertEquals("checkpoint", snapshotFuture.get(1,
TimeUnit.SECONDS));
+ } finally {
+ releaseRun.countDown();
+ addReaderFuture.get(1, TimeUnit.SECONDS);
+ executorService.shutdownNow();
+ }
+ }
}
