fix SessionDispatcherLeaderProcess and its test

Author: eemario

flink-runtime/src/main/java/org/apache/flink/runtime/dispatcher/runner/SessionDispatcherLeaderProcess.java

@@ -50,6 +50,8 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
@@ -309,23 +311,45 @@ private Collection<AbstractApplication> recoverApplications(
         log.info("Recover all persisted applications that are not finished, yet.");
         final Collection<ApplicationID> applicationIds = getApplicationIds();
         final Collection<AbstractApplication> recoveredApplications = new ArrayList<>();
-        final Collection<JobInfo> recoveredJobInfos =
-                recoveredJobs.stream()
-                        .map(
-                                executionPlan ->
-                                        new JobInfoImpl(
-                                                executionPlan.getJobID(), executionPlan.getName()))
-                        .collect(Collectors.toList());
-        final Collection<JobInfo> recoveredTerminalJobInfos =
-                dirtyJobResults.stream()
-                        .map(
-                                jobResult ->
-                                        new JobInfoImpl(
-                                                jobResult.getJobId(), jobResult.getJobName()))
-                        .collect(Collectors.toList());
+
+        final Map<ApplicationID, Collection<JobInfo>> recoveredJobInfosByApplication =
+                new HashMap<>();
+        for (ExecutionPlan executionPlan : recoveredJobs) {
+            ApplicationID applicationId =
+                    executionPlan
+                            .getApplicationId()
+                            .orElseThrow(
+                                    () ->
+                                            new IllegalStateException(
+                                                    "Application ID is missing in the recovered execution plan. This suggests the job was submitted through an unsupported or incomplete path."));
+            recoveredJobInfosByApplication
+                    .computeIfAbsent(applicationId, k -> new ArrayList<>())
+                    .add(new JobInfoImpl(executionPlan.getJobID(), executionPlan.getName()));
+        }
+
+        final Map<ApplicationID, Collection<JobInfo>> recoveredTerminalJobInfosByApplication =
+                new HashMap<>();
+        for (JobResult jobResult : dirtyJobResults) {
+            ApplicationID applicationId =
+                    jobResult
+                            .getApplicationId()
+                            .orElseThrow(
+                                    () ->
+                                            new IllegalStateException(
+                                                    "Application ID is missing in the recovered job result. This suggests the job was submitted through an unsupported or incomplete path."));
+            recoveredTerminalJobInfosByApplication
+                    .computeIfAbsent(applicationId, k -> new ArrayList<>())
+                    .add(new JobInfoImpl(jobResult.getJobId(), jobResult.getJobName()));
+        }
 
         for (ApplicationID applicationId : applicationIds) {
             if (!recoveredDirtyJobResults.contains(applicationId)) {
+                Collection<JobInfo> recoveredJobInfos =
+                        recoveredJobInfosByApplication.getOrDefault(
+                                applicationId, Collections.emptyList());
+                Collection<JobInfo> recoveredTerminalJobInfos =
+                        recoveredTerminalJobInfosByApplication.getOrDefault(
+                                applicationId, Collections.emptyList());
                 tryRecoverApplication(applicationId, recoveredJobInfos, recoveredTerminalJobInfos)
                         .ifPresent(recoveredApplications::add);
             } else {

flink-runtime/src/test/java/org/apache/flink/runtime/dispatcher/runner/SessionDispatcherLeaderProcessTest.java

@@ -18,9 +18,11 @@
 
 package org.apache.flink.runtime.dispatcher.runner;
 
+import org.apache.flink.api.common.ApplicationID;
 import org.apache.flink.api.common.JobID;
 import org.apache.flink.core.testutils.FlinkAssertions;
 import org.apache.flink.core.testutils.OneShotLatch;
+import org.apache.flink.runtime.application.SingleJobApplication;
 import org.apache.flink.runtime.blob.BlobServer;
 import org.apache.flink.runtime.client.DuplicateJobSubmissionException;
 import org.apache.flink.runtime.client.JobSubmissionException;
@@ -32,6 +34,9 @@
 import org.apache.flink.runtime.jobmanager.ExecutionPlanStore;
 import org.apache.flink.runtime.jobmaster.JobResult;
 import org.apache.flink.runtime.messages.Acknowledge;
+import org.apache.flink.runtime.messages.FlinkApplicationNotFoundException;
+import org.apache.flink.runtime.testutils.TestingApplicationResultStore;
+import org.apache.flink.runtime.testutils.TestingApplicationStore;
 import org.apache.flink.runtime.testutils.TestingExecutionPlanStore;
 import org.apache.flink.runtime.testutils.TestingJobResultStore;
 import org.apache.flink.runtime.util.TestingFatalErrorHandler;
@@ -96,13 +101,16 @@ class SessionDispatcherLeaderProcessTest {
     @BeforeAll
     static void setupClass() {
         ioExecutor = Executors.newSingleThreadExecutor();
+        JOB_GRAPH.setApplicationId(ApplicationID.fromHexString(JOB_GRAPH.getJobID().toHexString()));
     }
 
     @BeforeEach
     void setup() {
         fatalErrorHandler = new TestingFatalErrorHandler();
         executionPlanStore = TestingExecutionPlanStore.newBuilder().build();
         jobResultStore = TestingJobResultStore.builder().build();
+        applicationStore = TestingApplicationStore.newBuilder().build();
+        applicationResultStore = TestingApplicationResultStore.builder().build();
         dispatcherServiceFactory =
                 createFactoryBasedOnGenericSupplier(
                         () -> TestingDispatcherGatewayService.newBuilder().build());
@@ -183,6 +191,7 @@ void testRecoveryWithMultipleExecutionPlansAndOneMatchingDirtyJobResult() throws
         final JobResult matchingDirtyJobResult =
                 TestingJobResultStore.createSuccessfulJobResult(JOB_GRAPH.getJobID());
         final ExecutionPlan otherExecutionPlan = JobGraphTestUtils.emptyJobGraph();
+        otherExecutionPlan.setApplicationId(JOB_GRAPH.getApplicationId().get());
 
         testJobRecovery(
                 Arrays.asList(otherExecutionPlan, JOB_GRAPH),
@@ -566,11 +575,27 @@ void onAddedExecutionPlan_submitsRecoveredJob() throws Exception {
         final CompletableFuture<ExecutionPlan> submittedJobFuture = new CompletableFuture<>();
         final TestingDispatcherGateway testingDispatcherGateway =
                 TestingDispatcherGateway.newBuilder()
-                        .setSubmitFunction(
-                                submittedJob -> {
-                                    submittedJobFuture.complete(submittedJob);
-                                    return CompletableFuture.completedFuture(Acknowledge.get());
+                        .setRequestApplicationFunction(
+                                applicationId ->
+                                        FutureUtils.completedExceptionally(
+                                                new FlinkApplicationNotFoundException(
+                                                        applicationId)))
+                        .setSubmitApplicationFunction(
+                                application -> {
+                                    if (application instanceof SingleJobApplication) {
+                                        ExecutionPlan submittedJob =
+                                                ((SingleJobApplication) application)
+                                                        .getExecutionPlan();
+                                        submittedJobFuture.complete(submittedJob);
+                                        return CompletableFuture.completedFuture(Acknowledge.get());
+                                    }
+                                    return FutureUtils.completedExceptionally(
+                                            new UnsupportedOperationException());
                                 })
+                        .setSubmitFunction(
+                                ignored ->
+                                        FutureUtils.completedExceptionally(
+                                                new UnsupportedOperationException()))
                         .build();
 
         dispatcherServiceFactory =
@@ -715,11 +740,20 @@ private void runJobRecoveryFailureTest(FlinkException testException) throws Exce
     void onAddedExecutionPlan_failingRecoveredJobSubmission_failsFatally() throws Exception {
         final TestingDispatcherGateway dispatcherGateway =
                 TestingDispatcherGateway.newBuilder()
-                        .setSubmitFunction(
-                                jobGraph ->
+                        .setRequestApplicationFunction(
+                                applicationId ->
+                                        FutureUtils.completedExceptionally(
+                                                new FlinkApplicationNotFoundException(
+                                                        applicationId)))
+                        .setSubmitApplicationFunction(
+                                application ->
                                         FutureUtils.completedExceptionally(
                                                 new JobSubmissionException(
-                                                        jobGraph.getJobID(), "test exception")))
+                                                        JOB_GRAPH.getJobID(), "test exception")))
+                        .setSubmitFunction(
+                                ignored ->
+                                        FutureUtils.completedExceptionally(
+                                                new UnsupportedOperationException()))
                         .build();
 
         runOnAddedExecutionPlanTest(