This is an automated email from the ASF dual-hosted git repository. fhueske pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/flink.git
commit c6fc527d1e069f9d7cddc14ec6baf66b3ceac830 Author: Fabian Hueske <[email protected]> AuthorDate: Wed Jul 15 17:44:42 2026 +0200 [FLINK-39785][table] Add input-driven savepoint trigger to restore test framework Lets restore tests take the stop-with-savepoint at a point defined by an input-side signal rather than sink output, so operators that emit nothing at the point of interest can be captured. * RestoreTestBase: extract the trigger into an overridable awaitSavepointReady (default unchanged: waits for sinks to reach their before-restore rows) and retry stop-with-savepoint while the job is not yet fully running. * TestValues watermark-push-down NewSource: add a per-table emission barrier (TestValuesTableFactory#awaitSourceEmitted) completed as rows are emitted. Generated-By: Claude Opus 4.8 (1M context) --- .../factories/TestValuesRuntimeFunctions.java | 62 ++++++++++++++++++++++ .../planner/factories/TestValuesTableFactory.java | 11 ++++ .../plan/nodes/exec/testutils/RestoreTestBase.java | 48 +++++++++++++++-- 3 files changed, 116 insertions(+), 5 deletions(-) diff --git a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/factories/TestValuesRuntimeFunctions.java b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/factories/TestValuesRuntimeFunctions.java index 851275d4ef8..650d55740ac 100644 --- a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/factories/TestValuesRuntimeFunctions.java +++ b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/factories/TestValuesRuntimeFunctions.java @@ -124,6 +124,65 @@ public final class TestValuesRuntimeFunctions { private static final Map<String, List<BiConsumer<Integer, List<Row>>>> localRawResultsObservers = new HashMap<>(); + // [table_name, cumulative number of rows emitted by all subtasks of the source] + private static final Map<String, Integer> sourceEmittedCounts = new HashMap<>(); + // [table_name, [pending emission barriers]] + private static final Map<String, List<SourceEmissionBarrier>> sourceEmissionBarriers = + new HashMap<>(); + + /** A pending {@link #awaitSourceEmitted} request, completed once its target is reached. */ + private static final class SourceEmissionBarrier { + private final int target; + private final CompletableFuture<Void> future; + + private SourceEmissionBarrier(int target, CompletableFuture<Void> future) { + this.target = target; + this.future = future; + } + } + + /** + * Records that a source emitted a row and completes any barrier whose target row count has been + * reached. Called by the source runtime after every emitted element. + */ + static void notifySourceEmitted(String tableName) { + final List<CompletableFuture<Void>> toComplete = new ArrayList<>(); + synchronized (LOCK) { + final int count = sourceEmittedCounts.merge(tableName, 1, Integer::sum); + final List<SourceEmissionBarrier> barriers = sourceEmissionBarriers.get(tableName); + if (barriers != null) { + barriers.removeIf( + barrier -> { + if (count >= barrier.target) { + toComplete.add(barrier.future); + return true; + } + return false; + }); + } + } + // Complete outside the lock; the awaiting thread may re-enter this class. + toComplete.forEach(future -> future.complete(null)); + } + + /** + * Returns a future that completes once source {@code tableName} has emitted at least {@code + * targetCount} rows (cumulative across all subtasks). + */ + static CompletableFuture<Void> awaitSourceEmitted(String tableName, int targetCount) { + final CompletableFuture<Void> future = new CompletableFuture<>(); + synchronized (LOCK) { + if (sourceEmittedCounts.getOrDefault(tableName, 0) >= targetCount) { + future.complete(null); + } else { + sourceEmissionBarriers + .computeIfAbsent(tableName, n -> new ArrayList<>()) + .add(new SourceEmissionBarrier(targetCount, future)); + } + } + return future; + } + static List<String> getRawResultsAsStrings(String tableName) { return getRawResults(tableName).stream() .map(TestValuesRuntimeFunctions::rowToString) @@ -205,6 +264,8 @@ public final class TestValuesRuntimeFunctions { globalRetractResult.clear(); watermarkHistory.clear(); localRawResultsObservers.clear(); + sourceEmittedCounts.clear(); + sourceEmissionBarriers.clear(); } } @@ -334,6 +395,7 @@ public final class TestValuesRuntimeFunctions { generator.onEvent(next, Long.MIN_VALUE, output); generator.onPeriodicEmit(output); } + notifySourceEmitted(tableName); // If enabled, throttle emission of values if (sleepAfterElements > 0 diff --git a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/factories/TestValuesTableFactory.java b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/factories/TestValuesTableFactory.java index 4e260a3e75f..7049fc1c8c1 100644 --- a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/factories/TestValuesTableFactory.java +++ b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/factories/TestValuesTableFactory.java @@ -158,6 +158,7 @@ import java.util.Map; import java.util.Objects; import java.util.Optional; import java.util.Set; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BiConsumer; import java.util.function.Function; @@ -309,6 +310,16 @@ public final class TestValuesTableFactory TestValuesRuntimeFunctions.registerLocalRawResultsObserver(tableName, observer); } + /** + * Returns a future that completes once source {@code tableName} has emitted at least {@code + * targetCount} rows (cumulative across all subtasks). Useful for triggering a savepoint at a + * controlled point when the operator under test produces no output yet. For now only wired for + * the watermark-push-down {@code NewSource} runtime used by restore tests. + */ + public static CompletableFuture<Void> awaitSourceEmitted(String tableName, int targetCount) { + return TestValuesRuntimeFunctions.awaitSourceEmitted(tableName, targetCount); + } + public static List<Watermark> getWatermarkOutput(String tableName) { return TestValuesRuntimeFunctions.getWatermarks(tableName); } diff --git a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/nodes/exec/testutils/RestoreTestBase.java b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/nodes/exec/testutils/RestoreTestBase.java index f1cdc2ca0bc..83f910762dc 100644 --- a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/nodes/exec/testutils/RestoreTestBase.java +++ b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/nodes/exec/testutils/RestoreTestBase.java @@ -76,6 +76,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.stream.Stream; @@ -282,6 +283,21 @@ public abstract class RestoreTestBase implements TableTestProgramRunner { }); } + /** + * Awaits the point during {@link #generateTestSetupFiles} at which the stop-with-savepoint + * should be triggered. The default waits until every sink has produced its "before restore" + * expected rows. + * + * <p>Override for operators that produce no output at the point of interest or when the test + * logic requires more customized savepoint triggering (e.g., based on processed input). See + * {@link TestValuesTableFactory#awaitSourceEmitted(String, int)}. + */ + protected void awaitSavepointReady( + final TableTestProgram program, final List<CompletableFuture<?>> futures) + throws Exception { + awaitExpectedResults(program, futures, true); + } + private void awaitExpectedResults( final TableTestProgram program, final List<CompletableFuture<?>> futures, @@ -379,12 +395,9 @@ public abstract class RestoreTestBase implements TableTestProgramRunner { compiledPlan.writeToFile(getPlanPath(program, getLatestMetadata())); final TableResult tableResult = compiledPlan.execute(); - awaitExpectedResults(program, futures, true); + awaitSavepointReady(program, futures); final JobClient jobClient = tableResult.getJobClient().get(); - final String savepoint = - jobClient - .stopWithSavepoint(false, tmpDir.toString(), SavepointFormatType.DEFAULT) - .get(); + final String savepoint = stopWithSavepointWhenRunning(jobClient); CommonTestUtils.waitForJobStatus(jobClient, Collections.singletonList(JobStatus.FINISHED)); final Path savepointPath = Paths.get(new URI(savepoint)); final Path savepointDirPath = @@ -491,6 +504,31 @@ public abstract class RestoreTestBase implements TableTestProgramRunner { } } + /** + * Triggers a stop-with-savepoint, retrying while the job is not yet fully running. A savepoint + * gated on source emission (see {@link #awaitSavepointReady}) can be requested before all + * downstream tasks have reached RUNNING, which aborts with "Not all required tasks are + * currently running". Sink-gated triggers never hit this (sink output implies a running + * pipeline), so they succeed on the first attempt. + */ + private String stopWithSavepointWhenRunning(JobClient jobClient) throws Exception { + final long deadline = System.currentTimeMillis() + RESULT_AWAIT_TIMEOUT_MILLIS; + while (true) { + try { + return jobClient + .stopWithSavepoint(false, tmpDir.toString(), SavepointFormatType.DEFAULT) + .get(); + } catch (ExecutionException e) { + final boolean notAllRunning = + e.getMessage() != null && e.getMessage().contains("Not all required tasks"); + if (!notAllRunning || System.currentTimeMillis() >= deadline) { + throw e; + } + Thread.sleep(200); + } + } + } + private Path getPlanPath(TableTestProgram program, ExecNodeMetadata metadata) { return Paths.get( getTestResourceDirectory(program, metadata) + "/plan/" + program.id + ".json");
