[Spark] Make cancel() cancel the Spark jobs and stop only a session the runner created

.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
 }

.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"
 }

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<StreamingQuery> queries = new ArrayList<>();
-  private boolean stopped = false;
 
   StreamingEvaluationContext(
       Collection<? extends NamedDataset<?>> 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<StreamingQuery> toStop;
     synchronized (lock) {
-      if (stopped) {
+      if (isStopped()) {
         return;
       }
-      stopped = true;
+      super.stop();
       toStop = new ArrayList<>(queries);
     }
     for (StreamingQuery query : toStop) {

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<? extends @Nullable EvaluationContext> 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<? extends @Nullable EvaluationContext> 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;
   }
 }

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<EvaluationContext> 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;
-  }
 }

runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/EvaluationContext.java

@@ -52,6 +52,7 @@ public interface NamedDataset<T> {
 
   private final Collection<? extends NamedDataset<?>> leaves;
   private final SparkSession session;
+  private volatile boolean stopped = false;
 
   protected EvaluationContext(Collection<? extends NamedDataset<?>> leaves, SparkSession session) {
     this.leaves = leaves;
@@ -63,9 +64,13 @@ protected Collection<? extends NamedDataset<?>> 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 <T> void evaluate(String name, Dataset<T> ds) {
   }
 
   /**
-   * Stops any ongoing streaming execution triggered by this context.
-   *
-   * <p>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;

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<SparkSession, Integer> 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<SparkSession> session) {
+    return session.isDefined() && !session.get().sparkContext().isStopped();
   }
 
   /** Creates Spark session builder with some optimizations for local mode, e.g. in tests. */