Repository: apache/flink
File: flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java (Java)

Code before:
        manuallyTriggeredScheduledExecutor.triggerAll();

        CompletableFuture<Void> ackTask =
                CompletableFuture.runAsync(
                        () -> {
                            try {
                                ackCheckpoint(
                                        1L,
                                        coordinator,
                                        jobVertexID,
                                        graph,
                                        handle(),
                                        handle(),
                                        handle());
                            } catch (Exception e) {
                                throw new RuntimeException(e);
                            }
                        });

        assertThat(tracker.getReportStartedFuture().get(20, TimeUnit.SECONDS))
                .as("reportCompletedCheckpoint should be started soon when checkpoint is acked.")
                .isNull();

        for (int i = 0; i < 30; i++) {
            assertThat(checkpointFuture)
                    .as(
                            "Checkpoint future should not complete while reportCompletedCheckpoint is blocked")
                    .isNotDone();
            Thread.sleep(100);
        }

        tracker.getReportBlockingFuture().complete(null);

        CompletedCheckpoint result = checkpointFuture.get(5, TimeUnit.SECONDS);
        assertThat(result)
                .as("Checkpoint future should complete after reportCompletedCheckpoint finishes")
                .isNotNull();

        ackTask.get(5, TimeUnit.SECONDS);
    }

    /**
     * A controllable checkpoint stats tracker for testing purposes. Allows precise control over
     * when reportCompletedCheckpoint() completes, enabling verification of execution order and
     * timing in tests.
     */
    private static class ControllableCheckpointStatsTracker extends DefaultCheckpointStatsTracker {
        private final CompletableFuture<Void> reportStartedFuture;
        private final CompletableFuture<Void> reportBlockingFuture;

        public ControllableCheckpointStatsTracker() {

Review comment (suggestion):
Similarly to above, I am not sure you can confirm whether expected change did not occur because of being blocked vs corresponding thread being inactive. Will be better to wait indefinitely here

Code after:
        CompletableFuture<CompletedCheckpoint> checkpointFuture =
                coordinator.triggerCheckpoint(false);
        manuallyTriggeredScheduledExecutor.triggerAll();

        CompletableFuture<Void> ackTask =
                CompletableFuture.runAsync(
                        () -> {
                            try {
                                ackCheckpoint(
                                        1L,
                                        coordinator,
                                        jobVertexID,
                                        graph,
                                        handle(),
                                        handle(),
                                        handle());
                            } catch (Exception e) {
                                throw new RuntimeException(e);
                            }
                        });

        assertThat(tracker.getReportStartedFuture().get(TIMEOUT_SECONDS, TimeUnit.SECONDS))
                .as("reportCompletedCheckpoint should be started soon when checkpoint is acked.")
                .isNull();

        for (int i = 0; i < 30; i++) {
            assertThat(checkpointFuture)
                    .as(
                            "Checkpoint future should not complete while reportCompletedCheckpoint is blocked")
                    .isNotDone();
            Thread.sleep(100);
        }

        tracker.getReportBlockingFuture().complete(null);

        CompletedCheckpoint result = checkpointFuture.get(TIMEOUT_SECONDS, TimeUnit.SECONDS);
        assertThat(result)
                .as("Checkpoint future should complete after reportCompletedCheckpoint finishes")
                .isNotNull();

        ackTask.get(TIMEOUT_SECONDS, TimeUnit.SECONDS);
    }

    /**
     * A controllable checkpoint stats tracker for testing purposes. Allows precise control over
     * when reportCompletedCheckpoint() completes, enabling verification of execution order and
     * timing in tests.
     */
    private static class ControllableCheckpointStatsTracker extends DefaultCheckpointStatsTracker {
        private final CompletableFuture<Void> reportStartedFuture;
        private final CompletableFuture<Void> reportBlockingFuture;