address comments

Author: rmetzger

flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerSimpleITCase.java

@@ -20,7 +20,9 @@
 
 import org.apache.flink.api.common.ExecutionConfig;
 import org.apache.flink.api.common.JobExecutionResult;
+import org.apache.flink.api.common.JobStatus;
 import org.apache.flink.api.common.restartstrategy.RestartStrategies;
+import org.apache.flink.api.common.time.Deadline;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.configuration.JobManagerOptions;
 import org.apache.flink.runtime.execution.Environment;
@@ -33,6 +35,7 @@
 import org.apache.flink.runtime.jobmaster.JobResult;
 import org.apache.flink.runtime.minicluster.MiniCluster;
 import org.apache.flink.runtime.testtasks.NoOpInvokable;
+import org.apache.flink.runtime.testutils.CommonTestUtils;
 import org.apache.flink.runtime.testutils.MiniClusterResource;
 import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
 import org.apache.flink.util.FlinkRuntimeException;
@@ -43,6 +46,7 @@
 
 import java.io.IOException;
 import java.time.Duration;
+import java.time.temporal.ChronoUnit;
 
 import static org.junit.Assert.assertTrue;
 
@@ -104,6 +108,34 @@ private JobGraph createJobGraph() {
         return JobGraphTestUtils.streamingJobGraph(source, sink);
     }
 
+    @Test
+    public void testJobCancellationWhileRestartingSucceeds() throws Exception {
+        final long timeInRestartingState = 10000L;
+
+        final MiniCluster miniCluster = MINI_CLUSTER_RESOURCE.getMiniCluster();
+        final JobVertex alwaysFailingOperator = new JobVertex("Always failing operator");
+        alwaysFailingOperator.setInvokableClass(AlwaysFailingInvokable.class);
+        alwaysFailingOperator.setParallelism(1);
+
+        final JobGraph jobGraph = JobGraphTestUtils.streamingJobGraph(alwaysFailingOperator);
+        ExecutionConfig executionConfig = new ExecutionConfig();
+        // configure a high delay between attempts: We'll stay in RESTARTING for 10 seconds.
+        executionConfig.setRestartStrategy(
+                RestartStrategies.fixedDelayRestart(Integer.MAX_VALUE, timeInRestartingState));
+        jobGraph.setExecutionConfig(executionConfig);
+
+        miniCluster.submitJob(jobGraph).join();
+
+        // wait until we are in RESTARTING state
+        CommonTestUtils.waitUntilCondition(
+                () -> miniCluster.getJobStatus(jobGraph.getJobID()).get() == JobStatus.RESTARTING,
+                Deadline.fromNow(Duration.of(timeInRestartingState, ChronoUnit.MILLIS)),
+                5);
+
+        // now cancel while in RESTARTING state
+        miniCluster.cancelJob(jobGraph.getJobID()).get();
+    }
+
     @Test
     public void testGlobalFailoverIfTaskFails() throws Throwable {
         final MiniCluster miniCluster = MINI_CLUSTER_RESOURCE.getMiniCluster();
@@ -160,4 +192,15 @@ private static void reset() {
             hasFailed = false;
         }
     }
+
+    public static final class AlwaysFailingInvokable extends AbstractInvokable {
+        public AlwaysFailingInvokable(Environment environment) {
+            super(environment);
+        }
+
+        @Override
+        public void invoke() throws Exception {
+            throw new FlinkRuntimeException("Test failure.");
+        }
+    }
 }

flink-tests/src/test/java/org/apache/flink/test/scheduling/AdaptiveSchedulerITCase.java

@@ -28,7 +28,6 @@
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.configuration.JobManagerOptions;
 import org.apache.flink.core.execution.JobClient;
-import org.apache.flink.core.testutils.OneShotLatch;
 import org.apache.flink.runtime.state.FunctionInitializationContext;
 import org.apache.flink.runtime.state.FunctionSnapshotContext;
 import org.apache.flink.runtime.testutils.CommonTestUtils;
@@ -39,7 +38,6 @@
 import org.apache.flink.streaming.api.datastream.DataStreamSource;
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
 import org.apache.flink.streaming.api.functions.sink.DiscardingSink;
-import org.apache.flink.streaming.api.functions.source.ParallelSourceFunction;
 import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
 import org.apache.flink.test.util.MiniClusterWithClientResource;
 import org.apache.flink.util.FlinkException;
@@ -111,6 +109,7 @@ public void cancelRunningJobs() {
     /** Tests that the adaptive scheduler can recover stateful operators. */
     @Test
     public void testGlobalFailoverCanRecoverState() throws Exception {
+
         StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
         env.setParallelism(PARALLELISM);
 
@@ -232,43 +231,6 @@ public void testStopWithSavepointFailOnFirstSavepointSucceedOnSecond() throws Ex
         assertThat(savepoint, containsString(savepointDirectory.getAbsolutePath()));
     }
 
-    @Test
-    public void testCancellationOfJobInRestartLoop() throws Exception {
-        final long timeInRestartingState = 10000L;
-        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
-
-        // configure a high delay between attempts: We'll stay in RESTARTING for 10 seconds.
-        env.setRestartStrategy(
-                RestartStrategies.fixedDelayRestart(Integer.MAX_VALUE, timeInRestartingState));
-        env.addSource(new NotifyOnRunningAndFailingSource()).addSink(new DiscardingSink<>());
-
-        JobClient client = env.executeAsync();
-
-        NotifyOnRunningAndFailingSource.runningLatch.await();
-
-        // wait until we are in RESTARTING state
-        CommonTestUtils.waitUntilCondition(
-                () -> client.getJobStatus().get() == JobStatus.RESTARTING,
-                Deadline.fromNow(Duration.of(timeInRestartingState, ChronoUnit.MILLIS)),
-                5);
-
-        // now cancel while in RESTARTING state
-        client.cancel().get();
-    }
-
-    private static class NotifyOnRunningAndFailingSource implements ParallelSourceFunction<String> {
-        private static final OneShotLatch runningLatch = new OneShotLatch();
-
-        @Override
-        public void run(SourceContext<String> ctx) throws Exception {
-            runningLatch.trigger();
-            throw new RuntimeException();
-        }
-
-        @Override
-        public void cancel() {}
-    }
-
     private boolean isDirectoryEmpty(File directory) {
         File[] files = directory.listFiles();
         if (files.length > 0) {