[FLINK-38272][runtime] Fix unstable BatchJobRecoveryTest

Author: yinqi.yl

flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptivebatch/BatchJobRecoveryTest.java

@@ -264,7 +264,11 @@ void testRecoverFromJMFailover() throws Exception {
         }
 
         waitUntilWriteExecutionVertexFinishedEventPersisted(5);
-        runInMainThread(() -> jobEventStore.stop(false));
+        runInMainThread(
+                () -> {
+                    scheduler.cancel();
+                    jobEventStore.stop(false);
+                });
 
         // register all produced partitions
         registerPartitions(scheduler);
@@ -360,7 +364,11 @@ void testJobVertexUnFinishedAndOperatorCoordinatorNotSupportBatchSnapshot() thro
                 getCurrentAttemptIds(scheduler.getExecutionJobVertex(MIDDLE_ID));
 
         waitUntilWriteExecutionVertexFinishedEventPersisted(6);
-        runInMainThread(() -> jobEventStore.stop(false));
+        runInMainThread(
+                () -> {
+                    scheduler.cancel();
+                    jobEventStore.stop(false);
+                });
 
         // register partitions, the partition of source task 0 is lost, and it will be restarted
         // if middle task 0 need be restarted.
@@ -462,6 +470,7 @@ void testJobVertexFinishedAndOperatorCoordinatorNotSupportBatchSnapshotAndPartit
         waitUntilWriteExecutionVertexFinishedEventPersisted(5);
         runInMainThread(
                 () -> {
+                    scheduler.cancel();
                     jobEventStore.stop(false);
                 });
 
@@ -548,7 +557,11 @@ void testRecoverFromJMFailoverAndPartitionsUnavailable() throws Exception {
                 getCurrentAttemptIds(scheduler.getExecutionJobVertex(SOURCE_ID));
 
         waitUntilWriteExecutionVertexFinishedEventPersisted(5);
-        runInMainThread(() -> jobEventStore.stop(false));
+        runInMainThread(
+                () -> {
+                    scheduler.cancel();
+                    jobEventStore.stop(false);
+                });
 
         int losePartitionsTaskIndex = 0;
 
@@ -619,7 +632,11 @@ void testRecoverDecidedParallelismFromTheSameJobGraphInstance() throws Exception
                 getCurrentAttemptIds(scheduler.getExecutionJobVertex(SINK_ID));
 
         waitUntilWriteExecutionVertexFinishedEventPersisted(12);
-        runInMainThread(() -> jobEventStore.stop(false));
+        runInMainThread(
+                () -> {
+                    scheduler.cancel();
+                    jobEventStore.stop(false);
+                });
 
         // start a new scheduler and try to recover.
         AdaptiveBatchScheduler newScheduler = createScheduler(jobGraph);
@@ -682,6 +699,7 @@ void testPartitionNotFoundTwiceAfterJMFailover() throws Exception {
         waitUntilWriteExecutionVertexFinishedEventPersisted(5);
         runInMainThread(
                 () -> {
+                    scheduler.cancel();
                     jobEventStore.stop(false);
                 });