[FLINK-39270] Fix savepoint/last-state upgrade state loss under slow JM

Author: gyfora

flink-kubernetes-operator-api/src/main/java/org/apache/flink/kubernetes/operator/api/status/CommonStatus.java

@@ -133,4 +133,7 @@ public boolean isJobCancellable() {
         }
         return RECONCILING != jobState && !jobState.isGloballyTerminalState();
     }
+
+    @JsonIgnore
+    public abstract boolean isJmAccessible();
 }

flink-kubernetes-operator-api/src/main/java/org/apache/flink/kubernetes/operator/api/status/FlinkDeploymentStatus.java

@@ -107,4 +107,10 @@ public List<Condition> setRunningConditionIfRequired() {
         setConditions(newConditions);
         return newConditions;
     }
+
+    @JsonIgnore
+    @Override
+    public boolean isJmAccessible() {
+        return jobManagerDeploymentStatus.isRestApiAvailable();
+    }
 }

flink-kubernetes-operator-api/src/main/java/org/apache/flink/kubernetes/operator/api/status/FlinkSessionJobStatus.java

@@ -20,6 +20,7 @@
 import org.apache.flink.annotation.Experimental;
 import org.apache.flink.kubernetes.operator.api.spec.FlinkSessionJobSpec;
 
+import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import lombok.AllArgsConstructor;
 import lombok.Data;
@@ -42,4 +43,10 @@ public class FlinkSessionJobStatus extends CommonStatus<FlinkSessionJobSpec> {
     /** Status of the last reconcile operation. */
     private FlinkSessionJobReconciliationStatus reconciliationStatus =
             new FlinkSessionJobReconciliationStatus();
+
+    @Override
+    @JsonIgnore
+    public boolean isJmAccessible() {
+        return true;
+    }
 }

flink-kubernetes-operator/src/main/java/org/apache/flink/kubernetes/operator/reconciler/deployment/AbstractJobReconciler.java

@@ -216,16 +216,13 @@ protected JobUpgrade getJobUpgrade(FlinkResourceContext<CR> ctx, Configuration d
         var resource = ctx.getResource();
         var status = resource.getStatus();
         var upgradeMode = resource.getSpec().getJob().getUpgradeMode();
-        boolean terminal = ReconciliationUtils.isJobInTerminalState(status);
 
         if (upgradeMode == UpgradeMode.STATELESS) {
             LOG.info("Stateless job, ready for upgrade");
-            return JobUpgrade.stateless(terminal);
+            return JobUpgrade.stateless(ReconciliationUtils.isJobInTerminalState(status));
         }
 
-        var flinkService = ctx.getFlinkService();
-        if (ReconciliationUtils.isJobCancelled(status)
-                || (terminal && !flinkService.isHaMetadataAvailable(ctx.getObserveConfig()))) {
+        if (useStatusSavepoint(ctx, status)) {
 
             if (!SnapshotUtils.lastSavepointKnown(status)) {
                 throw new UpgradeFailureException(
@@ -303,6 +300,38 @@ protected JobUpgrade getJobUpgrade(FlinkResourceContext<CR> ctx, Configuration d
         return JobUpgrade.unavailable();
     }
 
+    /**
+     * Core logic to determine whether the savepoint recorded in the status is safe to use for
+     * upgrades: no other savepoint/checkpoint may be in progress and what's there is the latest.
+     */
+    @VisibleForTesting
+    static boolean useStatusSavepoint(FlinkResourceContext<?> ctx, CommonStatus<?> status) {
+
+        // Only terminal job state is fixed/guaranteed
+        if (!ReconciliationUtils.isJobInTerminalState(status)) {
+            return false;
+        }
+
+        // Observed to be fully canceled with state information
+        if (ReconciliationUtils.isJobCancelled(status)) {
+            return true;
+        }
+
+        var service = ctx.getFlinkService();
+        var conf = ctx.getObserveConfig();
+
+        // At this point our job is in terminal state (failed or stopped). Normally the jobmanager
+        // removes any HA state that would affect failover/restarts but this is not fully guaranteed
+        // (the JM may be slow due to CPU or memory pressure).
+        // We accept the recorded savepoint in cases where the HA metadata doesn't exist (it was a
+        // clean job shutdown) or the JM is still accessible which guarantees that the observer
+        // updated the savepoint information.
+        // This still guards against cases where the job + JM fails after restart in a way that the
+        // operator cannot access the new state (in that case HA metadata would be available and JM
+        // inaccessible)
+        return !service.isHaMetadataAvailable(conf) || status.isJmAccessible();
+    }
+
     @VisibleForTesting
     protected JobUpgrade getUpgradeModeBasedOnStateAge(
             FlinkResourceContext<CR> ctx, Configuration deployConfig, boolean cancellable)

flink-kubernetes-operator/src/test/java/org/apache/flink/kubernetes/operator/reconciler/deployment/ApplicationReconcilerUpgradeModeTest.java

@@ -313,6 +313,57 @@ void testUpgradeUsesLatestSnapshot(boolean useLegacyFields) throws Exception {
         assertEquals(savepointPath, flinkService.listJobs().get(0).f0);
     }
 
+    @Test
+    public void testTerminalJobWithHaMetaAvailable() throws Exception {
+        // Bootstrap a running deployment with LAST_STATE upgrade mode (HA config is required)
+        var deployment = buildApplicationCluster(FlinkVersion.v1_19, UpgradeMode.LAST_STATE);
+
+        reconciler.reconcile(deployment, context);
+        verifyAndSetRunningJobsToStatus(deployment, flinkService.listJobs());
+
+        // Simulate: job reached a terminal state while JM is still READY and HA metadata is
+        // available (e.g. job manager cannot delete HA metadata promptly, such as under memory or
+        // cpu pressure)
+        flinkService.clear();
+        flinkService.setHaDataAvailable(true);
+        deployment.getStatus().setJobManagerDeploymentStatus(JobManagerDeploymentStatus.READY);
+        deployment
+                .getStatus()
+                .getJobStatus()
+                .setState(org.apache.flink.api.common.JobStatus.FINISHED);
+
+        // Happy path: a valid savepoint path is recorded in status
+        // Expectation: redeploy from that recorded savepoint
+        final String recordedSavepointPath = "ha_recorded_savepoint";
+        deployment.getStatus().getJobStatus().setUpgradeSavepointPath(recordedSavepointPath);
+        deployment.getSpec().setRestartNonce(100L);
+
+        reconciler.reconcile(deployment, context);
+        reconciler.reconcile(deployment, context);
+
+        assertEquals(1, flinkService.getRunningCount());
+        assertEquals(recordedSavepointPath, flinkService.listJobs().get(0).f0);
+
+        // Error path: only the LAST_STATE dummy marker is in the status, meaning the job was
+        // previously upgraded via HA last-state without recording an actual savepoint path.
+        // When the job is now terminal with JM accessible, there is nothing to restore from.
+        flinkService.clear();
+        deployment.getSpec().setRestartNonce(200L);
+        deployment.getStatus().setJobManagerDeploymentStatus(JobManagerDeploymentStatus.READY);
+        deployment
+                .getStatus()
+                .getJobStatus()
+                .setState(org.apache.flink.api.common.JobStatus.FINISHED);
+        deployment
+                .getStatus()
+                .getJobStatus()
+                .setUpgradeSavepointPath(ApplicationReconciler.LAST_STATE_DUMMY_SP_PATH);
+
+        assertThatThrownBy(() -> reconciler.reconcile(deployment, context))
+                .isInstanceOf(UpgradeFailureException.class)
+                .hasMessageContaining("Job is in terminal state but last checkpoint is unknown");
+    }
+
     private FlinkDeployment cloneDeploymentWithUpgradeMode(
             FlinkDeployment deployment, UpgradeMode upgradeMode) {
         FlinkDeployment result = ReconciliationUtils.clone(deployment);