This is an automated email from the ASF dual-hosted git repository. MartijnVisser pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/flink-connector-kafka.git
commit 01e1a7744ebdb64de3cb00b50d2603e48227a3e0 Author: MartijnVisser <[email protected]> AuthorDate: Sat Oct 24 15:20:16 2026 +0200 [FLINK-31006][Connectors/Kafka] Signal no more splits after a restored bounded enumerator finds no partition changes When the enumerator is recreated from state, its one-time discovery finds every subscribed partition already assigned, so checkPartitionChanges returns early on the empty change and never reaches handlePartitionSplitChanges, the only place that sets noMoreNewPartitionSplits. The re-registered readers therefore wait for a NoMoreSplitsEvent that never arrives and a bounded source never finishes. The same early return hangs a bounded source that subscribes to no existing partition on a fresh start. Let an empty change through for a bounded source with one-time discovery, so that it reaches handlePartitionSplitChanges and signals the registered readers. Unbounded sources keep returning early: treating an empty change as a finished discovery would make partitions found later start from the earliest offset rather than the configured one. A bounded source that started without partitions now snapshots initialDiscoveryFinished as true too, but it finishes straight away, so a restore from that state is not expected in practice. Signalling from addReader on the restored initialDiscoveryFinished flag would finish a reader before the post-restore discovery assigns it a partition created while the job was down, which is why apache/flink#21909 was rejected. testRestoredBoundedEnumeratorAssignsBeforeItSignals pins that ordering. The early return came in with 0b38d1a9; before it, the discovery callback always ran. KafkaSourceBoundedRestoreITCase covers the reported symptom over stop-with-savepoint and restore. In KafkaSourceEnumeratorTest, a parameterised test covers the restore and the empty subscription, and two more pin the unbounded source that keeps skipping an empty change and the assign-before-signal ordering. DynamicKafkaSourceEnumeratorTest covers the same restore through the sub-enumerators. Generated-by: Claude Code (Claude Opus 5.5) --- .../source/reader/DynamicKafkaSourceReader.java | 7 +- .../source/enumerator/KafkaSourceEnumerator.java | 12 +- .../DynamicKafkaSourceEnumeratorTest.java | 38 ++- .../source/KafkaSourceBoundedRestoreITCase.java | 255 +++++++++++++++++++++ .../enumerator/KafkaSourceEnumeratorTest.java | 195 +++++++++++++++- 5 files changed, 498 insertions(+), 9 deletions(-) diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/reader/DynamicKafkaSourceReader.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/reader/DynamicKafkaSourceReader.java index 0d2c9c73..f1b6174a 100644 --- a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/reader/DynamicKafkaSourceReader.java +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/dynamic/source/reader/DynamicKafkaSourceReader.java @@ -405,11 +405,8 @@ public class DynamicKafkaSourceReader<T> implements SourceReader<T, DynamicKafka pendingSplits.clear(); } - // Replay only on the first metadata update. On a later metadata change the reader must - // wait for the enumerator to signal again after the new assignments (that re-signal is - // currently missing for a sub-enumerator recreated with partitions already assigned; see - // FLINK-31006), so replaying here would finish an active reader before the new topic's - // splits arrive. + // Replay only on the first metadata update. A later metadata change can bring new splits, + // so replaying the earlier signal here would finish an active reader before they arrive. if (isNoMoreSplits && firstMetadataUpdate) { notifyNoMoreSplits(); } diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumerator.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumerator.java index e66a73c3..e28fc5cc 100644 --- a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumerator.java +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumerator.java @@ -144,7 +144,10 @@ public class KafkaSourceEnumerator // This flag will be marked as true if periodically partition discovery is disabled AND the // initializing partition discovery has finished. private boolean noMoreNewPartitionSplits = false; - // this flag will be marked as true if initial partitions are discovered after enumerator starts + // Set once a discovery result reaches handlePartitionSplitChanges: a non-empty one, or an empty + // one of a bounded source with one-time discovery. Partitions found after that start from the + // earliest offset rather than the configured one, so an unbounded source's empty change must + // not set it. private boolean initialDiscoveryFinished; private final boolean topicIntegrityCheckEnabled; @@ -365,7 +368,12 @@ public class KafkaSourceEnumerator final PartitionChange partitionChange = getPartitionChange(fetchedPartitions, !initialDiscoveryFinished); - if (partitionChange.isEmpty()) { + // A bounded source must reach handlePartitionSplitChanges even when nothing changed, as + // that is where its readers are signalled. After a restore every subscribed partition is + // already assigned, so the change is empty. + final boolean boundedOneTimeDiscovery = + boundedness == Boundedness.BOUNDED && partitionDiscoveryIntervalMs <= 0; + if (partitionChange.isEmpty() && !boundedOneTimeDiscovery) { return; } context.callAsync( diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/DynamicKafkaSourceEnumeratorTest.java b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/DynamicKafkaSourceEnumeratorTest.java index e100aedf..68de569a 100644 --- a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/DynamicKafkaSourceEnumeratorTest.java +++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/dynamic/source/enumerator/DynamicKafkaSourceEnumeratorTest.java @@ -2160,6 +2160,12 @@ public class DynamicKafkaSourceEnumeratorTest { private DynamicKafkaSourceEnumerator createBoundedEnumerator( SplitEnumeratorContext<DynamicKafkaSourceSplit> context) { + return createBoundedEnumerator(context, new DynamicKafkaSourceEnumState()); + } + + private DynamicKafkaSourceEnumerator createBoundedEnumerator( + SplitEnumeratorContext<DynamicKafkaSourceSplit> context, + DynamicKafkaSourceEnumState enumState) { Properties properties = new Properties(); properties.setProperty(KafkaSourceOptions.PARTITION_DISCOVERY_INTERVAL_MS.key(), "0"); properties.setProperty( @@ -2173,7 +2179,7 @@ public class DynamicKafkaSourceEnumeratorTest { OffsetsInitializer.latest(), properties, Boundedness.BOUNDED, - new DynamicKafkaSourceEnumState(), + enumState, new TestKafkaEnumContextProxyFactory()); } @@ -2202,6 +2208,36 @@ public class DynamicKafkaSourceEnumeratorTest { applyPropertiesConsumer); } + /** + * A bounded DynamicKafkaSource restored with every partition already assigned must still tell + * its readers that no more splits are coming. Each sub-enumerator runs its one-time discovery + * and finds nothing new. One that returns early on that empty change never marks the discovery + * as finished, so the readers wait forever and the job never finishes. + */ + @Test + public void testBoundedRestoreSignalsNoMoreSplitsWithoutPartitionChanges() throws Throwable { + try (MockSplitEnumeratorContext<DynamicKafkaSourceSplit> context = + new MockSplitEnumeratorContext<>(NUM_SUBTASKS); + DynamicKafkaSourceEnumerator enumerator = + createBoundedEnumerator(context, getCheckpointState())) { + enumerator.start(); + + for (int reader = 0; reader < NUM_SUBTASKS; reader++) { + mockRegisterReaderAndSendReaderStartupEvent(context, enumerator, reader); + } + runAllOneTimeCallables(context); + + assertThat(context.getSplitsAssignmentSequence()) + .as("Every partition is already assigned, so nothing may be assigned again") + .isEmpty(); + for (int reader = 0; reader < NUM_SUBTASKS; reader++) { + assertThat(context.hasNoMoreSplits(reader)) + .as("Reader %s must be told that no more splits are coming", reader) + .isTrue(); + } + } + } + private DynamicKafkaSourceEnumerator createEnumerator( SplitEnumeratorContext<DynamicKafkaSourceSplit> context, KafkaMetadataService kafkaMetadataService, diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/source/KafkaSourceBoundedRestoreITCase.java b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/source/KafkaSourceBoundedRestoreITCase.java new file mode 100644 index 00000000..d7b61346 --- /dev/null +++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/source/KafkaSourceBoundedRestoreITCase.java @@ -0,0 +1,255 @@ +/* + * 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.flink.connector.kafka.source; + +import org.apache.flink.api.common.JobID; +import org.apache.flink.api.common.eventtime.WatermarkStrategy; +import org.apache.flink.api.common.functions.MapFunction; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.RestartStrategyOptions; +import org.apache.flink.configuration.StateRecoveryOptions; +import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer; +import org.apache.flink.connector.kafka.source.reader.deserializer.KafkaRecordDeserializationSchema; +import org.apache.flink.connector.kafka.testutils.KafkaSourceTestEnv; +import org.apache.flink.core.execution.SavepointFormatType; +import org.apache.flink.core.testutils.CommonTestUtils; +import org.apache.flink.runtime.client.JobStatusMessage; +import org.apache.flink.runtime.jobgraph.JobGraph; +import org.apache.flink.runtime.jobmaster.JobResult; +import org.apache.flink.runtime.minicluster.MiniCluster; +import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration; +import org.apache.flink.streaming.api.datastream.DataStream; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.streaming.api.functions.sink.v2.DiscardingSink; +import org.apache.flink.test.junit5.InjectMiniCluster; +import org.apache.flink.test.junit5.MiniClusterExtension; + +import org.apache.kafka.clients.producer.ProducerRecord; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.serialization.IntegerDeserializer; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.api.parallel.ResourceLock; + +import java.nio.file.Path; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * A bounded {@link KafkaSource} must still finish when its enumerator is recreated from state. This + * test restores from a savepoint. A retained checkpoint, a JobManager failover and a global + * failover after a completed checkpoint recreate the enumerator from state the same way; a regional + * failover keeps the running enumerator and is not affected. + * + * <p>The recreated enumerator finds every subscribed partition already assigned. On an unfixed + * enumerator that empty partition change makes {@code checkPartitionChanges} return before it + * reaches the only place that marks the discovery as finished, so the restored readers consume up + * to their stopping offset and then wait forever for a {@code NoMoreSplitsEvent} that is never + * sent. + * + * <p>Do not run this test in batch mode, which removes the checkpoint coordinator that the + * savepoint needs. + */ +@ResourceLock("KafkaTestBase") +public class KafkaSourceBoundedRestoreITCase { + + private static final String TOPIC = "KafkaSourceBoundedRestoreITCase-topic"; + + /** Records produced before the savepoint. */ + private static final int RECORDS_BEFORE_SAVEPOINT = 10; + + /** Records produced while the job is stopped, which the restored job must still consume. */ + private static final int RECORDS_AFTER_SAVEPOINT = 5; + + private static final int TOTAL_RECORDS = RECORDS_BEFORE_SAVEPOINT + RECORDS_AFTER_SAVEPOINT; + private static final Duration TIMEOUT = Duration.ofMinutes(2); + + /** Values seen by the pipeline in the current run, in the order they were seen. */ + private static final List<Integer> COLLECTED = new CopyOnWriteArrayList<>(); + + @TempDir private Path savepointBasePath; + + @RegisterExtension + public static final MiniClusterExtension MINI_CLUSTER = + new MiniClusterExtension( + new MiniClusterResourceConfiguration.Builder() + .setNumberTaskManagers(1) + .setNumberSlotsPerTaskManager(2) + .build()); + + @BeforeEach + public void setup() throws Throwable { + COLLECTED.clear(); + KafkaSourceTestEnv.setup(); + KafkaSourceTestEnv.createTestTopic(TOPIC, 1, 1); + produce(0, RECORDS_BEFORE_SAVEPOINT); + } + + @AfterEach + public void tearDown(@InjectMiniCluster MiniCluster miniCluster) throws Exception { + try { + // A job left running by a failed test would otherwise lose its broker under it. + for (JobStatusMessage job : + miniCluster.listJobs().get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS)) { + if (!job.getJobState().isGloballyTerminalState()) { + miniCluster + .cancelJob(job.getJobId()) + .get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS); + miniCluster + .requestJobResult(job.getJobId()) + .get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS); + } + } + } finally { + KafkaSourceTestEnv.tearDown(); + } + } + + private static void produce(int fromValue, int count) throws Throwable { + List<ProducerRecord<String, Integer>> records = new ArrayList<>(); + for (int i = fromValue; i < fromValue + count; i++) { + records.add(new ProducerRecord<>(TOPIC, 0, "key-" + i, i)); + } + KafkaSourceTestEnv.produceToKafka(records); + } + + /** + * The job graph is built by one helper for both runs so that the restored job has the identical + * topology, and therefore the identical operator IDs, as the job the savepoint was taken from. + */ + private JobGraph getJobGraph(Configuration extraConf) { + KafkaSource<Integer> source = + KafkaSource.<Integer>builder() + .setBootstrapServers(KafkaSourceTestEnv.brokerConnectionStrings) + .setTopics(TOPIC) + .setGroupId("KafkaSourceBoundedRestoreITCase") + .setStartingOffsets(OffsetsInitializer.earliest()) + // Beyond what has been produced so far, so the first run cannot finish on + // its own and the savepoint is always taken from a running source. + .setBounded( + OffsetsInitializer.offsets( + Collections.singletonMap( + new TopicPartition(TOPIC, 0), + (long) TOTAL_RECORDS))) + .setDeserializer( + KafkaRecordDeserializationSchema.valueOnly( + IntegerDeserializer.class)) + .build(); + + Configuration configuration = new Configuration(); + configuration.addAll(extraConf); + // A hang must surface as a timeout on the job, not as restart churn. + configuration.set(RestartStrategyOptions.RESTART_STRATEGY, "disable"); + + StreamExecutionEnvironment env = + StreamExecutionEnvironment.getExecutionEnvironment(configuration); + env.setParallelism(1); + DataStream<Integer> stream = + env.fromSource(source, WatermarkStrategy.noWatermarks(), "kafka-source") + .uid("kafka-source") + .map( + (MapFunction<Integer, Integer>) + value -> { + COLLECTED.add(value); + return value; + }) + .uid("collector"); + stream.sinkTo(new DiscardingSink<>()).uid("sink"); + return env.getStreamGraph().getJobGraph(); + } + + /** + * Waits for the job to finish. A job that fails or is cancelled instead rethrows its own cause, + * so that an unrelated failure is reported as itself rather than as the hang under test. + */ + private static void awaitJobFinished(MiniCluster miniCluster, JobID jobId) throws Exception { + final JobResult jobResult; + try { + jobResult = + miniCluster + .requestJobResult(jobId) + .get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS); + } catch (TimeoutException e) { + throw new AssertionError( + "The restored bounded job did not finish within " + TIMEOUT, e); + } + jobResult.toJobExecutionResult(KafkaSourceBoundedRestoreITCase.class.getClassLoader()); + } + + @Test + public void testBoundedSourceFinishesAfterRestoreFromSavepoint( + @InjectMiniCluster MiniCluster miniCluster) throws Throwable { + JobGraph firstJobGraph = getJobGraph(new Configuration()); + JobID firstJobId = firstJobGraph.getJobID(); + miniCluster.submitJob(firstJobGraph).get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS); + + CommonTestUtils.waitUtil( + () -> COLLECTED.size() >= RECORDS_BEFORE_SAVEPOINT, + TIMEOUT, + Duration.ofMillis(50), + "The first run did not consume the records produced before the savepoint"); + + String savepointPath = + miniCluster + .stopWithSavepoint( + firstJobId, + savepointBasePath.toFile().toString(), + false, + SavepointFormatType.CANONICAL) + .get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS); + assertThat(savepointPath).isNotBlank(); + assertThat(COLLECTED) + .as("The first run must consume the records produced before the savepoint once") + .containsExactlyElementsOf( + IntStream.range(0, RECORDS_BEFORE_SAVEPOINT) + .boxed() + .collect(Collectors.toList())); + COLLECTED.clear(); + + produce(RECORDS_BEFORE_SAVEPOINT, RECORDS_AFTER_SAVEPOINT); + + Configuration restoreConf = new Configuration(); + restoreConf.set(StateRecoveryOptions.SAVEPOINT_PATH, savepointPath); + JobGraph secondJobGraph = getJobGraph(restoreConf); + JobID secondJobId = secondJobGraph.getJobID(); + miniCluster.submitJob(secondJobGraph).get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS); + + // The restored source reaches its stopping offset and then needs the enumerator to tell it + // that no more splits are coming. Without that signal the job hangs here. + awaitJobFinished(miniCluster, secondJobId); + + assertThat(COLLECTED) + .as("The restored run must resume from the savepoint without replaying records") + .containsExactlyElementsOf( + IntStream.range(RECORDS_BEFORE_SAVEPOINT, TOTAL_RECORDS) + .boxed() + .collect(Collectors.toList())); + } +} diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumeratorTest.java b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumeratorTest.java index 07c513ef..cb7d1385 100644 --- a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumeratorTest.java +++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/source/enumerator/KafkaSourceEnumeratorTest.java @@ -44,7 +44,9 @@ import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.parallel.ResourceLock; import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.MethodSource; import java.util.ArrayList; import java.util.Arrays; @@ -59,6 +61,7 @@ import java.util.Set; import java.util.StringJoiner; import java.util.regex.Pattern; import java.util.stream.Collectors; +import java.util.stream.Stream; import static org.apache.flink.connector.kafka.source.split.KafkaPartitionSplit.MIGRATED; import static org.assertj.core.api.Assertions.assertThat; @@ -765,6 +768,29 @@ public class KafkaSourceEnumeratorTest { Properties overrideProperties, OffsetsInitializer startingOffsetsInitializer, boolean initialDiscoveryFinished) { + return createEnumerator( + enumContext, + partitionDiscoveryIntervalMs, + topicsToSubscribe, + assignedSplits, + unassignedInitialSplits, + overrideProperties, + startingOffsetsInitializer, + Boundedness.CONTINUOUS_UNBOUNDED, + initialDiscoveryFinished); + } + + /** Create the enumerator, controlling the boundedness of the source. */ + private KafkaSourceEnumerator createEnumerator( + MockSplitEnumeratorContext<KafkaPartitionSplit> enumContext, + long partitionDiscoveryIntervalMs, + Collection<String> topicsToSubscribe, + Collection<KafkaPartitionSplit> assignedSplits, + Collection<KafkaPartitionSplit> unassignedInitialSplits, + Properties overrideProperties, + OffsetsInitializer startingOffsetsInitializer, + Boundedness boundedness, + boolean initialDiscoveryFinished) { // Use a TopicPatternSubscriber so that no exception if a subscribed topic hasn't been // created yet. StringJoiner topicNameJoiner = new StringJoiner("|"); @@ -787,7 +813,7 @@ public class KafkaSourceEnumeratorTest { stoppingOffsetsInitializer, props, enumContext, - Boundedness.CONTINUOUS_UNBOUNDED, + boundedness, new KafkaSourceEnumState( assignedSplits, unassignedInitialSplits, initialDiscoveryFinished)); } @@ -979,6 +1005,173 @@ public class KafkaSourceEnumeratorTest { } } + /** + * A bounded enumerator whose one-time discovery finds no partition change must still tell its + * readers that no more splits are coming, both when it is restored with every subscribed + * partition already assigned and when it starts without any partition to subscribe to. An + * unfixed enumerator returns early on the empty change, before it reaches the only place that + * sets {@code noMoreNewPartitionSplits}, so the readers wait for a {@code NoMoreSplitsEvent} + * that never arrives and the job never finishes. + * + * <p>The restored case must hold every subscribed partition. With one left out its partition + * change is no longer empty, and it passes without the fix. + */ + @ParameterizedTest(name = "{0}") + @MethodSource("boundedSourcesWithoutPartitionChanges") + public void testBoundedEnumeratorSignalsNoMoreSplitsWithoutPartitionChanges( + String description, + Collection<String> topics, + Set<KafkaPartitionSplit> restoredAssignedSplits, + boolean initialDiscoveryFinished) + throws Throwable { + try (RecordingSplitEnumeratorContext context = + new RecordingSplitEnumeratorContext(NUM_SUBTASKS); + KafkaSourceEnumerator enumerator = + createEnumerator( + context, + -1, + topics, + restoredAssignedSplits, + Collections.emptySet(), + new Properties(), + OffsetsInitializer.earliest(), + Boundedness.BOUNDED, + initialDiscoveryFinished)) { + enumerator.start(); + + // A reader that registers before the discovery callback returns. + registerReader(context, enumerator, READER0); + runOneTimePartitionDiscovery(context); + // A reader that registers after it. + registerReader(context, enumerator, READER1); + + assertThat(context.events) + .as("Each reader is told exactly once that the input ended, nothing else") + .containsExactly( + "signalNoMoreSplits:" + READER0, "signalNoMoreSplits:" + READER1); + } + } + + private static Stream<Arguments> boundedSourcesWithoutPartitionChanges() { + return Stream.of( + Arguments.of( + "restored with every partition assigned", + PRE_EXISTING_TOPICS, + KafkaSourceTestEnv.getPartitionsForTopics(PRE_EXISTING_TOPICS).stream() + .map(tp -> new KafkaPartitionSplit(tp, 0L)) + .collect(Collectors.toSet()), + true), + // A single explicit name, because the helper joins the names into a pattern and an + // empty collection would compile to a pattern that matches every topic. + Arguments.of( + "started without any partition", + Collections.singleton("topic-that-does-not-exist"), + Collections.emptySet(), + false)); + } + + /** + * Unlike the bounded source in {@link + * #testBoundedEnumeratorSignalsNoMoreSplitsWithoutPartitionChanges}, an unbounded source with + * one-time discovery must keep returning early on an empty partition change. It has nothing to + * tell its readers, and treating the empty change as a finished discovery would make a later + * restore initialise the partitions that appeared meanwhile from the earliest offset instead of + * the configured one. This test fails if the condition in checkPartitionChanges is ever widened + * past bounded sources. + */ + @Test + public void testUnboundedSourceKeepsSkippingAnEmptyPartitionChange() throws Throwable { + Collection<String> noSuchTopic = Collections.singleton("topic-that-does-not-exist"); + + try (MockSplitEnumeratorContext<KafkaPartitionSplit> context = + new MockSplitEnumeratorContext<>(NUM_SUBTASKS); + KafkaSourceEnumerator enumerator = + createEnumerator( + context, + -1, + noSuchTopic, + Collections.emptySet(), + Collections.emptySet(), + new Properties(), + OffsetsInitializer.earliest(), + Boundedness.CONTINUOUS_UNBOUNDED, + false)) { + enumerator.start(); + registerReader(context, enumerator, READER0); + runOneTimePartitionDiscovery(context); + + assertThat(enumerator.snapshotState(1L).initialDiscoveryFinished()) + .as("An empty change must not mark the discovery as finished") + .isFalse(); + } + } + + /** + * A restored bounded enumerator must not tell a reader that the input ended before the + * discovery that follows the restore has assigned that reader its splits. A partition created + * while the job was down is only found by that discovery, so signalling from addReader on the + * strength of the restored initialDiscoveryFinished flag would finish the reader first and + * assign to it afterwards. This test pins the ordering that avoids it. + */ + @Test + public void testRestoredBoundedEnumeratorAssignsBeforeItSignals() throws Throwable { + // Only TOPIC1 was assigned before the savepoint; TOPIC2 stands in for the partitions that + // appeared while the job was down and that only the post-restore discovery can find. + Set<KafkaPartitionSplit> restoredAssignedSplits = + KafkaSourceTestEnv.getPartitionsForTopics(Collections.singleton(TOPIC1)).stream() + .map(tp -> new KafkaPartitionSplit(tp, 0L)) + .collect(Collectors.toSet()); + assertThat(restoredAssignedSplits).isNotEmpty(); + + try (RecordingSplitEnumeratorContext context = + new RecordingSplitEnumeratorContext(NUM_SUBTASKS); + KafkaSourceEnumerator enumerator = + createEnumerator( + context, + -1, + PRE_EXISTING_TOPICS, + restoredAssignedSplits, + Collections.emptySet(), + new Properties(), + OffsetsInitializer.earliest(), + Boundedness.BOUNDED, + true)) { + enumerator.start(); + // Registers before the discovery callback, which is the ordering a restore produces. + registerReader(context, enumerator, READER0); + runOneTimePartitionDiscovery(context); + + assertThat(context.events) + .as("The reader must get its new splits before it is told the input ended") + .containsExactly("assignSplits:" + READER0, "signalNoMoreSplits:" + READER0); + } + } + + /** Records the order in which splits are assigned and no-more-splits is signalled. */ + private static class RecordingSplitEnumeratorContext + extends MockSplitEnumeratorContext<KafkaPartitionSplit> { + private final List<String> events = new ArrayList<>(); + + private RecordingSplitEnumeratorContext(int parallelism) { + super(parallelism); + } + + @Override + public void assignSplits(SplitsAssignment<KafkaPartitionSplit> newSplitAssignments) { + newSplitAssignments + .assignment() + .keySet() + .forEach(subtask -> events.add("assignSplits:" + subtask)); + super.assignSplits(newSplitAssignments); + } + + @Override + public void signalNoMoreSplits(int subtask) { + events.add("signalNoMoreSplits:" + subtask); + super.signalNoMoreSplits(subtask); + } + } + @Test public void testCheckSourceIntegrityDefaultValue() throws Exception { final boolean defaultCheckSourceIntegrity = false;
