diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Spark4.json index e3d6056a5de9..b26833333238 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": 1 + "modification": 2 } diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json index cad8d98b8ea5..373c31ff2341 100644 --- a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json +++ b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming.json @@ -8,5 +8,6 @@ "https://github.com/apache/beam/pull/34080": "noting that PR #34080 should run this test", "https://github.com/apache/beam/pull/34155": "noting that PR #34155 should run this test", "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/39793": "noting that PR #39793 should run this test", + "https://github.com/apache/beam/pull/40103": "noting that PR #40103 should run this test" } diff --git a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java index ff7839bddb20..2ca54a1e9481 100644 --- a/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java +++ b/runners/spark/4/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/StreamingEvaluationContext.java @@ -56,10 +56,9 @@ public class StreamingEvaluationContext extends EvaluationContext { private final SparkStructuredStreamingPipelineOptions options; - // Guards queries and stopped. + // Guards queries and the stopped flag. private final Object lock = new Object(); private final List queries = new ArrayList<>(); - private boolean stopped = false; StreamingEvaluationContext( Collection> leaves, @@ -92,7 +91,7 @@ public void evaluate() { continue; } synchronized (lock) { - if (stopped) { + if (isStopped()) { break; } } @@ -105,7 +104,7 @@ public void evaluate() { boolean alreadyStopped; synchronized (lock) { queries.add(query); - alreadyStopped = stopped; + alreadyStopped = isStopped(); } if (alreadyStopped) { stopQuery(query); @@ -134,10 +133,10 @@ public void evaluate() { public void stop() { List toStop; synchronized (lock) { - if (stopped) { + if (isStopped()) { return; } - stopped = true; + super.stop(); toStop = new ArrayList<>(queries); } for (StreamingQuery query : toStop) { 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 b592b6fb742d..c483bbcf3cf3 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 @@ -25,6 +25,7 @@ import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Supplier; import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator; import org.apache.beam.runners.spark.structuredstreaming.translation.EvaluationContext; @@ -35,25 +36,37 @@ import org.apache.spark.SparkException; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +/** + * Result of a pipeline submitted to the {@link SparkStructuredStreamingRunner}. The pipeline runs + * asynchronously on a dedicated thread. + */ public class SparkStructuredStreamingPipelineResult implements PipelineResult { + private static final Logger LOG = + LoggerFactory.getLogger(SparkStructuredStreamingPipelineResult.class); + private final Future pipelineExecution; // Supplies the context of the translated pipeline, null until translation has completed. private final Supplier evaluationContext; private final MetricsAccumulator metrics; - private final @Nullable Runnable onTerminalState; - private PipelineResult.State state; + private final AtomicBoolean cancelRequested; + private final Runnable cancelSparkJobs; + private volatile PipelineResult.State state; SparkStructuredStreamingPipelineResult( Future pipelineExecution, Supplier evaluationContext, MetricsAccumulator metrics, - final @Nullable Runnable onTerminalState) { + AtomicBoolean cancelRequested, + Runnable cancelSparkJobs) { this.pipelineExecution = pipelineExecution; this.evaluationContext = evaluationContext; this.metrics = metrics; - this.onTerminalState = onTerminalState; + this.cancelRequested = cancelRequested; + this.cancelSparkJobs = cancelSparkJobs; // pipelineExecution is expected to have started executing eagerly. this.state = State.RUNNING; } @@ -77,13 +90,6 @@ private static RuntimeException unwrapCause(Throwable exception) { : new Pipeline.PipelineExecutionException(firstNonNull(next, exception)); } - private State awaitTermination(Duration duration) - throws TimeoutException, ExecutionException, InterruptedException { - pipelineExecution.get(duration.getMillis(), TimeUnit.MILLISECONDS); - // Throws an exception if the job is not finished successfully in the given time. - return PipelineResult.State.DONE; - } - @Override public PipelineResult.State getState() { return state; @@ -94,18 +100,31 @@ public PipelineResult.State waitUntilFinish() { return waitUntilFinish(Duration.millis(Long.MAX_VALUE)); } + /** + * Waits up to {@code duration} for the execution thread. A pipeline that ends after {@link + * #cancel()} is CANCELLED, any other failure is rethrown and the pipeline is FAILED. + */ @Override public State waitUntilFinish(final Duration duration) { try { - State finishState = awaitTermination(duration); - offerNewState(finishState); + pipelineExecution.get(duration.getMillis(), TimeUnit.MILLISECONDS); + state = cancelRequested.get() ? State.CANCELLED : State.DONE; } catch (final TimeoutException e) { // ignore. } catch (final ExecutionException e) { - offerNewState(PipelineResult.State.FAILED); + if (cancelRequested.get()) { + LOG.warn("Pipeline execution failed after cancel", e.getCause()); + state = State.CANCELLED; + return state; + } + state = State.FAILED; throw unwrapCause(firstNonNull(e.getCause(), e)); + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + state = State.FAILED; + throw unwrapCause(e); } catch (final Exception e) { - offerNewState(PipelineResult.State.FAILED); + state = State.FAILED; throw unwrapCause(e); } @@ -117,26 +136,17 @@ public MetricResults metrics() { return asAttemptedOnlyMetricResults(metrics.value()); } + /** Requests cancellation of the pipeline and returns immediately. */ @Override public PipelineResult.State cancel() throws IOException { - EvaluationContext ctx = evaluationContext.get(); - if (ctx != null) { - ctx.stop(); - } - pipelineExecution.cancel(true); - offerNewState(PipelineResult.State.CANCELLED); - return state; - } - - private void offerNewState(State newState) { - State oldState = this.state; - this.state = newState; - if (!oldState.isTerminal() && newState.isTerminal() && onTerminalState != null) { - try { - onTerminalState.run(); - } catch (Exception e) { - throw unwrapCause(e); + if (!state.isTerminal() && cancelRequested.compareAndSet(false, true)) { + EvaluationContext ctx = evaluationContext.get(); + if (ctx != null) { + ctx.stop(); } + cancelSparkJobs.run(); + state = State.CANCELLED; } + return state; } } diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java index f78026847fad..0f5ddffae344 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingRunner.java @@ -17,12 +17,13 @@ */ package org.apache.beam.runners.spark.structuredstreaming; +import java.util.UUID; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.ThreadFactory; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; -import javax.annotation.Nullable; import org.apache.beam.runners.core.metrics.MetricsPusher; import org.apache.beam.runners.core.metrics.NoOpMetricsSink; import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator; @@ -42,6 +43,7 @@ import org.apache.beam.sdk.util.construction.SplittableParDo; import org.apache.beam.sdk.util.construction.graph.ProjectionPushdownOptimizer; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder; +import org.apache.spark.SparkContext; import org.apache.spark.SparkEnv$; import org.apache.spark.metrics.MetricsSystem; import org.apache.spark.sql.SparkSession; @@ -145,27 +147,47 @@ public SparkStructuredStreamingPipelineResult run(final Pipeline pipeline) { PipelineTranslator.detectStreamingMode(pipeline, options); - final SparkSession sparkSession = SparkSessionFactory.getOrCreateSession(options); + final boolean releaseSession = !options.getUseActiveSparkSession(); + final SparkSession sparkSession = SparkSessionFactory.acquire(options); + final SparkContext sc = sparkSession.sparkContext(); final MetricsAccumulator metrics = MetricsAccumulator.getInstance(sparkSession); - // Set once the pipeline is translated, so the result can stop an ongoing (streaming) - // evaluation on cancel. Remains null until translation completes. + // Null until translation completes. final AtomicReference ctxRef = new AtomicReference<>(); + final AtomicBoolean cancelRequested = new AtomicBoolean(false); + + final String jobName = options.getJobName(); + final String jobGroupId = "beam-" + jobName + "-" + UUID.randomUUID(); + final Runnable cancelSparkJobs = + () -> { + try { + sc.cancelJobGroup(jobGroupId); + } catch (IllegalStateException e) { + // Context stopped concurrently. + } + }; final Future submissionFuture = runAsync( () -> { - EvaluationContext ctx = translatePipeline(sparkSession, pipeline); - ctxRef.set(ctx); - ctx.evaluate(); + try { + // Interrupts running tasks on cancel, as Spark's StreamExecution does. + sc.setJobGroup(jobGroupId, "Beam " + jobName, true); + EvaluationContext ctx = translatePipeline(sparkSession, pipeline); + ctxRef.set(ctx); + if (!cancelRequested.get()) { + ctx.evaluate(); + } + } finally { + if (releaseSession) { + SparkSessionFactory.release(sparkSession); + } + } }); final SparkStructuredStreamingPipelineResult result = new SparkStructuredStreamingPipelineResult( - submissionFuture, - ctxRef::get, - metrics, - sparkStopFn(sparkSession, options.getUseActiveSparkSession())); + submissionFuture, ctxRef::get, metrics, cancelRequested, cancelSparkJobs); if (options.getEnableSparkMetricSinks()) { registerMetricsSource(options.getAppName(), metrics); @@ -228,8 +250,4 @@ private static Future runAsync(Runnable task) { execService.shutdown(); return future; } - - private static @Nullable Runnable sparkStopFn(SparkSession session, boolean isProvided) { - return !isProvided ? () -> session.stop() : null; - } } diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/EvaluationContext.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/EvaluationContext.java index 0e677051fb61..fb559ab8d9c1 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/EvaluationContext.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/EvaluationContext.java @@ -52,6 +52,7 @@ public interface NamedDataset { private final Collection> leaves; private final SparkSession session; + private volatile boolean stopped = false; protected EvaluationContext(Collection> leaves, SparkSession session) { this.leaves = leaves; @@ -63,9 +64,13 @@ protected Collection> leaves() { return leaves; } - /** Trigger evaluation of all leaf datasets. */ + /** Trigger evaluation of all leaf datasets. Returns early once {@link #stop()} was called. */ public void evaluate() { for (NamedDataset ds : leaves) { + if (stopped) { + LOG.info("Evaluation stopped, skipping remaining datasets"); + return; + } final Dataset dataset = ds.dataset(); if (dataset == null) { continue; @@ -119,11 +124,16 @@ public static void evaluate(String name, Dataset ds) { } /** - * Stops any ongoing streaming execution triggered by this context. - * - *

This is a no-op for batch pipelines. + * Stops the evaluation after the current leaf dataset. Streaming contexts override this to stop + * their queries. */ - public void stop() {} + public void stop() { + stopped = true; + } + + protected boolean isStopped() { + return stopped; + } public SparkSession getSparkSession() { return session; diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java index 148188bb15a2..4222857f4f30 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SparkSessionFactory.java @@ -28,6 +28,7 @@ import java.util.Collection; import java.util.Collections; import java.util.HashMap; +import java.util.Map; import javax.annotation.Nullable; import org.apache.beam.repackaged.core.org.apache.commons.lang3.ArrayUtils; import org.apache.beam.runners.core.construction.SerializablePipelineOptions; @@ -90,6 +91,7 @@ import org.apache.spark.sql.execution.datasources.v2.DataWritingSparkTaskResult; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import scala.Option; public class SparkSessionFactory { @@ -113,15 +115,49 @@ public class SparkSessionFactory { "/com.esotericsoftware/kryo-shaded", "/com/esotericsoftware/kryo-shaded"); - /** - * Gets active {@link SparkSession} or creates one using {@link - * SparkStructuredStreamingPipelineOptions}. - */ - public static SparkSession getOrCreateSession(SparkStructuredStreamingPipelineOptions options) { + // Builder.getOrCreate adopts an existing session without applying the pipeline's configuration. + // A pipeline must not stop a session it did not create, and the next pipeline needs the + // previous one stopped to get its own configuration, so sessions created here are counted. + private static final Map OWNED_SESSIONS = new HashMap<>(); + + /** Returns the {@link SparkSession} for a pipeline, paired with {@link #release}. */ + public static synchronized SparkSession acquire(SparkStructuredStreamingPipelineOptions options) { if (options.getUseActiveSparkSession()) { return SparkSession.active(); } - return sessionBuilder(options.getSparkMaster(), options).getOrCreate(); + boolean noUsableSession = + !isUsable(SparkSession.getActiveSession()) && !isUsable(SparkSession.getDefaultSession()); + SparkSession session = sessionBuilder(options.getSparkMaster(), options).getOrCreate(); + Integer count = OWNED_SESSIONS.get(session); + if (count != null) { + OWNED_SESSIONS.put(session, count + 1); + LOG.info("Pipeline options will not be applied to the shared SparkSession"); + } else if (noUsableSession) { + OWNED_SESSIONS.put(session, 1); + } + return session; + } + + /** + * Releases a session from {@link #acquire} and stops it when no longer used. The stop runs under + * the lock, a pipeline starting meanwhile creates a new session. + */ + public static synchronized void release(SparkSession session) { + Integer count = OWNED_SESSIONS.get(session); + if (count == null) { + return; + } + if (count > 1) { + OWNED_SESSIONS.put(session, count - 1); + return; + } + OWNED_SESSIONS.remove(session); + LOG.info("Stopping SparkSession created by the runner"); + session.stop(); + } + + private static boolean isUsable(Option session) { + return session.isDefined() && !session.get().sparkContext().isStopped(); } /** Creates Spark session builder with some optimizations for local mode, e.g. in tests. */ 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 new file mode 100644 index 000000000000..4cd4858f0001 --- /dev/null +++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResultTest.java @@ -0,0 +1,66 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.spark.structuredstreaming; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.is; +import static org.junit.Assert.assertFalse; + +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.PipelineResult.State; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for the cancel and wait semantics of {@link SparkStructuredStreamingPipelineResult}. */ +@RunWith(JUnit4.class) +public class SparkStructuredStreamingPipelineResultTest { + + private final AtomicInteger cancelSparkJobsCalls = new AtomicInteger(); + + private SparkStructuredStreamingPipelineResult result(Future execution) { + return new SparkStructuredStreamingPipelineResult( + execution, + () -> null, + new MetricsAccumulator(), + new AtomicBoolean(), + cancelSparkJobsCalls::incrementAndGet); + } + + @Test + public void testCancelRunsJobCancelHookOnce() throws Exception { + SparkStructuredStreamingPipelineResult result = result(new CompletableFuture<>()); + assertThat(result.cancel(), is(State.CANCELLED)); + assertThat(result.cancel(), is(State.CANCELLED)); + assertThat(cancelSparkJobsCalls.get(), is(1)); + } + + @Test + public void testCancelIsAsynchronous() throws Exception { + CompletableFuture execution = new CompletableFuture<>(); + SparkStructuredStreamingPipelineResult result = result(execution); + assertThat(result.cancel(), is(State.CANCELLED)); + assertFalse(execution.isDone()); + execution.completeExceptionally(new IllegalStateException("job cancelled")); + assertThat(result.waitUntilFinish(), is(State.CANCELLED)); + } +} 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 b44df7bf101b..647c4334ff15 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 @@ -20,10 +20,15 @@ import static org.hamcrest.CoreMatchers.instanceOf; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.is; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import java.io.Serializable; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import org.apache.beam.runners.spark.io.CreateStream; +import org.apache.beam.runners.spark.structuredstreaming.translation.SparkSessionFactory; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.coders.StringUtf8Coder; @@ -36,6 +41,8 @@ import org.apache.beam.sdk.transforms.SimpleFunction; import org.apache.beam.sdk.values.PBegin; import org.apache.beam.sdk.values.PCollection; +import org.apache.spark.TaskContext; +import org.apache.spark.sql.SparkSession; import org.joda.time.Duration; import org.junit.Ignore; import org.junit.Rule; @@ -62,6 +69,33 @@ private static class MyCustomException extends RuntimeException { private static final String FAILED_THE_BATCH_INTENTIONALLY = "Failed the batch intentionally"; + private static final long DEADLINE_SECONDS = 60; + + // Shared with the DoFn running in Spark's local executor threads, reset per test. + private static volatile CountDownLatch started = new CountDownLatch(1); + + /** Signals started, then blocks until the task is killed. */ + private static class BlockingDoFn extends DoFn { + @ProcessElement + public void processElement(ProcessContext c) throws InterruptedException { + started.countDown(); + while (!TaskContext.get().isInterrupted()) { + Thread.sleep(50); + } + c.output(c.element()); + } + } + + private SparkStructuredStreamingPipelineResult runBlockingPipeline() throws InterruptedException { + started = new CountDownLatch(1); + Pipeline pipeline = Pipeline.create(getBatchOptions()); + pipeline.apply(Create.of("one", "two")).apply(ParDo.of(new BlockingDoFn())); + SparkStructuredStreamingPipelineResult result = + (SparkStructuredStreamingPipelineResult) pipeline.run(); + assertTrue("DoFn did not start", started.await(DEADLINE_SECONDS, TimeUnit.SECONDS)); + return result; + } + private ParDo.SingleOutput printParDo(final String prefix) { return ParDo.of( new DoFn() { @@ -151,6 +185,7 @@ private void testTimeoutPipeline(final SparkStructuredStreamingPipelineOptions o assertThat(result.getState(), is(PipelineResult.State.RUNNING)); result.cancel(); + assertThat(result.waitUntilFinish(), is(PipelineResult.State.CANCELLED)); } private void testCanceledPipeline(final SparkStructuredStreamingPipelineOptions options) @@ -164,6 +199,7 @@ private void testCanceledPipeline(final SparkStructuredStreamingPipelineOptions result.cancel(); assertThat(result.getState(), is(PipelineResult.State.CANCELLED)); + assertThat(result.waitUntilFinish(), is(PipelineResult.State.CANCELLED)); } private void testRunningPipeline(final SparkStructuredStreamingPipelineOptions options) @@ -177,6 +213,7 @@ private void testRunningPipeline(final SparkStructuredStreamingPipelineOptions o assertThat(result.getState(), is(PipelineResult.State.RUNNING)); result.cancel(); + assertThat(result.waitUntilFinish(), is(PipelineResult.State.CANCELLED)); } @Ignore("TODO: Reactivate with streaming.") @@ -222,4 +259,39 @@ public void testStreamingPipelineTimeoutState() throws Exception { public void testBatchPipelineTimeoutState() throws Exception { testTimeoutPipeline(getBatchOptions()); } + + @Test + public void testBatchCancelStopsRunningJob() throws Exception { + SparkStructuredStreamingPipelineResult result = runBlockingPipeline(); + assertThat(result.cancel(), is(PipelineResult.State.CANCELLED)); + assertThat(result.waitUntilFinish(), is(PipelineResult.State.CANCELLED)); + assertTrue("owned session not stopped", SparkSession.getDefaultSession().isEmpty()); + } + + @Test + public void testCancelKeepsSharedSession() throws Exception { + SparkSession session = SparkSessionFactory.sessionBuilder("local[1]").getOrCreate(); + try { + SparkStructuredStreamingPipelineResult result = runBlockingPipeline(); + assertThat(result.cancel(), is(PipelineResult.State.CANCELLED)); + assertThat(result.waitUntilFinish(), is(PipelineResult.State.CANCELLED)); + assertFalse("shared session stopped", session.sparkContext().isStopped()); + } finally { + session.stop(); + } + } + + /** The second pipeline shares the first session or creates a new one, both must end cleanly. */ + @Test + public void testCancelFollowedImmediatelyBySecondPipeline() throws Exception { + SparkStructuredStreamingPipelineResult first = runBlockingPipeline(); + assertThat(first.cancel(), is(PipelineResult.State.CANCELLED)); + Pipeline secondPipeline = Pipeline.create(getBatchOptions()); + secondPipeline.apply(Create.of("a", "b")).apply(printParDo("second")); + SparkStructuredStreamingPipelineResult second = + (SparkStructuredStreamingPipelineResult) secondPipeline.run(); + assertThat(first.waitUntilFinish(), is(PipelineResult.State.CANCELLED)); + assertThat(second.waitUntilFinish(), is(PipelineResult.State.DONE)); + assertTrue("session not stopped", SparkSession.getDefaultSession().isEmpty()); + } }