diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/LazyFlinkSourceSplitEnumerator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/LazyFlinkSourceSplitEnumerator.java index f298e6534b03..ec41680f2574 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/LazyFlinkSourceSplitEnumerator.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/LazyFlinkSourceSplitEnumerator.java @@ -109,7 +109,14 @@ private void initializeSplits() { this::splitBeamSource, (sourceSplits, error) -> { if (error != null) { - throw new RuntimeException("Failed to start source enumerator.", error); + throw new RuntimeException( + "Failed to initialize splits for bounded source " + beamSource.getClass().getName(), + error); + } + if (sourceSplits == null) { + throw new IllegalStateException( + "Split initialization returned null without an error for bounded source " + + beamSource.getClass().getName()); } pendingSplits.addAll(sourceSplits); splitsInitialized = true; diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumeratorTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumeratorTest.java index c4d220d22138..17a14c317bf8 100644 --- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumeratorTest.java +++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumeratorTest.java @@ -30,6 +30,7 @@ import java.util.List; import java.util.Map; import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BiConsumer; import org.apache.beam.runners.core.construction.SerializablePipelineOptions; import org.apache.beam.runners.flink.FlinkPipelineOptions; import org.apache.beam.runners.flink.translation.utils.SerdeUtils; @@ -43,9 +44,11 @@ import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.values.KV; import org.apache.flink.api.connector.source.SplitEnumerator; +import org.apache.flink.api.connector.source.SplitEnumeratorContext; import org.apache.flink.connector.testutils.source.reader.TestingSplitEnumeratorContext; import org.apache.flink.core.io.SimpleVersionedSerializer; import org.junit.Test; +import org.mockito.Mockito; /** Unit tests for the Flink source split enumerators. */ public class FlinkSourceSplitEnumeratorTest { @@ -54,6 +57,31 @@ public class FlinkSourceSplitEnumeratorTest { private static final int SOURCE_PARALLELISM = 2; private static final int REQUESTED_SPLITS = 4; + @Test + @SuppressWarnings("unchecked") + public void testLazyInitializationRejectsNullResultWithoutError() throws Exception { + SplitEnumeratorContext> context = + Mockito.mock(SplitEnumeratorContext.class); + Mockito.doAnswer( + invocation -> { + BiConsumer>, Throwable> callback = + invocation.getArgument(1); + callback.accept(null, null); + return null; + }) + .when(context) + .callAsync(Mockito.any(), Mockito.any()); + try (LazyFlinkSourceSplitEnumerator enumerator = + new LazyFlinkSourceSplitEnumerator<>( + context, + TestEstimatedSizeBoundedSource.create(1L, 1), + FlinkPipelineOptions.defaults(), + 1)) { + IllegalStateException thrown = assertThrows(IllegalStateException.class, enumerator::start); + assertTrue(thrown.getMessage().contains("returned null without an error")); + } + } + @Test public void testSelectsAssignmentModeFromEstimatedSizeAndConfiguration() throws Exception { assertEquals(0L, (long) FlinkPipelineOptions.defaults().getSourceStaticSplitThresholdMb());