[FLINK-40559][clients] Honor --claimMode in application mode

flink-clients/src/main/java/org/apache/flink/client/cli/CliFrontendParser.java

@@ -652,36 +652,43 @@ private static void printCustomCliOptions(
     }
 
     public static SavepointRestoreSettings createSavepointRestoreSettings(CommandLine commandLine) {
+        final Boolean allowNonRestoredState =
+                commandLine.hasOption(SAVEPOINT_ALLOW_NON_RESTORED_OPTION.getOpt())
+                        ? Boolean.TRUE
+                        : null;
+        // Null exactly when neither --claimMode nor --restoreMode was passed, so it also answers
+        // "was a claim mode explicitly requested?". Keeping an unset value null means an explicit
+        // "--claimMode NO_CLAIM" stays distinguishable from no override at all.
+        final RecoveryClaimMode recoveryClaimMode = parseRecoveryClaimMode(commandLine);
+
         if (commandLine.hasOption(SAVEPOINT_PATH_OPTION.getOpt())) {
             final String savepointPath = commandLine.getOptionValue(SAVEPOINT_PATH_OPTION.getOpt());
-            final Boolean allowNonRestoredState =
-                    commandLine.hasOption(SAVEPOINT_ALLOW_NON_RESTORED_OPTION.getOpt())
-                            ? Boolean.TRUE
-                            : null;
-            final RecoveryClaimMode recoveryClaimMode;
-            if (commandLine.hasOption(SAVEPOINT_CLAIM_MODE)) {
-                recoveryClaimMode =
-                        ConfigurationUtils.convertValue(
-                                commandLine.getOptionValue(SAVEPOINT_CLAIM_MODE),
-                                RecoveryClaimMode.class);
-            } else if (commandLine.hasOption(SAVEPOINT_RESTORE_MODE)) {
-                recoveryClaimMode =
-                        ConfigurationUtils.convertValue(
-                                commandLine.getOptionValue(SAVEPOINT_RESTORE_MODE),
-                                RecoveryClaimMode.class);
-                System.out.printf(
-                        "The option '%s' is deprecated. Please use '%s' instead.%n",
-                        SAVEPOINT_RESTORE_MODE.getLongOpt(), SAVEPOINT_CLAIM_MODE.getLongOpt());
-            } else {
-                recoveryClaimMode = null;
-            }
             return SavepointRestoreSettings.forPath(
                     savepointPath, allowNonRestoredState, recoveryClaimMode);
+        } else if (recoveryClaimMode != null) {
+            return SavepointRestoreSettings.forRecoveryClaimMode(
+                    allowNonRestoredState, recoveryClaimMode);
         } else {
             return SavepointRestoreSettings.none();
         }
     }
 
+    @Nullable
+    private static RecoveryClaimMode parseRecoveryClaimMode(CommandLine commandLine) {
+        if (commandLine.hasOption(SAVEPOINT_CLAIM_MODE)) {
+            return ConfigurationUtils.convertValue(
+                    commandLine.getOptionValue(SAVEPOINT_CLAIM_MODE), RecoveryClaimMode.class);
+        }
+        if (commandLine.hasOption(SAVEPOINT_RESTORE_MODE)) {
+            System.out.printf(
+                    "The option '%s' is deprecated. Please use '%s' instead.%n",
+                    SAVEPOINT_RESTORE_MODE.getLongOpt(), SAVEPOINT_CLAIM_MODE.getLongOpt());
+            return ConfigurationUtils.convertValue(
+                    commandLine.getOptionValue(SAVEPOINT_RESTORE_MODE), RecoveryClaimMode.class);
+        }
+        return null;
+    }
+
     // --------------------------------------------------------------------------------------------
     //  Line Parsing
     // --------------------------------------------------------------------------------------------

flink-clients/src/test/java/org/apache/flink/client/cli/CliFrontendRunTest.java

@@ -22,19 +22,25 @@
 import org.apache.flink.client.cli.util.DummyClusterClientServiceLoader;
 import org.apache.flink.client.deployment.ClusterClientServiceLoader;
 import org.apache.flink.client.deployment.DefaultClusterClientServiceLoader;
+import org.apache.flink.client.deployment.executors.PipelineExecutorUtils;
 import org.apache.flink.client.program.PackagedProgram;
 import org.apache.flink.client.program.TestingClusterClient;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.configuration.CoreOptions;
+import org.apache.flink.configuration.StateRecoveryOptions;
 import org.apache.flink.core.execution.RecoveryClaimMode;
 import org.apache.flink.runtime.jobgraph.SavepointRestoreSettings;
+import org.apache.flink.streaming.api.environment.CheckpointConfig;
+import org.apache.flink.streaming.api.graph.StreamGraph;
+import org.apache.flink.streaming.api.graph.StreamGraphGenerator;
 
 import org.apache.commons.cli.CommandLine;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.EnumSource;
 import org.junit.jupiter.params.provider.MethodSource;
 
 import java.util.Arrays;
@@ -183,6 +189,73 @@ private void testRestoreMode(String flag, String arg, RecoveryClaimMode expected
         assertThat(savepointSettings.allowNonRestoredState()).isTrue();
     }
 
+    @Test
+    void testClaimRestoreModeParsingWithoutSavepoint() throws Exception {
+        testRestoreModeWithoutSavepoint("--claimMode", "claim", RecoveryClaimMode.CLAIM);
+    }
+
+    @Test
+    void testNoClaimRestoreModeParsingWithoutSavepoint() throws Exception {
+        testRestoreModeWithoutSavepoint("-rm", "no_claim", RecoveryClaimMode.NO_CLAIM);
+    }
+
+    @Test
+    void testNoRestoreModeAndNoSavepointParsing() throws Exception {
+        String[] parameters = {getTestJarPath()};
+        CommandLine commandLine =
+                CliFrontendParser.parse(CliFrontendParser.RUN_OPTIONS, parameters, true);
+
+        assertThat(CliFrontendParser.createSavepointRestoreSettings(commandLine))
+                .isEqualTo(SavepointRestoreSettings.none());
+    }
+
+    private void testRestoreModeWithoutSavepoint(
+            String flag, String arg, RecoveryClaimMode expectedMode) throws Exception {
+        String[] parameters = {flag, arg, getTestJarPath()};
+        CommandLine commandLine =
+                CliFrontendParser.parse(CliFrontendParser.RUN_OPTIONS, parameters, true);
+
+        SavepointRestoreSettings savepointSettings =
+                ExecutionConfigAccessor.fromProgramOptions(
+                                ProgramOptions.create(commandLine), Collections.emptyList())
+                        .getSavepointRestoreSettings();
+
+        assertThat(savepointSettings.restoreSavepoint()).isFalse();
+        assertThat(savepointSettings.getRestorePath()).isNull();
+        assertThat(savepointSettings.getRecoveryClaimMode()).isEqualTo(expectedMode);
+
+        // containsKey rather than get: an explicitly requested mode must be written even when it
+        // equals the default, which is what distinguishes it from no override at all.
+        Configuration configuration = new Configuration();
+        SavepointRestoreSettings.toConfiguration(savepointSettings, configuration);
+        assertThat(configuration.containsKey(StateRecoveryOptions.RESTORE_MODE.key())).isTrue();
+        assertThat(configuration.get(StateRecoveryOptions.RESTORE_MODE)).isEqualTo(expectedMode);
+    }
+
+    @ParameterizedTest
+    @EnumSource(RecoveryClaimMode.class)
+    void testPathlessRecoveryClaimModeThroughApplicationGraph(RecoveryClaimMode mode)
+            throws Exception {
+        Configuration configuration = new Configuration();
+        SavepointRestoreSettings expected =
+                SavepointRestoreSettings.forRecoveryClaimMode(null, mode);
+        SavepointRestoreSettings.toConfiguration(expected, configuration);
+
+        StreamGraph graph =
+                new StreamGraphGenerator(
+                                Collections.emptyList(),
+                                new ExecutionConfig(),
+                                new CheckpointConfig(),
+                                configuration)
+                        .generate();
+
+        assertThat(graph.getSavepointRestoreSettings()).isEqualTo(expected);
+        assertThat(
+                        PipelineExecutorUtils.getStreamGraph(graph, configuration)
+                                .getSavepointRestoreSettings())
+                .isEqualTo(expected);
+    }
+
     @Test
     void testUnrecognizedOption() {
         assertThatThrownBy(

flink-container/src/main/java/org/apache/flink/container/entrypoint/StandaloneApplicationClusterConfigurationParserFactory.java

@@ -81,6 +81,8 @@ public Options getOptions() {
         options.addOption(DYNAMIC_PROPERTY_OPTION);
         options.addOption(CliFrontendParser.SAVEPOINT_PATH_OPTION);
         options.addOption(CliFrontendParser.SAVEPOINT_ALLOW_NON_RESTORED_OPTION);
+        options.addOption(CliFrontendParser.SAVEPOINT_CLAIM_MODE);
+        options.addOption(CliFrontendParser.SAVEPOINT_RESTORE_MODE);
 
         return options;
     }

flink-container/src/test/java/org/apache/flink/container/entrypoint/StandaloneApplicationClusterConfigurationParserFactoryTest.java

@@ -25,6 +25,7 @@
 import org.apache.flink.configuration.JobManagerOptions;
 import org.apache.flink.configuration.PipelineOptionsInternal;
 import org.apache.flink.configuration.RestOptions;
+import org.apache.flink.core.execution.RecoveryClaimMode;
 import org.apache.flink.runtime.entrypoint.FlinkParseException;
 import org.apache.flink.runtime.entrypoint.parser.CommandLineParser;
 import org.apache.flink.runtime.jobgraph.SavepointRestoreSettings;
@@ -33,6 +34,8 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 
 import java.io.File;
 import java.io.IOException;
@@ -218,6 +221,29 @@ void testSavepointRestoreSettingsParsing() throws FlinkParseException {
         assertThat(savepointRestoreSettings.allowNonRestoredState()).isTrue();
     }
 
+    @ParameterizedTest
+    @ValueSource(strings = {"--claimMode", "--restoreMode"})
+    void testClaimModeParsing(String option) throws FlinkParseException {
+        final String restorePath = "foobar";
+        final String[] args = {
+            "-c",
+            confDirPath,
+            "-j",
+            JOB_CLASS_NAME,
+            "-s",
+            restorePath,
+            option,
+            RecoveryClaimMode.CLAIM.name()
+        };
+
+        final SavepointRestoreSettings savepointRestoreSettings =
+                commandLineParser.parse(args).getSavepointRestoreSettings();
+
+        assertThat(savepointRestoreSettings.getRestorePath()).isEqualTo(restorePath);
+        assertThat(savepointRestoreSettings.getRecoveryClaimMode())
+                .isEqualTo(RecoveryClaimMode.CLAIM);
+    }
+
     @Test
     void testSetJobIdManually() throws FlinkParseException {
         final JobID jobId = new JobID();

flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/SavepointRestoreSettings.java

@@ -148,6 +148,13 @@ public String toString() {
                     + ", recoveryClaimMode="
                     + getRecoveryClaimMode()
                     + ')';
+        } else if (recoveryClaimMode != null) {
+            return "SavepointRestoreSettings.forRecoveryClaimMode("
+                    + "allowNonRestoredState="
+                    + allowNonRestoredState()
+                    + ", recoveryClaimMode="
+                    + getRecoveryClaimMode()
+                    + ')';
         } else {
             return "SavepointRestoreSettings.none()";
         }
@@ -189,6 +196,22 @@ public static SavepointRestoreSettings forPath(
                 savepointPath, allowNonRestoredState, recoveryClaimMode);
     }
 
+    /**
+     * Creates restore settings that carry an explicitly requested {@link RecoveryClaimMode} without
+     * a savepoint restore path, as used when {@code --claimMode} is passed but {@code
+     * --fromSavepoint} is not. Such a claim mode would otherwise have to be represented as {@link
+     * #none()}, which discards it.
+     *
+     * @param allowNonRestoredState whether to allow non-restored state, or {@code null} if not
+     *     explicitly set by the user.
+     * @param recoveryClaimMode how to restore from the checkpoint.
+     */
+    public static SavepointRestoreSettings forRecoveryClaimMode(
+            @Nullable Boolean allowNonRestoredState, @Nonnull RecoveryClaimMode recoveryClaimMode) {
+        checkNotNull(recoveryClaimMode, "Recovery claim mode.");
+        return new SavepointRestoreSettings(null, allowNonRestoredState, recoveryClaimMode);
+    }
+
     // -------------------------- Parsing to and from a configuration object
     // ------------------------------------
 
@@ -212,15 +235,18 @@ public static void toConfiguration(
 
     public static SavepointRestoreSettings fromConfiguration(final ReadableConfig configuration) {
         final String savepointPath = configuration.get(StateRecoveryOptions.SAVEPOINT_PATH);
-        if (savepointPath == null) {
+        final RecoveryClaimMode recoveryClaimMode =
+                configuration.getOptional(StateRecoveryOptions.RESTORE_MODE).orElse(null);
+        if (savepointPath == null && recoveryClaimMode == null) {
             return SavepointRestoreSettings.none();
         }
         final Boolean allowNonRestored =
                 configuration
                         .getOptional(StateRecoveryOptions.SAVEPOINT_IGNORE_UNCLAIMED_STATE)
                         .orElse(null);
-        final RecoveryClaimMode recoveryClaimMode =
-                configuration.getOptional(StateRecoveryOptions.RESTORE_MODE).orElse(null);
-        return SavepointRestoreSettings.forPath(savepointPath, allowNonRestored, recoveryClaimMode);
+        return savepointPath == null
+                ? SavepointRestoreSettings.forRecoveryClaimMode(allowNonRestored, recoveryClaimMode)
+                : SavepointRestoreSettings.forPath(
+                        savepointPath, allowNonRestored, recoveryClaimMode);
     }
 }