[Spark] Let cancel() wait for the execution thread before stopping the session

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<? extends @Nullable EvaluationContext> evaluationContext;
   private final MetricsAccumulator metrics;
   private final @Nullable Runnable onTerminalState;
+  private final ExecutorService executor;
   private PipelineResult.State state;
 
   SparkStructuredStreamingPipelineResult(
       Future<?> pipelineExecution,
       Supplier<? extends @Nullable EvaluationContext> 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,13 +129,28 @@ 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();
     if (ctx != null) {
       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;
   }

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<EvaluationContext> 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) {

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());
+  }
+}