Fix stale ExecutorClassLoader leak and local JAR staging in portable Spark runner

.github/trigger_files/beam_PostCommit_Java_PVR_Spark4_StructuredStreaming.json

@@ -1,3 +1,4 @@
 {
-  "comment": "Modify this file in a trivial way to cause this test suite to run."
+  "comment": "Modify this file in a trivial way to cause this test suite to run.",
+  "modification": 1
 }

runners/java-job-service/src/main/java/org/apache/beam/runners/jobsubmission/JobInvoker.java

@@ -49,8 +49,21 @@ JobInvocation invoke(RunnerApi.Pipeline pipeline, Struct options, @Nullable Stri
   }
 
   private ListeningExecutorService createExecutorService(String name) {
+    final ClassLoader classLoader = getClass().getClassLoader();
+    ThreadFactory baseFactory = Executors.defaultThreadFactory();
     ThreadFactory threadFactory =
-        new ThreadFactoryBuilder().setNameFormat(name).setDaemon(true).build();
+        new ThreadFactoryBuilder()
+            .setThreadFactory(
+                r -> {
+                  Thread t = baseFactory.newThread(r);
+                  // Prevent job invoker threads from inheriting a transient task classloader
+                  // (e.g. a closed Spark ExecutorClassLoader) from gRPC executor threads.
+                  t.setContextClassLoader(classLoader);
+                  return t;
+                })
+            .setNameFormat(name)
+            .setDaemon(true)
+            .build();
     return MoreExecutors.listeningDecorator(Executors.newCachedThreadPool(threadFactory));
   }
 

runners/spark/src/main/java/org/apache/beam/runners/spark/SparkCommonPipelineOptions.java

@@ -90,11 +90,16 @@ public String create(PipelineOptions options) {
    */
   @Internal
   static void prepareFilesToStage(SparkCommonPipelineOptions options) {
-    if (!options.getSparkMaster().matches("local\\[?\\d*]?")) {
+    if (!isLocalSparkMaster(options)) {
       PipelineResources.prepareFilesForStaging(options);
     }
   }
 
+  @Internal
+  static boolean isLocalSparkMaster(SparkCommonPipelineOptions options) {
+    return options.getSparkMaster().matches("local\\[?[0-9\\*,]*\\]?");
+  }
+
   /**
    * Returns Spark's default storage level for the Dataset or RDD API based on the respective
    * runner.

runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineRunner.java

@@ -115,6 +115,8 @@ public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo)
             ? trimmedPipeline
             : GreedyPipelineFuser.fuse(trimmedPipeline).toPipeline();
 
+    // Avoid inheriting a stale Spark ExecutorClassLoader leaked onto shared executor threads.
+    Thread.currentThread().setContextClassLoader(SparkPipelineRunner.class.getClassLoader());
     prepareFilesToStage(pipelineOptions);
     PortablePipelineResult result;
     final JavaSparkContext jsc = SparkContextFactory.getSparkContext(pipelineOptions);
@@ -129,6 +131,12 @@ public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo)
     final ExecutorService executorService =
         Executors.newSingleThreadExecutor(
             new ThreadFactoryBuilder()
+                .setThreadFactory(
+                    r -> {
+                      Thread t = new Thread(r);
+                      t.setContextClassLoader(SparkPipelineRunner.class.getClassLoader());
+                      return t;
+                    })
                 .setDaemon(true)
                 .setNameFormat("DefaultSparkRunner-thread")
                 .build());

runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkContextFactory.java

@@ -19,12 +19,13 @@
 
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
 
-import javax.annotation.Nullable;
+import org.apache.beam.runners.spark.SparkCommonPipelineOptions;
 import org.apache.beam.runners.spark.SparkContextOptions;
 import org.apache.beam.runners.spark.SparkPipelineOptions;
 import org.apache.beam.runners.spark.coders.SparkRunnerKryoRegistrator;
 import org.apache.spark.SparkConf;
 import org.apache.spark.api.java.JavaSparkContext;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -136,13 +137,20 @@ private static JavaSparkContext createSparkContext(SparkPipelineOptions options)
       conf.setMaster(options.getSparkMaster());
     }
 
-    if (options.getFilesToStage() != null && !options.getFilesToStage().isEmpty()) {
+    // Skip staging jars in local mode since the driver and executor share the same JVM classpath.
+    // In portable runs, filesToStage is populated with the full classpath, which would otherwise
+    // be copied into a per-context temporary userFiles directory that is deleted on context stop.
+    if (options.getFilesToStage() != null
+        && !options.getFilesToStage().isEmpty()
+        && !SparkCommonPipelineOptions.isLocalSparkMaster(options)) {
       conf.setJars(options.getFilesToStage().toArray(new String[0]));
     }
 
     conf.setAppName(options.getAppName());
     // register immutable collections serializers because the SDK uses them.
     conf.set("spark.kryo.registrator", SparkRunnerKryoRegistrator.class.getName());
+    // Ensure SparkContext does not inherit a stale ExecutorClassLoader from a previous Spark job.
+    Thread.currentThread().setContextClassLoader(SparkContextFactory.class.getClassLoader());
     return new JavaSparkContext(conf);
   }
 }

runners/spark/src/test/java/org/apache/beam/runners/spark/SparkPipelineOptionsTest.java

@@ -38,4 +38,21 @@ public void testCheckpointDirDefaultValue() {
         "/tmp/" + sparkCommonPipelineOptions.getJobName(),
         sparkCommonPipelineOptions.getCheckpointDir());
   }
+
+  @Test
+  public void testPrepareFilesToStageSkipsStagingInLocalMode() {
+    for (String master : new String[] {"local", "local[4]", "local[*]", "local[2,3]"}) {
+      SparkCommonPipelineOptions options =
+          PipelineOptionsFactory.create().as(SparkCommonPipelineOptions.class);
+      options.setSparkMaster(master);
+      SparkCommonPipelineOptions.prepareFilesToStage(options);
+      assertEquals(null, options.getFilesToStage());
+
+      java.util.List<String> explicitFiles =
+          java.util.Collections.singletonList("/tmp/nonexistent.jar");
+      options.setFilesToStage(explicitFiles);
+      SparkCommonPipelineOptions.prepareFilesToStage(options);
+      assertEquals(explicitFiles, options.getFilesToStage());
+    }
+  }
 }