[FLINK-34227][runtime] Adds assertion to chained operation in JobMast…

…er#stopJobExecution
apache · Feb 4, 2025 · 37f227f · 37f227f
1 parent 5adcab9
commit 37f227f
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/JobMaster.java b/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/JobMaster.java
@@ -1197,6 +1197,8 @@ private CompletableFuture<Void> stopJobExecution(final Exception cause) {
         return FutureUtils.runAfterwards(
                 terminationFuture,
                 () -> {
+                    getMainThreadExecutor().assertRunningInMainThread();
+
                     shuffleMaster.unregisterJob(executionPlan.getJobID());
                     disconnectTaskManagerResourceManagerConnections(cause);
                     stopJobMasterServices();

diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmaster/JobMasterTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmaster/JobMasterTest.java
@@ -1825,10 +1825,18 @@ void testJobMasterDisconnectsOldTaskExecutorIfNewSessionIsSeen() throws Exceptio
 
     @Test
     void testJobMasterOnlyTerminatesAfterTheSchedulerHasClosed() throws Exception {
-        final CompletableFuture<Void> schedulerTerminationFuture = new CompletableFuture<>();
+        final OneShotLatch terminationLatch = new OneShotLatch();
         final TestingSchedulerNG testingSchedulerNG =
                 TestingSchedulerNG.newBuilder()
-                        .setCloseAsyncSupplier(() -> schedulerTerminationFuture)
+                        .setCloseAsyncSupplier(
+                                () -> {
+                                    // just returning a future that's completed in the test code
+                                    // would fail due to the mainThread assertion completing such a
+                                    // future in the main thread is not achievable that easily
+                                    // because the main thread is not exposed in the JobMasterTest
+                                    terminationLatch.awaitQuietly();
+                                    return FutureUtils.completedVoidFuture();
+                                })
                         .build();
 
         try (final JobMaster jobMaster =
@@ -1847,7 +1855,7 @@ void testJobMasterOnlyTerminatesAfterTheSchedulerHasClosed() throws Exception {
                     .as("Expected TimeoutException because the JobMaster should not terminate.")
                     .isInstanceOf(TimeoutException.class);
 
-            schedulerTerminationFuture.complete(null);
+            terminationLatch.trigger();
 
             jobMasterTerminationFuture.get();
         }