From dc62a5d6ac1223bca10c863d0408197214b75f90 Mon Sep 17 00:00:00 2001 From: Tobias Kaymak Date: Fri, 11 Sep 2026 07:35:13 +0000 Subject: [PATCH] [Spark] Let cancel() wait for the execution thread before stopping the session SparkStructuredStreamingPipelineResult.cancel() interrupted the execution thread and ran the terminal state callback at once, which stops the SparkSession. The thread kept translating or evaluating on a stopped SparkContext. When that happened during the static initialization of PipelineTranslatorBatch the class was poisoned for the JVM and every later batch pipeline failed with NoClassDefFoundError, seen in the Spark Versions PreCommit on #40090. runAsync now hands the result its single thread executor and cancel() waits for it to terminate, bounded at 60 seconds, before the callback runs. --- ...parkStructuredStreamingPipelineResult.java | 29 +++++- .../SparkStructuredStreamingRunner.java | 22 ++++- ...StructuredStreamingPipelineResultTest.java | 91 +++++++++++++++++++ 3 files changed, 136 insertions(+), 6 deletions(-) create mode 100644 runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/SparkStructuredStreamingPipelineResultTest.java 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()); + } +}