[FLINK-39133][runtime] Fix DispatcherTest to properly handle async operations for applications

Author: eemario

flink-runtime/src/main/java/org/apache/flink/runtime/dispatcher/Dispatcher.java

@@ -709,29 +709,31 @@ public void notifyApplicationStatusChange(
                                     })
                             .collect(Collectors.toList());
 
-            // wait for all jobs to be stored
-            FutureUtils.combineAll(jobFutures)
-                    .thenAcceptAsync(
-                            combinedJobs -> {
-                                Map<JobID, ExecutionGraphInfo> jobs = new HashMap<>();
-                                for (ExecutionGraphInfo executionGraphInfo : combinedJobs) {
-                                    jobs.put(executionGraphInfo.getJobId(), executionGraphInfo);
-                                    partialExecutionGraphInfoStore.remove(
-                                            executionGraphInfo.getJobId());
-                                }
-
-                                ArchivedApplication archivedApplication =
-                                        new ArchivedApplication(
-                                                application.getApplicationId(),
-                                                application.getName(),
-                                                application.getApplicationStatus(),
-                                                stateTimestamps,
-                                                jobs);
-
-                                applications.remove(applicationId);
-                                writeToArchivedApplicationStore(archivedApplication);
-                                CompletableFuture<?> applicationArchivingFuture =
-                                        historyServerArchivist
+            // wait for all jobs to be stored, then archive the application
+            CompletableFuture<?> applicationArchivingFuture =
+                    FutureUtils.combineAll(jobFutures)
+                            .thenComposeAsync(
+                                    combinedJobs -> {
+                                        Map<JobID, ExecutionGraphInfo> jobs = new HashMap<>();
+                                        for (ExecutionGraphInfo executionGraphInfo : combinedJobs) {
+                                            jobs.put(
+                                                    executionGraphInfo.getJobId(),
+                                                    executionGraphInfo);
+                                            partialExecutionGraphInfoStore.remove(
+                                                    executionGraphInfo.getJobId());
+                                        }
+
+                                        ArchivedApplication archivedApplication =
+                                                new ArchivedApplication(
+                                                        application.getApplicationId(),
+                                                        application.getName(),
+                                                        application.getApplicationStatus(),
+                                                        stateTimestamps,
+                                                        jobs);
+
+                                        applications.remove(applicationId);
+                                        writeToArchivedApplicationStore(archivedApplication);
+                                        return historyServerArchivist
                                                 .archiveApplication(archivedApplication)
                                                 .exceptionally(
                                                         throwable -> {
@@ -741,10 +743,9 @@ public void notifyApplicationStatusChange(
                                                                     throwable);
                                                             return null;
                                                         });
-                                applicationArchivingFutures.put(
-                                        applicationId, applicationArchivingFuture);
-                            },
-                            getMainThreadExecutor());
+                                    },
+                                    getMainThreadExecutor());
+            applicationArchivingFutures.put(applicationId, applicationArchivingFuture);
         }
     }
 

flink-runtime/src/test/java/org/apache/flink/runtime/dispatcher/DispatcherTest.java

@@ -638,6 +638,10 @@ public void testApplicationStatusChange_ThrowsIfDuplicateTerminalStatus() throws
         dispatcher = createTestingDispatcherBuilder().build(rpcService);
         dispatcher.start();
         final ApplicationID applicationId = mockApplicationStatusChange(ApplicationState.FINISHED);
+        // wait for archive to complete
+        dispatcher
+                .getApplicationArchivingFuture(applicationId)
+                .get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
         assertThrows(
                 IllegalStateException.class,
                 () ->