[Spark] Resolve the structured streaming pipeline state from the execution, cancel() returns the observed state

.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "modification": 5
+  "modification": 6
 }

.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json

@@ -10,6 +10,7 @@
   "https://github.com/apache/beam/pull/35159": "moving WindowedValue and making an interface",
   "https://github.com/apache/beam/pull/39793": "noting that PR #39793 should run this test",
   "https://github.com/apache/beam/pull/40103": "noting that PR #40103 should run this test",
+  "https://github.com/apache/beam/issues/40120": "cancel() returns the observed state",
   "https://github.com/apache/beam/issues/40427": "Spark ValidatesRunner tests run serially",
   "https://github.com/apache/beam/pull/40281": "noting that PR #40281 should run this test",
   "https://github.com/apache/beam/issues/36841": "Spark 4 streaming Impulse and PAssert"

runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java

@@ -117,8 +117,8 @@ public void cancelStopsTheQueryAndReportsCancelled() throws Exception {
 
     awaitActiveQueries(count -> count > 0, "no streaming query started");
 
-    PipelineResult.State cancelledState = result.cancel();
-    assertEquals(PipelineResult.State.CANCELLED, cancelledState);
+    result.cancel();
+    assertEquals(PipelineResult.State.CANCELLED, StreamingTestUtils.waitUntilFinish(result));
     assertEquals(PipelineResult.State.CANCELLED, result.getState());
 
     awaitActiveQueries(count -> count == 0, "the streaming query did not stop after cancel");
@@ -148,8 +148,8 @@ public void timeoutKeepsRunningState() throws Exception {
       assertEquals(PipelineResult.State.RUNNING, stateAfterTimeout);
       assertEquals(PipelineResult.State.RUNNING, result.getState());
 
-      PipelineResult.State cancelledState = result.cancel();
-      assertEquals(PipelineResult.State.CANCELLED, cancelledState);
+      result.cancel();
+      assertEquals(PipelineResult.State.CANCELLED, StreamingTestUtils.waitUntilFinish(result));
       assertEquals(PipelineResult.State.CANCELLED, result.getState());
     } finally {
       if (result.getState() == PipelineResult.State.RUNNING) {

runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResult.java

@@ -21,6 +21,7 @@
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects.firstNonNull;
 
 import java.io.IOException;
+import java.util.concurrent.CancellationException;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
@@ -33,6 +34,7 @@
 import org.apache.beam.sdk.PipelineResult;
 import org.apache.beam.sdk.metrics.MetricResults;
 import org.apache.beam.sdk.util.UserCodeException;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Futures;
 import org.apache.spark.SparkException;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Duration;
@@ -92,6 +94,18 @@ private static RuntimeException unwrapCause(Throwable exception) {
 
   @Override
   public PipelineResult.State getState() {
+    if (!state.isTerminal() && pipelineExecution.isDone()) {
+      if (cancelRequested.get()) {
+        state = State.CANCELLED;
+      } else {
+        try {
+          Futures.getDone(pipelineExecution);
+          state = State.DONE;
+        } catch (ExecutionException | CancellationException e) {
+          state = State.FAILED;
+        }
+      }
+    }
     return state;
   }
 
@@ -108,17 +122,13 @@ public PipelineResult.State waitUntilFinish() {
   public State waitUntilFinish(final Duration duration) {
     try {
       pipelineExecution.get(duration.getMillis(), TimeUnit.MILLISECONDS);
-      state = cancelRequested.get() ? State.CANCELLED : State.DONE;
     } catch (final TimeoutException e) {
       // ignore.
     } catch (final ExecutionException e) {
-      if (cancelRequested.get()) {
-        LOG.warn("Pipeline execution failed after cancel", e.getCause());
-        state = State.CANCELLED;
-        return state;
+      if (getState() != State.CANCELLED) {
+        throw unwrapCause(firstNonNull(e.getCause(), e));
       }
-      state = State.FAILED;
-      throw unwrapCause(firstNonNull(e.getCause(), e));
+      LOG.warn("Pipeline execution failed after cancel", e.getCause());
     } catch (final InterruptedException e) {
       Thread.currentThread().interrupt();
       state = State.FAILED;
@@ -128,25 +138,27 @@ public State waitUntilFinish(final Duration duration) {
       throw unwrapCause(e);
     }
 
-    return state;
+    return getState();
   }
 
   @Override
   public MetricResults metrics() {
     return asAttemptedOnlyMetricResults(metrics.value());
   }
 
-  /** Requests cancellation of the pipeline and returns immediately. */
+  /**
+   * Requests cancellation and returns the current state. The state becomes CANCELLED once the
+   * execution ends.
+   */
   @Override
   public PipelineResult.State cancel() throws IOException {
-    if (!state.isTerminal() && cancelRequested.compareAndSet(false, true)) {
+    if (!pipelineExecution.isDone() && cancelRequested.compareAndSet(false, true)) {
       EvaluationContext ctx = evaluationContext.get();
       if (ctx != null) {
         ctx.stop();
       }
       cancelSparkJobs.run();
-      state = State.CANCELLED;
     }
-    return state;
+    return getState();
   }
 }

runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResultTest.java

@@ -20,13 +20,16 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.is;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
 
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Future;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator;
+import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.PipelineResult.State;
+import org.apache.beam.sdk.util.UserCodeException;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -49,18 +52,39 @@ private SparkStructuredStreamingPipelineResult result(Future<?> execution) {
   @Test
   public void testCancelRunsJobCancelHookOnce() throws Exception {
     SparkStructuredStreamingPipelineResult result = result(new CompletableFuture<>());
-    assertThat(result.cancel(), is(State.CANCELLED));
-    assertThat(result.cancel(), is(State.CANCELLED));
+    assertThat(result.cancel(), is(State.RUNNING));
+    assertThat(result.cancel(), is(State.RUNNING));
     assertThat(cancelSparkJobsCalls.get(), is(1));
   }
 
   @Test
   public void testCancelIsAsynchronous() throws Exception {
     CompletableFuture<Void> execution = new CompletableFuture<>();
     SparkStructuredStreamingPipelineResult result = result(execution);
-    assertThat(result.cancel(), is(State.CANCELLED));
+    assertThat(result.cancel(), is(State.RUNNING));
     assertFalse(execution.isDone());
     execution.completeExceptionally(new IllegalStateException("job cancelled"));
+    assertThat(result.getState(), is(State.CANCELLED));
     assertThat(result.waitUntilFinish(), is(State.CANCELLED));
   }
+
+  @Test
+  public void testCancelAfterNormalCompletionReturnsDone() throws Exception {
+    CompletableFuture<Void> execution = CompletableFuture.completedFuture(null);
+    SparkStructuredStreamingPipelineResult result = result(execution);
+    assertThat(result.cancel(), is(State.DONE));
+    assertThat(cancelSparkJobsCalls.get(), is(0));
+  }
+
+  @Test
+  public void testCancelAfterExceptionalCompletionReturnsFailed() throws Exception {
+    CompletableFuture<Void> execution = new CompletableFuture<>();
+    execution.completeExceptionally(
+        UserCodeException.wrap(new IllegalStateException("pipeline failed")));
+    SparkStructuredStreamingPipelineResult result = result(execution);
+    assertThat(result.cancel(), is(State.FAILED));
+    assertThat(cancelSparkJobsCalls.get(), is(0));
+    assertThat(result.getState(), is(State.FAILED));
+    assertThrows(Pipeline.PipelineExecutionException.class, result::waitUntilFinish);
+  }
 }

runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/StructuredStreamingPipelineStateTest.java

@@ -40,9 +40,7 @@
 import org.apache.spark.TaskContext;
 import org.apache.spark.sql.SparkSession;
 import org.joda.time.Duration;
-import org.junit.Rule;
 import org.junit.Test;
-import org.junit.rules.TestName;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
@@ -60,8 +58,6 @@ private static class MyCustomException extends RuntimeException {
   private final transient SparkStructuredStreamingPipelineOptions options =
       PipelineOptionsFactory.create().as(SparkStructuredStreamingPipelineOptions.class);
 
-  @Rule public transient TestName testName = new TestName();
-
   private static final String FAILED_THE_BATCH_INTENTIONALLY = "Failed the batch intentionally";
 
   private static final long DEADLINE_SECONDS = 60;
@@ -108,16 +104,6 @@ private SparkStructuredStreamingPipelineOptions getBatchOptions() {
     return options;
   }
 
-  private Pipeline getPipeline(final SparkStructuredStreamingPipelineOptions options) {
-
-    final Pipeline pipeline = Pipeline.create(options);
-    final String name = testName.getMethodName();
-
-    pipeline.apply(Create.of("one", "two")).setCoder(StringUtf8Coder.of()).apply(printParDo(name));
-
-    return pipeline;
-  }
-
   private void testFailedPipeline(final SparkStructuredStreamingPipelineOptions options)
       throws Exception {
 
@@ -152,13 +138,9 @@ public String apply(final String input) {
     fail("An injected failure did not affect the pipeline as expected.");
   }
 
-  private void testTimeoutPipeline(final SparkStructuredStreamingPipelineOptions options)
-      throws Exception {
+  private void testTimeoutPipeline() throws Exception {
 
-    final Pipeline pipeline = getPipeline(options);
-
-    final SparkStructuredStreamingPipelineResult result =
-        (SparkStructuredStreamingPipelineResult) pipeline.run();
+    final SparkStructuredStreamingPipelineResult result = runBlockingPipeline();
 
     result.waitUntilFinish(Duration.millis(1));
 
@@ -168,27 +150,19 @@ private void testTimeoutPipeline(final SparkStructuredStreamingPipelineOptions o
     assertThat(result.waitUntilFinish(), is(PipelineResult.State.CANCELLED));
   }
 
-  private void testCanceledPipeline(final SparkStructuredStreamingPipelineOptions options)
-      throws Exception {
+  private void testCanceledPipeline() throws Exception {
 
-    final Pipeline pipeline = getPipeline(options);
-
-    final SparkStructuredStreamingPipelineResult result =
-        (SparkStructuredStreamingPipelineResult) pipeline.run();
+    final SparkStructuredStreamingPipelineResult result = runBlockingPipeline();
 
     result.cancel();
 
-    assertThat(result.getState(), is(PipelineResult.State.CANCELLED));
     assertThat(result.waitUntilFinish(), is(PipelineResult.State.CANCELLED));
+    assertThat(result.getState(), is(PipelineResult.State.CANCELLED));
   }
 
-  private void testRunningPipeline(final SparkStructuredStreamingPipelineOptions options)
-      throws Exception {
+  private void testRunningPipeline() throws Exception {
 
-    final Pipeline pipeline = getPipeline(options);
-
-    final SparkStructuredStreamingPipelineResult result =
-        (SparkStructuredStreamingPipelineResult) pipeline.run();
+    final SparkStructuredStreamingPipelineResult result = runBlockingPipeline();
 
     assertThat(result.getState(), is(PipelineResult.State.RUNNING));
 
@@ -198,12 +172,12 @@ private void testRunningPipeline(final SparkStructuredStreamingPipelineOptions o
 
   @Test
   public void testBatchPipelineRunningState() throws Exception {
-    testRunningPipeline(getBatchOptions());
+    testRunningPipeline();
   }
 
   @Test
   public void testBatchPipelineCanceledState() throws Exception {
-    testCanceledPipeline(getBatchOptions());
+    testCanceledPipeline();
   }
 
   @Test
@@ -213,13 +187,13 @@ public void testBatchPipelineFailedState() throws Exception {
 
   @Test
   public void testBatchPipelineTimeoutState() throws Exception {
-    testTimeoutPipeline(getBatchOptions());
+    testTimeoutPipeline();
   }
 
   @Test
   public void testBatchCancelStopsRunningJob() throws Exception {
     SparkStructuredStreamingPipelineResult result = runBlockingPipeline();
-    assertThat(result.cancel(), is(PipelineResult.State.CANCELLED));
+    result.cancel();
     assertThat(result.waitUntilFinish(), is(PipelineResult.State.CANCELLED));
     assertTrue("owned session not stopped", SparkSession.getDefaultSession().isEmpty());
   }
@@ -229,7 +203,7 @@ public void testCancelKeepsSharedSession() throws Exception {
     SparkSession session = SparkSessionFactory.sessionBuilder("local[1]").getOrCreate();
     try {
       SparkStructuredStreamingPipelineResult result = runBlockingPipeline();
-      assertThat(result.cancel(), is(PipelineResult.State.CANCELLED));
+      result.cancel();
       assertThat(result.waitUntilFinish(), is(PipelineResult.State.CANCELLED));
       assertFalse("shared session stopped", session.sparkContext().isStopped());
     } finally {
@@ -241,7 +215,7 @@ public void testCancelKeepsSharedSession() throws Exception {
   @Test
   public void testCancelFollowedImmediatelyBySecondPipeline() throws Exception {
     SparkStructuredStreamingPipelineResult first = runBlockingPipeline();
-    assertThat(first.cancel(), is(PipelineResult.State.CANCELLED));
+    first.cancel();
     Pipeline secondPipeline = Pipeline.create(getBatchOptions());
     secondPipeline.apply(Create.of("a", "b")).apply(printParDo("second"));
     SparkStructuredStreamingPipelineResult second =