diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json index e0266d62f2e0..f1ba03a243ee 100644 --- a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json +++ b/.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 } diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json index 1b832d7f19cb..14a1420fa8b0 100644 --- a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json +++ b/.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" diff --git a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java b/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java index fd897db515a0..63a9f5ccbda6 100644 --- a/runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java +++ b/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) { diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResult.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResult.java index c483bbcf3cf3..c666e0b9a2bc 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResult.java +++ b/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,7 +138,7 @@ public State waitUntilFinish(final Duration duration) { throw unwrapCause(e); } - return state; + return getState(); } @Override @@ -136,17 +146,19 @@ 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(); } } diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResultTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResultTest.java index 4cd4858f0001..cfcf8a368527 100644 --- a/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResultTest.java +++ b/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,8 +52,8 @@ 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)); } @@ -58,9 +61,30 @@ public void testCancelRunsJobCancelHookOnce() throws Exception { public void testCancelIsAsynchronous() throws Exception { CompletableFuture 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 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 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); + } } diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/StructuredStreamingPipelineStateTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/StructuredStreamingPipelineStateTest.java index b883af5eb006..a77be97d8d99 100644 --- a/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/StructuredStreamingPipelineStateTest.java +++ b/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 =