This is an automated email from the ASF dual-hosted git repository. 1996fanrui pushed a commit to branch release-1.20 in repository https://gitbox.apache.org/repos/asf/flink.git
commit e9b9ac8fecab25008639411017c4676c06cd4d4b Author: Roman Khachatryan <[email protected]> AuthorDate: Wed Feb 18 19:21:31 2026 +0100 [FLINK-38613][tests] Force number of splits in UnalignedCheckpointRescaleWithMixedExchangesITCase Motivation: avoid sub-tasks finishing prematurely and fail the test assumption of all sub-tasks always running (cherry picked from commit 0f589a532b110e5da0012c90812fc7c450bbd55a) --- .../datagen/source/DataGeneratorSource.java | 38 +++++++++- ...dCheckpointRescaleWithMixedExchangesITCase.java | 80 +++++++--------------- 2 files changed, 60 insertions(+), 58 deletions(-) diff --git a/flink-connectors/flink-connector-datagen/src/main/java/org/apache/flink/connector/datagen/source/DataGeneratorSource.java b/flink-connectors/flink-connector-datagen/src/main/java/org/apache/flink/connector/datagen/source/DataGeneratorSource.java index 08ab05bc647..1145ba858f6 100644 --- a/flink-connectors/flink-connector-datagen/src/main/java/org/apache/flink/connector/datagen/source/DataGeneratorSource.java +++ b/flink-connectors/flink-connector-datagen/src/main/java/org/apache/flink/connector/datagen/source/DataGeneratorSource.java @@ -142,16 +142,50 @@ public class DataGeneratorSource<OUT> rateLimiterStrategy, ExecutionConfig.ClosureCleanerLevel.RECURSIVE, true); } + /** + * Instantiates a new {@code DataGeneratorSource}. + * + * @param generatorFunction The {@code GeneratorFunction} function. + * @param numberSource The number source. + * @param rateLimiterStrategy The strategy for rate limiting. + * @param typeInfo The type of the produced data points. + */ + public DataGeneratorSource( + GeneratorFunction<Long, OUT> generatorFunction, + NumberSequenceSource numberSource, + RateLimiterStrategy rateLimiterStrategy, + TypeInformation<OUT> typeInfo) { + this( + new GeneratorSourceReaderFactory<>(generatorFunction, rateLimiterStrategy), + generatorFunction, + typeInfo, + numberSource); + ClosureCleaner.clean( + rateLimiterStrategy, ExecutionConfig.ClosureCleanerLevel.RECURSIVE, true); + } + DataGeneratorSource( SourceReaderFactory<OUT, NumberSequenceSplit> sourceReaderFactory, GeneratorFunction<Long, OUT> generatorFunction, long count, TypeInformation<OUT> typeInfo) { + this( + sourceReaderFactory, + generatorFunction, + typeInfo, + // a noop source (0 elements) is used in Table tests + new NumberSequenceSource(0, count > 0 ? count - 1 : 0)); + } + + DataGeneratorSource( + SourceReaderFactory<OUT, NumberSequenceSplit> sourceReaderFactory, + GeneratorFunction<Long, OUT> generatorFunction, + TypeInformation<OUT> typeInfo, + NumberSequenceSource numberSource) { this.sourceReaderFactory = checkNotNull(sourceReaderFactory); this.generatorFunction = checkNotNull(generatorFunction); this.typeInfo = checkNotNull(typeInfo); - long to = count > 0 ? count - 1 : 0; // a noop source (0 elements) is used in Table tests - this.numberSource = new NumberSequenceSource(0, to); + this.numberSource = numberSource; ClosureCleaner.clean( generatorFunction, ExecutionConfig.ClosureCleanerLevel.RECURSIVE, true); ClosureCleaner.clean( diff --git a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/UnalignedCheckpointRescaleWithMixedExchangesITCase.java b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/UnalignedCheckpointRescaleWithMixedExchangesITCase.java index e443bbcfad4..de9acdf0fb5 100644 --- a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/UnalignedCheckpointRescaleWithMixedExchangesITCase.java +++ b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/UnalignedCheckpointRescaleWithMixedExchangesITCase.java @@ -21,6 +21,7 @@ package org.apache.flink.test.checkpointing; import org.apache.flink.api.common.JobStatus; import org.apache.flink.api.common.eventtime.WatermarkStrategy; import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.api.connector.source.lib.NumberSequenceSource; import org.apache.flink.api.connector.source.util.ratelimit.RateLimiterStrategy; import org.apache.flink.api.java.functions.KeySelector; import org.apache.flink.configuration.CheckpointingOptions; @@ -57,6 +58,7 @@ import java.io.IOException; import java.time.Duration; import java.util.Collection; import java.util.Collections; +import java.util.List; import java.util.Random; import static org.apache.flink.configuration.RestartStrategyOptions.RestartStrategyType.NO_RESTART_STRATEGY; @@ -172,16 +174,10 @@ public class UnalignedCheckpointRescaleWithMixedExchangesITCase extends TestLogg } private static JobClient createMultiOutputDAG(StreamExecutionEnvironment env) throws Exception { - DataGeneratorSource<Long> source = - new DataGeneratorSource<>( - index -> index, - Long.MAX_VALUE, - RateLimiterStrategy.perSecond(5000), - Types.LONG); int sourceParallelism = getRandomParallelism(); DataStream<Long> sourceStream = - env.fromSource(source, WatermarkStrategy.noWatermarks(), "Data Generator") + env.fromSource(createSource(), WatermarkStrategy.noWatermarks(), "Data Generator") .setParallelism(sourceParallelism); sourceStream @@ -209,25 +205,13 @@ public class UnalignedCheckpointRescaleWithMixedExchangesITCase extends TestLogg private static JobClient createMultiInputDAG(StreamExecutionEnvironment env) throws Exception { int source1Parallelism = getRandomParallelism(); - DataGeneratorSource<Long> source1 = - new DataGeneratorSource<>( - index -> index, - Long.MAX_VALUE, - RateLimiterStrategy.perSecond(5000), - Types.LONG); DataStream<Long> sourceStream1 = - env.fromSource(source1, WatermarkStrategy.noWatermarks(), "Source 1") + env.fromSource(createSource(), WatermarkStrategy.noWatermarks(), "Source 1") .setParallelism(source1Parallelism); int source2Parallelism = getRandomParallelism(); - DataGeneratorSource<Long> source2 = - new DataGeneratorSource<>( - index -> index, - Long.MAX_VALUE, - RateLimiterStrategy.perSecond(5000), - Types.LONG); DataStream<Long> sourceStream2 = - env.fromSource(source2, WatermarkStrategy.noWatermarks(), "Source 2") + env.fromSource(createSource(), WatermarkStrategy.noWatermarks(), "Source 2") .setParallelism(source2Parallelism); // Keep the same parallelism to ensure the ForwardPartitioner will be used. @@ -246,16 +230,10 @@ public class UnalignedCheckpointRescaleWithMixedExchangesITCase extends TestLogg private static JobClient createRescalePartitionerDAG(StreamExecutionEnvironment env) throws Exception { - DataGeneratorSource<Long> source = - new DataGeneratorSource<>( - index -> index, - Long.MAX_VALUE, - RateLimiterStrategy.perSecond(5000), - Types.LONG); int sourceParallelism = getRandomParallelism(); DataStream<Long> sourceStream = - env.fromSource(source, WatermarkStrategy.noWatermarks(), "Data Generator") + env.fromSource(createSource(), WatermarkStrategy.noWatermarks(), "Data Generator") .setParallelism(sourceParallelism); sourceStream @@ -285,25 +263,13 @@ public class UnalignedCheckpointRescaleWithMixedExchangesITCase extends TestLogg throws Exception { // Multi-input part int source1Parallelism = getRandomParallelism(); - DataGeneratorSource<Long> source1 = - new DataGeneratorSource<>( - index -> index, - Long.MAX_VALUE, - RateLimiterStrategy.perSecond(5000), - Types.LONG); DataStream<Long> sourceStream1 = - env.fromSource(source1, WatermarkStrategy.noWatermarks(), "Source 1") + env.fromSource(createSource(), WatermarkStrategy.noWatermarks(), "Source 1") .setParallelism(source1Parallelism); int source2Parallelism = getRandomParallelism(); - DataGeneratorSource<Long> source2 = - new DataGeneratorSource<>( - index -> index, - Long.MAX_VALUE, - RateLimiterStrategy.perSecond(5000), - Types.LONG); DataStream<Long> sourceStream2 = - env.fromSource(source2, WatermarkStrategy.noWatermarks(), "Source 2") + env.fromSource(createSource(), WatermarkStrategy.noWatermarks(), "Source 2") .setParallelism(source2Parallelism); // Keep the same parallelism to ensure the ForwardPartitioner will be used. @@ -350,27 +316,15 @@ public class UnalignedCheckpointRescaleWithMixedExchangesITCase extends TestLogg private static JobClient createPartEmptyHashExchangeDAG(StreamExecutionEnvironment env) throws Exception { int source1Parallelism = getRandomParallelism(); - DataGeneratorSource<Long> source1 = - new DataGeneratorSource<>( - index -> index, - Long.MAX_VALUE, - RateLimiterStrategy.perSecond(5000), - Types.LONG); DataStream<Long> sourceStream1 = - env.fromSource(source1, WatermarkStrategy.noWatermarks(), "Source 1") + env.fromSource(createSource(), WatermarkStrategy.noWatermarks(), "Source 1") .setParallelism(source1Parallelism); int source2Parallelism = getRandomParallelism(); - DataGeneratorSource<Long> source2 = - new DataGeneratorSource<>( - index -> index, - Long.MAX_VALUE, - RateLimiterStrategy.perSecond(5000), - Types.LONG); // Filter all records to simulate empty state exchange DataStream<Long> sourceStream2 = - env.fromSource(source2, WatermarkStrategy.noWatermarks(), "Source 2") + env.fromSource(createSource(), WatermarkStrategy.noWatermarks(), "Source 2") .setParallelism(source2Parallelism) .filter(value -> false) .setParallelism(source2Parallelism); @@ -393,6 +347,20 @@ public class UnalignedCheckpointRescaleWithMixedExchangesITCase extends TestLogg return RANDOM.nextInt(MAX_SLOTS) + 1; } + private static DataGeneratorSource<Long> createSource() { + return new DataGeneratorSource<>( + index -> index, + new NumberSequenceSource(0, Long.MAX_VALUE - 1) { + @Override + protected List<NumberSequenceSplit> splitNumberRange( + long from, long to, int numSplitsIgnored) { + return super.splitNumberRange(from, to, MAX_SLOTS); + } + }, + RateLimiterStrategy.perSecond(5000), + Types.LONG); + } + /** A simple CoMapFunction that sleeps for 1ms for each element. */ private static class SleepingCoMap implements CoMapFunction<Long, Long, Long> { @Override
