[FLINK-35552][runtime] Rework how CheckpointStatsTracker is constructed.

The checkpoint tracker doesn't live in the DefaultExecutionGraphFactory anymore but is moved into the AdaptiveScheduler. This will allow the scheduler to react to checkpoint-related events.

Author: XComp

flink-core/src/main/java/org/apache/flink/util/function/CachingSupplier.java

@@ -559,12 +559,6 @@ public CheckpointCoordinator getCheckpointCoordinator() {
         return checkpointCoordinator;
     }
 
-    @Nullable
-    @Override
-    public CheckpointStatsTracker getCheckpointStatsTracker() {
-        return checkpointStatsTracker;
-    }
-
     @Override
     public KvStateLocationRegistry getKvStateLocationRegistry() {
         return kvStateLocationRegistry;

flink-core/src/test/java/org/apache/flink/util/function/CachingSupplierTest.java

@@ -61,7 +61,6 @@
 import java.util.List;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledExecutorService;
-import java.util.function.Supplier;
 
 import static org.apache.flink.configuration.StateChangelogOptions.STATE_CHANGE_LOG_STORAGE;
 import static org.apache.flink.util.Preconditions.checkNotNull;
@@ -92,7 +91,7 @@ public static DefaultExecutionGraph buildGraph(
             long initializationTimestamp,
             VertexAttemptNumberStore vertexAttemptNumberStore,
             VertexParallelismStore vertexParallelismStore,
-            Supplier<CheckpointStatsTracker> checkpointStatsTrackerFactory,
+            CheckpointStatsTracker checkpointStatsTracker,
             boolean isDynamicGraph,
             ExecutionJobVertex.Factory executionJobVertexFactory,
             MarkPartitionFinishedStrategy markPartitionFinishedStrategy,
@@ -342,7 +341,7 @@ public static DefaultExecutionGraph buildGraph(
                     completedCheckpointStore,
                     rootBackend,
                     rootStorage,
-                    checkpointStatsTrackerFactory.get(),
+                    checkpointStatsTracker,
                     checkpointsCleaner,
                     jobManagerConfig.get(STATE_CHANGE_LOG_STORAGE));
         }

flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/DefaultExecutionGraph.java

@@ -96,9 +96,6 @@ void enableCheckpointing(
     @Nullable
     CheckpointCoordinator getCheckpointCoordinator();
 
-    @Nullable
-    CheckpointStatsTracker getCheckpointStatsTracker();
-
     KvStateLocationRegistry getKvStateLocationRegistry();
 
     void setJsonPlan(String jsonPlan);

flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/DefaultExecutionGraphBuilder.java

@@ -20,14 +20,12 @@
 
 import org.apache.flink.api.common.time.Time;
 import org.apache.flink.configuration.Configuration;
-import org.apache.flink.configuration.WebOptions;
 import org.apache.flink.runtime.blob.BlobWriter;
 import org.apache.flink.runtime.checkpoint.CheckpointCoordinator;
 import org.apache.flink.runtime.checkpoint.CheckpointIDCounter;
 import org.apache.flink.runtime.checkpoint.CheckpointStatsTracker;
 import org.apache.flink.runtime.checkpoint.CheckpointsCleaner;
 import org.apache.flink.runtime.checkpoint.CompletedCheckpointStore;
-import org.apache.flink.runtime.checkpoint.DefaultCheckpointStatsTracker;
 import org.apache.flink.runtime.deployment.TaskDeploymentDescriptorFactory;
 import org.apache.flink.runtime.executiongraph.DefaultExecutionGraphBuilder;
 import org.apache.flink.runtime.executiongraph.ExecutionDeploymentListener;
@@ -43,14 +41,12 @@
 import org.apache.flink.runtime.jobmaster.ExecutionDeploymentTrackerDeploymentListenerAdapter;
 import org.apache.flink.runtime.metrics.groups.JobManagerJobMetricGroup;
 import org.apache.flink.runtime.shuffle.ShuffleMaster;
-import org.apache.flink.util.function.CachingSupplier;
 
 import org.slf4j.Logger;
 
 import java.util.HashSet;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledExecutorService;
-import java.util.function.Supplier;
 
 import static org.apache.flink.util.Preconditions.checkNotNull;
 
@@ -67,7 +63,6 @@ public class DefaultExecutionGraphFactory implements ExecutionGraphFactory {
     private final BlobWriter blobWriter;
     private final ShuffleMaster<?> shuffleMaster;
     private final JobMasterPartitionTracker jobMasterPartitionTracker;
-    private final Supplier<CheckpointStatsTracker> checkpointStatsTrackerFactory;
     private final boolean isDynamicGraph;
     private final ExecutionJobVertex.Factory executionJobVertexFactory;
 
@@ -124,12 +119,6 @@ public DefaultExecutionGraphFactory(
         this.blobWriter = blobWriter;
         this.shuffleMaster = shuffleMaster;
         this.jobMasterPartitionTracker = jobMasterPartitionTracker;
-        this.checkpointStatsTrackerFactory =
-                new CachingSupplier<>(
-                        () ->
-                                new DefaultCheckpointStatsTracker(
-                                        configuration.get(WebOptions.CHECKPOINTS_HISTORY_SIZE),
-                                        jobManagerJobMetricGroup));
         this.isDynamicGraph = isDynamicGraph;
         this.executionJobVertexFactory = checkNotNull(executionJobVertexFactory);
         this.nonFinishedHybridPartitionShouldBeUnknown = nonFinishedHybridPartitionShouldBeUnknown;
@@ -141,6 +130,7 @@ public ExecutionGraph createAndRestoreExecutionGraph(
             CompletedCheckpointStore completedCheckpointStore,
             CheckpointsCleaner checkpointsCleaner,
             CheckpointIDCounter checkpointIdCounter,
+            CheckpointStatsTracker checkpointStatsTracker,
             TaskDeploymentDescriptorFactory.PartitionLocationConstraint partitionLocationConstraint,
             long initializationTimestamp,
             VertexAttemptNumberStore vertexAttemptNumberStore,
@@ -180,7 +170,7 @@ public ExecutionGraph createAndRestoreExecutionGraph(
                         initializationTimestamp,
                         vertexAttemptNumberStore,
                         vertexParallelismStore,
-                        checkpointStatsTrackerFactory,
+                        checkpointStatsTracker,
                         isDynamicGraph,
                         executionJobVertexFactory,
                         markPartitionFinishedStrategy,

flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionGraph.java

@@ -19,6 +19,7 @@
 package org.apache.flink.runtime.scheduler;
 
 import org.apache.flink.runtime.checkpoint.CheckpointIDCounter;
+import org.apache.flink.runtime.checkpoint.CheckpointStatsTracker;
 import org.apache.flink.runtime.checkpoint.CheckpointsCleaner;
 import org.apache.flink.runtime.checkpoint.CompletedCheckpointStore;
 import org.apache.flink.runtime.deployment.TaskDeploymentDescriptorFactory;
@@ -40,6 +41,8 @@ public interface ExecutionGraphFactory {
      * @param completedCheckpointStore completedCheckpointStore to pass to the CheckpointCoordinator
      * @param checkpointsCleaner checkpointsCleaner to pass to the CheckpointCoordinator
      * @param checkpointIdCounter checkpointIdCounter to pass to the CheckpointCoordinator
+     * @param checkpointStatsTracker The {@link CheckpointStatsTracker} that's used for collecting
+     *     the checkpoint-related statistics.
      * @param partitionLocationConstraint partitionLocationConstraint for this job
      * @param initializationTimestamp initializationTimestamp when the ExecutionGraph was created
      * @param vertexAttemptNumberStore vertexAttemptNumberStore keeping information about the vertex
@@ -57,6 +60,7 @@ ExecutionGraph createAndRestoreExecutionGraph(
             CompletedCheckpointStore completedCheckpointStore,
             CheckpointsCleaner checkpointsCleaner,
             CheckpointIDCounter checkpointIdCounter,
+            CheckpointStatsTracker checkpointStatsTracker,
             TaskDeploymentDescriptorFactory.PartitionLocationConstraint partitionLocationConstraint,
             long initializationTimestamp,
             VertexAttemptNumberStore vertexAttemptNumberStore,

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/DefaultExecutionGraphFactory.java

@@ -35,6 +35,7 @@
 import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
 import org.apache.flink.runtime.jobgraph.IntermediateDataSetID;
 import org.apache.flink.runtime.jobgraph.JobVertexID;
+import org.apache.flink.runtime.jobgraph.tasks.CheckpointCoordinatorConfiguration;
 import org.apache.flink.runtime.jobmanager.PartitionProducerDisposedException;
 import org.apache.flink.runtime.jobmaster.SerializedInputSplit;
 import org.apache.flink.runtime.messages.checkpoint.AcknowledgeCheckpoint;
@@ -81,23 +82,19 @@ public void reportCheckpointMetrics(
     public void reportInitializationMetrics(
             ExecutionAttemptID executionAttemptId,
             SubTaskInitializationMetrics initializationMetrics) {
-        if (executionGraph.getCheckpointStatsTracker() == null) {
+        final CheckpointCoordinatorConfiguration checkpointConfig =
+                executionGraph.getCheckpointCoordinatorConfiguration();
+        if (checkpointConfig == null || !checkpointConfig.isCheckpointingEnabled()) {
             // TODO: Consider to support reporting initialization stats without checkpointing
-            log.debug(
-                    "Ignoring reportInitializationMetrics if checkpoint coordinator is not present");
+            log.debug("Ignoring reportInitializationMetrics if checkpointing is not present");
             return;
         }
-        ioExecutor.execute(
-                () -> {
-                    try {
-                        executionGraph
-                                .getCheckpointStatsTracker()
-                                .reportInitializationMetrics(
-                                        executionAttemptId, initializationMetrics);
-                    } catch (Exception t) {
-                        log.warn("Error while reportInitializationMetrics", t);
-                    }
-                });
+
+        processCheckpointCoordinatorMessage(
+                "ReportInitializationMetrics",
+                coordinator ->
+                        coordinator.reportInitializationMetrics(
+                                executionAttemptId, initializationMetrics));
     }
 
     public void acknowledgeCheckpoint(

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/ExecutionGraphFactory.java

@@ -43,9 +43,11 @@
 import org.apache.flink.runtime.checkpoint.CheckpointRecoveryFactory;
 import org.apache.flink.runtime.checkpoint.CheckpointScheduling;
 import org.apache.flink.runtime.checkpoint.CheckpointStatsSnapshot;
+import org.apache.flink.runtime.checkpoint.CheckpointStatsTracker;
 import org.apache.flink.runtime.checkpoint.CheckpointsCleaner;
 import org.apache.flink.runtime.checkpoint.CompletedCheckpoint;
 import org.apache.flink.runtime.checkpoint.CompletedCheckpointStore;
+import org.apache.flink.runtime.checkpoint.DefaultCheckpointStatsTracker;
 import org.apache.flink.runtime.checkpoint.SubTaskInitializationMetrics;
 import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
 import org.apache.flink.runtime.concurrent.ComponentMainThreadExecutor;
@@ -221,11 +223,20 @@ public SchedulerBase(
         this.deploymentStateTimeMetrics =
                 new DeploymentStateTimeMetrics(jobGraph.getJobType(), jobStatusMetricsSettings);
 
+        final CheckpointStatsTracker checkpointStatsTracker =
+                SchedulerUtils.createCheckpointStatsTrackerIfCheckpointingIsEnabled(
+                        jobGraph,
+                        () ->
+                                new DefaultCheckpointStatsTracker(
+                                        jobMasterConfiguration.get(
+                                                WebOptions.CHECKPOINTS_HISTORY_SIZE),
+                                        jobManagerJobMetricGroup));
         this.executionGraph =
                 createAndRestoreExecutionGraph(
                         completedCheckpointStore,
                         checkpointsCleaner,
                         checkpointIdCounter,
+                        checkpointStatsTracker,
                         initializationTimestamp,
                         mainThreadExecutor,
                         jobStatusListener,
@@ -372,6 +383,7 @@ private ExecutionGraph createAndRestoreExecutionGraph(
             CompletedCheckpointStore completedCheckpointStore,
             CheckpointsCleaner checkpointsCleaner,
             CheckpointIDCounter checkpointIdCounter,
+            CheckpointStatsTracker checkpointStatsTracker,
             long initializationTimestamp,
             ComponentMainThreadExecutor mainThreadExecutor,
             JobStatusListener jobStatusListener,
@@ -384,6 +396,7 @@ private ExecutionGraph createAndRestoreExecutionGraph(
                         completedCheckpointStore,
                         checkpointsCleaner,
                         checkpointIdCounter,
+                        checkpointStatsTracker,
                         TaskDeploymentDescriptorFactory.PartitionLocationConstraint.fromJobType(
                                 jobGraph.getJobType()),
                         initializationTimestamp,

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/ExecutionGraphHandler.java

@@ -24,10 +24,12 @@
 import org.apache.flink.core.execution.RestoreMode;
 import org.apache.flink.runtime.checkpoint.CheckpointIDCounter;
 import org.apache.flink.runtime.checkpoint.CheckpointRecoveryFactory;
+import org.apache.flink.runtime.checkpoint.CheckpointStatsTracker;
 import org.apache.flink.runtime.checkpoint.CompletedCheckpointStore;
 import org.apache.flink.runtime.checkpoint.DeactivatedCheckpointCompletedCheckpointStore;
 import org.apache.flink.runtime.checkpoint.DeactivatedCheckpointIDCounter;
 import org.apache.flink.runtime.checkpoint.DefaultCompletedCheckpointStoreUtils;
+import org.apache.flink.runtime.checkpoint.NoOpCheckpointStatsTracker;
 import org.apache.flink.runtime.client.JobExecutionException;
 import org.apache.flink.runtime.executiongraph.DefaultExecutionGraphBuilder;
 import org.apache.flink.runtime.jobgraph.JobGraph;
@@ -36,6 +38,7 @@
 import org.slf4j.Logger;
 
 import java.util.concurrent.Executor;
+import java.util.function.Supplier;
 
 /** Utils class for Flink's scheduler implementations. */
 public final class SchedulerUtils {
@@ -107,6 +110,15 @@ public static CheckpointIDCounter createCheckpointIDCounterIfCheckpointingIsEnab
         }
     }
 
+    public static CheckpointStatsTracker createCheckpointStatsTrackerIfCheckpointingIsEnabled(
+            JobGraph jobGraph, Supplier<CheckpointStatsTracker> checkpointStatsTrackerFactory) {
+        if (DefaultExecutionGraphBuilder.isCheckpointingEnabled(jobGraph)) {
+            return checkpointStatsTrackerFactory.get();
+        } else {
+            return NoOpCheckpointStatsTracker.INSTANCE;
+        }
+    }
+
     private static CheckpointIDCounter createCheckpointIdCounter(
             CheckpointRecoveryFactory recoveryFactory, JobID jobId) throws Exception {
         return recoveryFactory.createCheckpointIDCounter(jobId);

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/SchedulerBase.java

@@ -43,9 +43,11 @@
 import org.apache.flink.runtime.checkpoint.CheckpointRecoveryFactory;
 import org.apache.flink.runtime.checkpoint.CheckpointScheduling;
 import org.apache.flink.runtime.checkpoint.CheckpointStatsSnapshot;
+import org.apache.flink.runtime.checkpoint.CheckpointStatsTracker;
 import org.apache.flink.runtime.checkpoint.CheckpointsCleaner;
 import org.apache.flink.runtime.checkpoint.CompletedCheckpoint;
 import org.apache.flink.runtime.checkpoint.CompletedCheckpointStore;
+import org.apache.flink.runtime.checkpoint.DefaultCheckpointStatsTracker;
 import org.apache.flink.runtime.checkpoint.SubTaskInitializationMetrics;
 import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
 import org.apache.flink.runtime.client.JobExecutionException;
@@ -328,6 +330,7 @@ public Duration getMaximumDelayForTriggeringRescale() {
     private final CheckpointsCleaner checkpointsCleaner;
     private final CompletedCheckpointStore completedCheckpointStore;
     private final CheckpointIDCounter checkpointIdCounter;
+    private final CheckpointStatsTracker checkpointStatsTracker;
 
     private final CompletableFuture<JobStatus> jobTerminationFuture = new CompletableFuture<>();
 
@@ -419,6 +422,13 @@ public AdaptiveScheduler(
         this.checkpointIdCounter =
                 SchedulerUtils.createCheckpointIDCounterIfCheckpointingIsEnabled(
                         jobGraph, checkpointRecoveryFactory);
+        this.checkpointStatsTracker =
+                SchedulerUtils.createCheckpointStatsTrackerIfCheckpointingIsEnabled(
+                        jobGraph,
+                        () ->
+                                new DefaultCheckpointStatsTracker(
+                                        configuration.get(WebOptions.CHECKPOINTS_HISTORY_SIZE),
+                                        jobManagerJobMetricGroup));
 
         this.slotAllocator = slotAllocator;
 
@@ -1293,6 +1303,7 @@ private ExecutionGraph createExecutionGraphAndRestoreState(
                 completedCheckpointStore,
                 checkpointsCleaner,
                 checkpointIdCounter,
+                checkpointStatsTracker,
                 TaskDeploymentDescriptorFactory.PartitionLocationConstraint.MUST_BE_KNOWN,
                 initializationTimestamp,
                 vertexAttemptNumberStore,