[FLINK-40200][Checkpoint] Checkpoint sync phase timeout

docs/layouts/shortcodes/generated/checkpointing_configuration.html

@@ -152,6 +152,12 @@
             <td>String</td>
             <td>The checkpoint storage implementation to be used to checkpoint state.<br />The implementation can be specified either via their shortcut  name, or via the class name of a <code class="highlighter-rouge">CheckpointStorageFactory</code>. If a factory is specified it is instantiated via its zero argument constructor and its <code class="highlighter-rouge">CheckpointStorageFactory#createFromConfig(ReadableConfig, ClassLoader)</code>  method is called.<br />Recognized shortcut names are 'jobmanager' and 'filesystem'.<br />'execution.checkpointing.storage' and 'execution.checkpointing.dir' are usually combined to configure the checkpoint location. By default,  the checkpoint meta data and actual program state will be stored in the JobManager's memory directly. When 'execution.checkpointing.storage' is set to 'jobmanager', if 'execution.checkpointing.dir' is configured, the meta data of checkpoints will be persisted to the path specified by 'execution.checkpointing.dir'. Otherwise, the meta data will be stored in the JobManager's memory. When 'execution.checkpointing.storage' is set to 'filesystem', a valid path must be configured to 'execution.checkpointing.dir', and the checkpoint meta data and actual program state will both be persisted to the path.</td>
         </tr>
+        <tr>
+            <td><h5>execution.checkpointing.sync-phase-timeout</h5></td>
+            <td style="word-wrap: break-word;">0 ms</td>
+            <td>Duration</td>
+            <td>A timeout for the blocking part of a checkpoint, after which a job failure is triggered to address thread blockages. Defaults to 0 (disabled).</td>
+        </tr>
         <tr>
             <td><h5>execution.checkpointing.timeout</h5></td>
             <td style="word-wrap: break-word;">10 min</td>

flink-core/src/main/java/org/apache/flink/configuration/CheckpointingOptions.java

@@ -583,6 +583,14 @@ public class CheckpointingOptions {
                                                     + "will timeout and checkpoint barrier will start working as unaligned checkpoint.")
                                     .build());
 
+    public static final ConfigOption<Duration> CHECKPOINTING_SYNC_PHASE_TIMEOUT =
+            ConfigOptions.key("execution.checkpointing.sync-phase-timeout")
+                    .durationType()
+                    .defaultValue(Duration.ofSeconds(0L))
+                    .withDescription(
+                            "A timeout for the blocking part of a checkpoint, after which a job failure is triggered "
+                                    + "to address thread blockages. Defaults to 0 (disabled).");
+
     public static final ConfigOption<Boolean> FORCE_UNALIGNED =
             ConfigOptions.key("execution.checkpointing.unaligned.forced")
                     .booleanType()

flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java

@@ -1879,20 +1879,23 @@ public void run() {
         }
     }
 
-    public static void logTaskThreadStackTrace(
-            Thread thread, String taskName, long timeoutMs, String action) {
+    public static String getTaskThreadStackTrace(Thread thread) {
         StackTraceElement[] stack = thread.getStackTrace();
         StringBuilder stackTraceStr = new StringBuilder();
         for (StackTraceElement e : stack) {
             stackTraceStr.append(e).append('\n');
         }
+        return stackTraceStr.toString();
+    }
 
+    public static void logTaskThreadStackTrace(
+            Thread thread, String taskName, long timeoutMs, String action) {
         LOG.warn(
                 "Task '{}' did not react to cancelling signal - {}; it is stuck for {} seconds in method:\n {}",
                 taskName,
                 action,
                 timeoutMs / 1000,
-                stackTraceStr);
+                getTaskThreadStackTrace(thread));
     }
 
     /** Various operation of notify checkpoint. */

flink-runtime/src/main/java/org/apache/flink/streaming/api/environment/CheckpointConfig.java

@@ -498,6 +498,17 @@ public void setAlignedCheckpointTimeout(Duration alignedCheckpointTimeout) {
                 CheckpointingOptions.ALIGNED_CHECKPOINT_TIMEOUT, alignedCheckpointTimeout);
     }
 
+    @Experimental
+    public Duration getCheckpointSyncPhaseTimeout() {
+        return configuration.get(CheckpointingOptions.CHECKPOINTING_SYNC_PHASE_TIMEOUT);
+    }
+
+    @Experimental
+    public void setCheckpointSyncPhaseTimeout(Duration checkpointSyncPhaseTimeout) {
+        configuration.set(
+                CheckpointingOptions.CHECKPOINTING_SYNC_PHASE_TIMEOUT, checkpointSyncPhaseTimeout);
+    }
+
     /**
      * @return the number of subtasks to share the same channel state file, as configured via {@link
      *     #setMaxSubtasksPerChannelStateFile(int)} or {@link
@@ -643,6 +654,9 @@ public void configure(ReadableConfig configuration) {
         configuration
                 .getOptional(CheckpointingOptions.PAUSE_SOURCES_UNTIL_FIRST_CHECKPOINT)
                 .ifPresent(this::setPauseSourcesUntilFirstCheckpoint);
+        configuration
+                .getOptional(CheckpointingOptions.CHECKPOINTING_SYNC_PHASE_TIMEOUT)
+                .ifPresent(this::setCheckpointSyncPhaseTimeout);
     }
 
     public void setPauseSourcesUntilFirstCheckpoint(boolean pauseCheckpointsIfTasksNotRunning) {

flink-runtime/src/main/java/org/apache/flink/streaming/runtime/tasks/SubtaskCheckpointCoordinatorImpl.java

@@ -18,6 +18,7 @@
 package org.apache.flink.streaming.runtime.tasks;
 
 import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.configuration.CheckpointingOptions;
 import org.apache.flink.runtime.checkpoint.CheckpointException;
 import org.apache.flink.runtime.checkpoint.CheckpointMetaData;
 import org.apache.flink.runtime.checkpoint.CheckpointMetricsBuilder;
@@ -37,13 +38,15 @@
 import org.apache.flink.runtime.state.CheckpointStreamFactory;
 import org.apache.flink.runtime.state.filesystem.FsMergingCheckpointStorageLocation;
 import org.apache.flink.runtime.taskmanager.AsyncExceptionHandler;
+import org.apache.flink.runtime.taskmanager.AsynchronousException;
 import org.apache.flink.runtime.taskmanager.Task;
 import org.apache.flink.streaming.api.operators.OperatorSnapshotFutures;
 import org.apache.flink.streaming.runtime.io.checkpointing.BarrierAlignmentUtil;
 import org.apache.flink.streaming.runtime.io.checkpointing.BarrierAlignmentUtil.Cancellable;
 import org.apache.flink.streaming.runtime.io.checkpointing.BarrierAlignmentUtil.DelayableTimer;
 import org.apache.flink.util.CollectionUtil;
 import org.apache.flink.util.ExceptionUtils;
+import org.apache.flink.util.FatalExitExceptionHandler;
 import org.apache.flink.util.FlinkRuntimeException;
 import org.apache.flink.util.IOUtils;
 import org.apache.flink.util.clock.Clock;
@@ -70,7 +73,10 @@
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
 
@@ -747,6 +753,15 @@ private boolean takeSnapshotSync(
                         checkpointId, checkpointOptions.getTargetLocation());
         storage = applyFileMergingCheckpoint(storage, checkpointOptions);
 
+        final long syncPhaseTimeoutMillis =
+                env.getJobConfiguration()
+                        .get(CheckpointingOptions.CHECKPOINTING_SYNC_PHASE_TIMEOUT)
+                        .toMillis();
+        CountDownLatch syncPhaseCompleted = new CountDownLatch(1);
+        if (syncPhaseTimeoutMillis > 0) {
+            startSyncPhaseTimeoutWatchdog(checkpointId, syncPhaseTimeoutMillis, syncPhaseCompleted);
+        }
+
         try {
             operatorChain.snapshotState(
                     operatorSnapshotsInProgress,
@@ -757,6 +772,7 @@ private boolean takeSnapshotSync(
                     storage);
 
         } finally {
+            syncPhaseCompleted.countDown();
             checkpointStorage.clearCacheFor(checkpointId);
         }
 
@@ -849,4 +865,51 @@ private static void logCheckpointProcessingDelay(CheckpointMetaData checkpointMe
                     delay);
         }
     }
+
+    private Thread startSyncPhaseTimeoutWatchdog(
+            long checkpointId, long syncPhaseTimeoutMillis, CountDownLatch syncPhaseCompleted) {
+        final Thread taskThread = Thread.currentThread();
+        final Thread syncPhaseTimeoutWatchDog =
+                new Thread(
+                        taskThread.getThreadGroup(),
+                        () -> {
+                            try {
+                                syncPhaseCompleted.await(
+                                        syncPhaseTimeoutMillis, TimeUnit.MILLISECONDS);
+                            } catch (InterruptedException e) {
+                                // task shutdown
+                                return;
+                            }
+                            if (syncPhaseCompleted.getCount() == 0) {
+                                return;
+                            }
+                            try {
+                                final String errorMessage =
+                                        String.format(
+                                                "Task %s did not complete the synchronous phase of checkpoint %s within %s ms.",
+                                                taskName, checkpointId, syncPhaseTimeoutMillis);
+                                if (LOG.isWarnEnabled()) {
+                                    LOG.warn(
+                                            errorMessage
+                                                    + "\n"
+                                                    + Task.getTaskThreadStackTrace(taskThread));
+                                }
+                                env.failExternally(
+                                        new AsynchronousException(
+                                                new TimeoutException(errorMessage)));
+                            } catch (Exception e) {
+                                LOG.error(
+                                        "Error handling sync phase timeout for checkpoint {}",
+                                        checkpointId,
+                                        e);
+                            }
+                        },
+                        String.format(
+                                "checkpoint %s sync phase watchdog for %s_%s",
+                                checkpointId, taskName, env.getTaskInfo().getIndexOfThisSubtask()));
+        syncPhaseTimeoutWatchDog.setDaemon(true);
+        syncPhaseTimeoutWatchDog.setUncaughtExceptionHandler(FatalExitExceptionHandler.INSTANCE);
+        syncPhaseTimeoutWatchDog.start();
+        return syncPhaseTimeoutWatchDog;
+    }
 }

flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/SubtaskCheckpointCoordinatorTest.java

@@ -18,9 +18,12 @@
 
 package org.apache.flink.streaming.runtime.tasks;
 
+import org.apache.flink.api.common.ExecutionConfig;
 import org.apache.flink.api.common.functions.MapFunction;
 import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
 import org.apache.flink.api.common.typeutils.base.StringSerializer;
+import org.apache.flink.configuration.CheckpointingOptions;
+import org.apache.flink.configuration.Configuration;
 import org.apache.flink.core.execution.SavepointFormatType;
 import org.apache.flink.core.testutils.OneShotLatch;
 import org.apache.flink.metrics.groups.OperatorMetricGroup;
@@ -73,6 +76,7 @@
 import org.junit.jupiter.api.Test;
 
 import java.io.IOException;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Map;
 import java.util.concurrent.CancellationException;
@@ -82,6 +86,7 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.RunnableFuture;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 
@@ -95,6 +100,7 @@
 /** Tests for {@link SubtaskCheckpointCoordinator}. */
 class SubtaskCheckpointCoordinatorTest {
     private static final CheckpointStorage CHECKPOINT_STORAGE = new JobManagerCheckpointStorage();
+    private static final long DEFAULT_SYNC_PHASE_TIMEOUT = 200L;
 
     @Test
     void testInitCheckpoint() throws IOException, CheckpointException {
@@ -378,6 +384,99 @@ void testNotifyCheckpointAbortedBeforeAsyncPhase() throws Exception {
         }
     }
 
+    private StreamMockEnvironment getEnvWithCheckpointSyncTimeout() {
+        return new StreamMockEnvironment(
+                new Configuration(
+                        new Configuration()
+                                .set(
+                                        CheckpointingOptions.CHECKPOINTING_SYNC_PHASE_TIMEOUT,
+                                        Duration.ofMillis(DEFAULT_SYNC_PHASE_TIMEOUT))),
+                new Configuration(),
+                new ExecutionConfig(),
+                1L,
+                new MockInputSplitProvider(),
+                1,
+                new TestTaskStateManager());
+    }
+
+    @Test
+    void testSyncPhaseCompletedOnTime() throws Exception {
+        // Block for half the timeout so sync phase completes on time.
+        final long operatorBlockingUntil = DEFAULT_SYNC_PHASE_TIMEOUT / 2;
+        StreamMockEnvironment env = getEnvWithCheckpointSyncTimeout();
+        OneShotLatch unblockSnapshotLatch = new OneShotLatch();
+        BlockingCheckpointOperator operator =
+                new BlockingCheckpointOperator(unblockSnapshotLatch, operatorBlockingUntil);
+
+        try (SubtaskCheckpointCoordinatorImpl subtaskCheckpointCoordinator =
+                (SubtaskCheckpointCoordinatorImpl)
+                        new MockSubtaskCheckpointCoordinatorBuilder().setEnvironment(env).build()) {
+            final OperatorChain<String, AbstractStreamOperator<String>> operatorChain =
+                    operatorChain(operator);
+
+            CheckpointOptions checkpointOptions =
+                    new CheckpointOptions(
+                            CheckpointType.CHECKPOINT,
+                            CheckpointStorageLocationReference.getDefault(),
+                            CheckpointOptions.AlignmentType.ALIGNED,
+                            CheckpointOptions.NO_ALIGNED_CHECKPOINT_TIME_OUT);
+            subtaskCheckpointCoordinator.checkpointState(
+                    new CheckpointMetaData(13L, System.currentTimeMillis()),
+                    checkpointOptions,
+                    new CheckpointMetricsBuilder()
+                            .setAlignmentDurationNanos(0L)
+                            .setBytesProcessedDuringAlignment(0L),
+                    operatorChain,
+                    false,
+                    () -> true);
+        }
+    }
+
+    @Test
+    void testSyncPhaseWatchDogFailsStuckTask() throws Exception {
+        final long checkpointId = 13L;
+        // block for twice the timeout so sync phase times out
+        final long operatorBlockingUntil = DEFAULT_SYNC_PHASE_TIMEOUT * 2;
+        AtomicReference<Throwable> errorRef = new AtomicReference<>();
+        StreamMockEnvironment env = getEnvWithCheckpointSyncTimeout();
+        env.setExternalExceptionHandler(errorRef::set);
+        OneShotLatch unblockSnapshotLatch = new OneShotLatch();
+        BlockingCheckpointOperator operator =
+                new BlockingCheckpointOperator(unblockSnapshotLatch, operatorBlockingUntil);
+
+        try (SubtaskCheckpointCoordinatorImpl subtaskCheckpointCoordinator =
+                (SubtaskCheckpointCoordinatorImpl)
+                        new MockSubtaskCheckpointCoordinatorBuilder().setEnvironment(env).build()) {
+            final OperatorChain<String, AbstractStreamOperator<String>> operatorChain =
+                    operatorChain(operator);
+
+            CheckpointOptions checkpointOptions =
+                    new CheckpointOptions(
+                            CheckpointType.CHECKPOINT,
+                            CheckpointStorageLocationReference.getDefault(),
+                            CheckpointOptions.AlignmentType.ALIGNED,
+                            CheckpointOptions.NO_ALIGNED_CHECKPOINT_TIME_OUT);
+            subtaskCheckpointCoordinator.checkpointState(
+                    new CheckpointMetaData(checkpointId, System.currentTimeMillis()),
+                    checkpointOptions,
+                    new CheckpointMetricsBuilder(),
+                    operatorChain,
+                    false,
+                    () -> true);
+            assertThat(errorRef.get()).isNotNull();
+            assertThat(errorRef.get().getCause())
+                    .isInstanceOf(TimeoutException.class)
+                    .hasMessageContaining(
+                            "did not complete the synchronous phase of checkpoint "
+                                    + checkpointId
+                                    + " within "
+                                    + DEFAULT_SYNC_PHASE_TIMEOUT
+                                    + " ms.");
+        } finally {
+            unblockSnapshotLatch.trigger();
+        }
+    }
+
     @Test
     void testBroadcastCancelCheckpointMarkerOnAbortingFromCoordinator() throws Exception {
         OneInputStreamTaskTestHarness<String, String> testHarness =
@@ -873,6 +972,33 @@ public void processLatencyMarker(LatencyMarker latencyMarker) {}
         public void processWatermarkStatus(WatermarkStatus watermarkStatus) throws Exception {}
     }
 
+    private static final class BlockingCheckpointOperator extends CheckpointOperator {
+
+        private final OneShotLatch unblockSnapshotLatch;
+        private final long waitFor;
+
+        BlockingCheckpointOperator(OneShotLatch unblockSnapshotLatch, long waitFor) {
+            super(new OperatorSnapshotFutures());
+            this.unblockSnapshotLatch = unblockSnapshotLatch;
+            this.waitFor = waitFor;
+        }
+
+        @Override
+        public OperatorSnapshotFutures snapshotState(
+                long checkpointId,
+                long timestamp,
+                CheckpointOptions checkpointOptions,
+                CheckpointStreamFactory storageLocation)
+                throws Exception {
+            try {
+                unblockSnapshotLatch.await(waitFor, TimeUnit.MILLISECONDS);
+            } catch (TimeoutException ignored) {
+
+            }
+            return super.snapshotState(checkpointId, timestamp, checkpointOptions, storageLocation);
+        }
+    }
+
     private static SubtaskCheckpointCoordinator coordinator(ChannelStateWriter channelStateWriter)
             throws IOException {
         return new SubtaskCheckpointCoordinatorImpl(

flink-test-utils-parent/flink-test-utils/src/main/java/org/apache/flink/streaming/util/TestStreamEnvironment.java

@@ -165,6 +165,11 @@ private static void randomizeConfiguration(MiniCluster miniCluster, Configuratio
             if (!conf.contains(CheckpointingOptions.FILE_MERGING_ENABLED)) {
                 randomize(conf, CheckpointingOptions.FILE_MERGING_ENABLED, true);
             }
+            randomize(
+                    conf,
+                    CheckpointingOptions.CHECKPOINTING_SYNC_PHASE_TIMEOUT,
+                    CheckpointingOptions.CHECKPOINTING_TIMEOUT.defaultValue(),
+                    Duration.ofMillis(0));
         }
 
         randomize(