[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": 4
+  "modification": 5
 }

.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 GroupByKey"

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

@@ -152,13 +152,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 {
-
-    final Pipeline pipeline = getPipeline(options);
+  private void testTimeoutPipeline() throws Exception {
 
-    final SparkStructuredStreamingPipelineResult result =
-        (SparkStructuredStreamingPipelineResult) pipeline.run();
+    final SparkStructuredStreamingPipelineResult result = runBlockingPipeline();
 
     result.waitUntilFinish(Duration.millis(1));
 
@@ -168,27 +164,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 {
-
-    final Pipeline pipeline = getPipeline(options);
+  private void testRunningPipeline() throws Exception {
 
-    final SparkStructuredStreamingPipelineResult result =
-        (SparkStructuredStreamingPipelineResult) pipeline.run();
+    final SparkStructuredStreamingPipelineResult result = runBlockingPipeline();
 
     assertThat(result.getState(), is(PipelineResult.State.RUNNING));
 
@@ -198,12 +186,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 +201,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 +217,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 +229,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 =