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..14fadd55fc4f 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 @@ -22,6 +22,7 @@ import java.io.IOException; import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; @@ -35,25 +36,36 @@ 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; public class SparkStructuredStreamingPipelineResult implements PipelineResult { + private static final Logger LOG = + LoggerFactory.getLogger(SparkStructuredStreamingPipelineResult.class); + + /** Upper bound on how long {@link #cancel()} waits for the execution thread to end. */ + private static final long CANCEL_WAIT_SECONDS = 60; + 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 final ExecutorService executor; private PipelineResult.State state; SparkStructuredStreamingPipelineResult( Future pipelineExecution, Supplier evaluationContext, MetricsAccumulator metrics, - final @Nullable Runnable onTerminalState) { + final @Nullable Runnable onTerminalState, + ExecutorService executor) { this.pipelineExecution = pipelineExecution; this.evaluationContext = evaluationContext; this.metrics = metrics; this.onTerminalState = onTerminalState; + this.executor = executor; // pipelineExecution is expected to have started executing eagerly. this.state = State.RUNNING; } @@ -117,6 +129,11 @@ public MetricResults metrics() { return asAttemptedOnlyMetricResults(metrics.value()); } + /** + * Cancels the execution and waits up to {@link #CANCEL_WAIT_SECONDS} for the execution thread to + * end before the terminal state callback stops the session. An interrupted caller returns without + * the callback, the state stays RUNNING. + */ @Override public PipelineResult.State cancel() throws IOException { EvaluationContext ctx = evaluationContext.get(); @@ -124,6 +141,16 @@ public PipelineResult.State cancel() throws IOException { ctx.stop(); } pipelineExecution.cancel(true); + try { + if (!executor.awaitTermination(CANCEL_WAIT_SECONDS, TimeUnit.SECONDS)) { + LOG.warn( + "Pipeline execution still running {} s after cancel, stopping the session anyway.", + CANCEL_WAIT_SECONDS); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return state; + } offerNewState(PipelineResult.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..8cce4c18e866 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 @@ -152,7 +152,7 @@ public SparkStructuredStreamingPipelineResult run(final Pipeline pipeline) { // evaluation on cancel. Remains null until translation completes. final AtomicReference ctxRef = new AtomicReference<>(); - final Future submissionFuture = + final Submission submission = runAsync( () -> { EvaluationContext ctx = translatePipeline(sparkSession, pipeline); @@ -162,10 +162,11 @@ public SparkStructuredStreamingPipelineResult run(final Pipeline pipeline) { final SparkStructuredStreamingPipelineResult result = new SparkStructuredStreamingPipelineResult( - submissionFuture, + submission.future, ctxRef::get, metrics, - sparkStopFn(sparkSession, options.getUseActiveSparkSession())); + sparkStopFn(sparkSession, options.getUseActiveSparkSession()), + submission.executor); if (options.getEnableSparkMetricSinks()) { registerMetricsSource(options.getAppName(), metrics); @@ -217,7 +218,18 @@ private void startMetricsPusher( } } - private static Future runAsync(Runnable task) { + /** The future of a submitted pipeline and the executor that runs it, for cancel to await. */ + private static final class Submission { + private final Future future; + private final ExecutorService executor; + + Submission(Future future, ExecutorService executor) { + this.future = future; + this.executor = executor; + } + } + + private static Submission runAsync(Runnable task) { ThreadFactory factory = new ThreadFactoryBuilder() .setDaemon(true) @@ -226,7 +238,7 @@ private static Future runAsync(Runnable task) { ExecutorService execService = Executors.newSingleThreadExecutor(factory); Future future = execService.submit(task); execService.shutdown(); - return future; + return new Submission(future, execService); } private static @Nullable Runnable sparkStopFn(SparkSession session, boolean isProvided) { 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..0ef31df3c4a9 --- /dev/null +++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResultTest.java @@ -0,0 +1,91 @@ +/* + * 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.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Uninterruptibles.sleepUninterruptibly; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator; +import org.apache.beam.sdk.PipelineResult; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link SparkStructuredStreamingPipelineResult#cancel()}. */ +@RunWith(JUnit4.class) +public class SparkStructuredStreamingPipelineResultTest { + + private static SparkStructuredStreamingPipelineResult result( + Future f, Runnable cb, ExecutorService exec) { + return new SparkStructuredStreamingPipelineResult( + f, () -> null, new MetricsAccumulator(), cb, exec); + } + + private static Future submit(ExecutorService exec, Runnable r) throws Exception { + CountDownLatch started = new CountDownLatch(1); + Future f = + exec.submit( + () -> { + started.countDown(); + r.run(); + }); + exec.shutdown(); + started.await(); + return f; + } + + @Test + public void testCancelWaitsForExecutionThread() throws Exception { + CountDownLatch block = new CountDownLatch(1); + AtomicBoolean finished = new AtomicBoolean(); + AtomicBoolean observed = new AtomicBoolean(); + ExecutorService exec = Executors.newSingleThreadExecutor(); + Future f = + submit( + exec, + () -> { + try { + block.await(); + } catch (InterruptedException ignored) { + } + sleepUninterruptibly(300, TimeUnit.MILLISECONDS); + finished.set(true); + }); + SparkStructuredStreamingPipelineResult res = + result(f, () -> observed.set(finished.get()), exec); + assertEquals(PipelineResult.State.CANCELLED, res.cancel()); + assertTrue(observed.get()); + assertEquals(PipelineResult.State.CANCELLED, res.getState()); + } + + @Test + public void testCancelAwaitsTaskIgnoringInterrupt() throws Exception { + ExecutorService exec = Executors.newSingleThreadExecutor(); + Future f = submit(exec, () -> sleepUninterruptibly(200, TimeUnit.MILLISECONDS)); + SparkStructuredStreamingPipelineResult res = result(f, null, exec); + assertEquals(PipelineResult.State.CANCELLED, res.cancel()); + assertTrue(exec.isTerminated()); + } +}