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..e149dc18a383 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; @@ -32,28 +33,41 @@ 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.base.Throwables; 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 + * on a dedicated thread, {@link #cancel()} stops it and joins that 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 +91,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 +101,33 @@ 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.info( + "Pipeline execution ended with an exception after cancel: {}", + String.valueOf(Throwables.getRootCause(e).getMessage())); + 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 +139,44 @@ public MetricResults metrics() { return asAttemptedOnlyMetricResults(metrics.value()); } + /** + * Cancels the Spark jobs of the pipeline and blocks until the execution thread has ended. An + * execution that already ended keeps its state. An interrupted caller keeps the interrupt flag + * and gets the current state. + */ @Override - public PipelineResult.State cancel() throws IOException { + public synchronized PipelineResult.State cancel() throws IOException { + if (state.isTerminal()) { + return state; + } + if (pipelineExecution.isDone()) { + try { + pipelineExecution.get(); + state = cancelRequested.get() ? State.CANCELLED : State.DONE; + } catch (ExecutionException e) { + state = cancelRequested.get() ? State.CANCELLED : State.FAILED; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return state; + } + cancelRequested.set(true); 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); - } + cancelSparkJobs.run(); + try { + pipelineExecution.get(); + } catch (ExecutionException e) { + LOG.info( + "Pipeline execution ended with an exception after cancel: {}", + String.valueOf(Throwables.getRootCause(e).getMessage())); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return state; } + 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..7aa12bbb869d 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,59 @@ public SparkStructuredStreamingPipelineResult run(final Pipeline pipeline) { PipelineTranslator.detectStreamingMode(pipeline, options); - final SparkSession sparkSession = SparkSessionFactory.getOrCreateSession(options); - final MetricsAccumulator metrics = MetricsAccumulator.getInstance(sparkSession); + final boolean releaseSession = !options.getUseActiveSparkSession(); + final SparkSession sparkSession = SparkSessionFactory.acquire(options); + final SparkContext sc = sparkSession.sparkContext(); + final MetricsAccumulator metrics; + try { + metrics = MetricsAccumulator.getInstance(sparkSession); + } catch (RuntimeException e) { + if (releaseSession) { + SparkSessionFactory.release(sparkSession); + } + throw e; + } - // 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 { + if (!sc.isStopped()) { + sc.cancelJobGroup(jobGroupId); + } + } catch (IllegalStateException e) { + // Context stopped concurrently. + } + }; final Future submissionFuture = runAsync( () -> { - EvaluationContext ctx = translatePipeline(sparkSession, pipeline); - ctxRef.set(ctx); - ctx.evaluate(); + try { + sc.setJobGroup(jobGroupId, "Beam " + jobName, true); + EvaluationContext ctx = translatePipeline(sparkSession, pipeline); + ctxRef.set(ctx); + if (!cancelRequested.get()) { + ctx.evaluate(); + } + } finally { + if (!sc.isStopped()) { + sc.clearJobGroup(); + } + 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 +262,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..792b7426eea7 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,12 @@ 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; + } 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..28fc7f127e49 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,47 @@ public class SparkSessionFactory { "/com.esotericsoftware/kryo-shaded", "/com/esotericsoftware/kryo-shaded"); + // Users per session created here, guarded by the class lock. + private static final Map OWNED_SESSIONS = new HashMap<>(); + /** - * Gets active {@link SparkSession} or creates one using {@link - * SparkStructuredStreamingPipelineOptions}. + * Returns the {@link SparkSession} for a pipeline, pair with {@link #release}. Without {@code + * useActiveSparkSession} the session is created here unless a usable one exists, only sessions + * created here are stopped on release. */ - public static SparkSession getOrCreateSession(SparkStructuredStreamingPipelineOptions options) { + public static synchronized SparkSession acquire(SparkStructuredStreamingPipelineOptions options) { if (options.getUseActiveSparkSession()) { return SparkSession.active(); } - return sessionBuilder(options.getSparkMaster(), options).getOrCreate(); + // Spark 3 also returns stopped sessions. + boolean noUsableSession = + !isUsable(SparkSession.getActiveSession()) && !isUsable(SparkSession.getDefaultSession()); + SparkSession session = sessionBuilder(options.getSparkMaster(), options).getOrCreate(); + if (noUsableSession) { + OWNED_SESSIONS.put(session, 1); + } else { + OWNED_SESSIONS.computeIfPresent(session, (unused, count) -> count + 1); + } + return session; + } + + /** Releases a session from {@link #acquire}, stops it when it was created here and unused. */ + 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..a042a37c64d6 --- /dev/null +++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResultTest.java @@ -0,0 +1,147 @@ +/* + * 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.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +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.After; +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 ExecutorService executor = Executors.newSingleThreadExecutor(); + private final AtomicInteger cancelSparkJobsCalls = new AtomicInteger(); + + @After + public void shutdownExecutor() { + executor.shutdownNow(); + } + + private SparkStructuredStreamingPipelineResult result(Future execution, Runnable onCancel) { + Runnable cancelSparkJobs = + () -> { + cancelSparkJobsCalls.incrementAndGet(); + onCancel.run(); + }; + return new SparkStructuredStreamingPipelineResult( + execution, () -> null, new MetricsAccumulator(), new AtomicBoolean(), cancelSparkJobs); + } + + @Test + public void testCancelJoinsExecutionThread() throws Exception { + CountDownLatch jobsCancelled = new CountDownLatch(1); + AtomicBoolean finished = new AtomicBoolean(); + Future execution = + executor.submit( + () -> { + jobsCancelled.await(); + finished.set(true); + return null; + }); + SparkStructuredStreamingPipelineResult result = result(execution, jobsCancelled::countDown); + + assertThat(result.cancel(), is(State.CANCELLED)); + assertTrue("cancel returned before the execution thread ended", finished.get()); + assertThat(result.getState(), is(State.CANCELLED)); + assertThat(cancelSparkJobsCalls.get(), is(1)); + + assertThat(result.cancel(), is(State.CANCELLED)); + assertThat(cancelSparkJobsCalls.get(), is(1)); + } + + @Test + public void testCancelAfterCompletionKeepsTerminalState() throws Exception { + SparkStructuredStreamingPipelineResult result = + result(CompletableFuture.completedFuture(null), () -> {}); + + assertThat(result.waitUntilFinish(), is(State.DONE)); + assertThat(result.cancel(), is(State.DONE)); + assertThat(result.getState(), is(State.DONE)); + assertThat(cancelSparkJobsCalls.get(), is(0)); + } + + @Test + public void testCancelOfUnobservedCompletionReportsDone() throws Exception { + SparkStructuredStreamingPipelineResult result = + result(CompletableFuture.completedFuture(null), () -> {}); + + assertThat(result.cancel(), is(State.DONE)); + assertThat(cancelSparkJobsCalls.get(), is(0)); + } + + @Test + public void testCancelOfUnobservedFailureReportsFailed() throws Exception { + CompletableFuture failed = new CompletableFuture<>(); + failed.completeExceptionally(new IllegalStateException("boom")); + SparkStructuredStreamingPipelineResult result = result(failed, () -> {}); + + assertThat(result.cancel(), is(State.FAILED)); + assertThat(cancelSparkJobsCalls.get(), is(0)); + assertThrows(RuntimeException.class, result::waitUntilFinish); + } + + @Test + public void testInterruptedCancelKeepsStateAndInterruptFlag() throws Exception { + CountDownLatch jobsCancelled = new CountDownLatch(1); + Future execution = + executor.submit( + () -> { + jobsCancelled.await(); + return null; + }); + SparkStructuredStreamingPipelineResult result = result(execution, () -> {}); + + Thread.currentThread().interrupt(); + assertThat(result.cancel(), is(State.RUNNING)); + assertTrue("interrupt flag not restored", Thread.interrupted()); + + jobsCancelled.countDown(); + assertThat(result.cancel(), is(State.CANCELLED)); + } + + @Test + public void testFailureAfterCancelIsCancelled() throws Exception { + CountDownLatch jobsCancelled = new CountDownLatch(1); + Future execution = + executor.submit( + () -> { + jobsCancelled.await(); + throw new IllegalStateException("job cancelled"); + }); + SparkStructuredStreamingPipelineResult result = result(execution, jobsCancelled::countDown); + + assertThat(result.cancel(), is(State.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..7f3749aa4324 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,7 +41,10 @@ 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.After; import org.junit.Ignore; import org.junit.Rule; import org.junit.Test; @@ -62,6 +70,41 @@ 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); + private static volatile CountDownLatch release = new CountDownLatch(1); + + /** Signals started, then blocks until the task is interrupted or release is counted down. */ + private static class BlockingDoFn extends DoFn { + @ProcessElement + public void processElement(ProcessContext c) throws InterruptedException { + started.countDown(); + while (!TaskContext.get().isInterrupted() && !release.await(50, TimeUnit.MILLISECONDS)) { + // wait for cancel + } + c.output(c.element()); + } + } + + @After + public void releaseBlockedDoFn() { + release.countDown(); + } + + private SparkStructuredStreamingPipelineResult runBlockingPipeline( + SparkStructuredStreamingPipelineOptions options) throws InterruptedException { + started = new CountDownLatch(1); + release = new CountDownLatch(1); + Pipeline pipeline = Pipeline.create(options); + 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() { @@ -222,4 +265,43 @@ public void testStreamingPipelineTimeoutState() throws Exception { public void testBatchPipelineTimeoutState() throws Exception { testTimeoutPipeline(getBatchOptions()); } + + @Test + public void testBatchCancelStopsRunningJob() throws Exception { + SparkStructuredStreamingPipelineResult result = runBlockingPipeline(getBatchOptions()); + assertThat(result.cancel(), is(PipelineResult.State.CANCELLED)); + assertThat(result.getState(), 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(getBatchOptions()); + assertThat(result.cancel(), is(PipelineResult.State.CANCELLED)); + assertFalse("shared session stopped", session.sparkContext().isStopped()); + } finally { + session.stop(); + } + } + + @Test + public void testActiveSessionPipelineKeepsRunnerSession() throws Exception { + SparkStructuredStreamingPipelineResult owner = runBlockingPipeline(getBatchOptions()); + SparkSession session = SparkSession.getDefaultSession().get(); + + SparkStructuredStreamingPipelineOptions active = + PipelineOptionsFactory.create().as(SparkStructuredStreamingPipelineOptions.class); + active.setRunner(SparkStructuredStreamingRunner.class); + active.setUseActiveSparkSession(true); + Pipeline guest = Pipeline.create(active); + guest.apply(Create.of("guest")).apply(printParDo("guest")); + assertThat(guest.run().waitUntilFinish(), is(PipelineResult.State.DONE)); + assertFalse("guest pipeline stopped the runner session", session.sparkContext().isStopped()); + + release.countDown(); + assertThat(owner.waitUntilFinish(), is(PipelineResult.State.DONE)); + assertTrue("owner did not stop its session", session.sparkContext().isStopped()); + } }