Run portable ValidatesRunner tests one at a time

.github/trigger_files/beam_PostCommit_Java_PVR_Flink_Batch.json

@@ -1,6 +1,6 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
   "https://github.com/apache/beam/pull/32440": "test new datastream runner for batch",
-  "modification": 2,
+  "modification": 8,
   "https://github.com/apache/beam/pull/39990": "removing dead code from FnApiDoFnRunner"
 }

.github/trigger_files/beam_PostCommit_Java_PVR_Flink_Streaming.json

@@ -1,6 +1,6 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "modification": 2,
+  "modification": 3,
   "https://github.com/apache/beam/pull/32440": "test new datastream runner for batch",
   "https://github.com/apache/beam/pull/39990": "removing dead code from FnApiDoFnRunner"
 }

runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContextFactory.java

@@ -22,18 +22,16 @@
 import org.apache.beam.runners.fnexecution.control.DefaultExecutableStageContext;
 import org.apache.beam.runners.fnexecution.control.ExecutableStageContext;
 import org.apache.beam.runners.fnexecution.control.ReferenceCountingExecutableStageContextFactory;
+import org.apache.beam.runners.fnexecution.control.StageBundleFactory;
 import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
+import org.apache.beam.sdk.util.construction.graph.ExecutableStage;
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
 
 /** Singleton class that contains one {@link ExecutableStageContext.Factory} per job. */
 public class FlinkExecutableStageContextFactory implements ExecutableStageContext.Factory {
 
   private static final FlinkExecutableStageContextFactory instance =
       new FlinkExecutableStageContextFactory();
-  // This map should only ever have a single element, as each job will have its own
-  // classloader and therefore its own instance of FlinkExecutableStageContextFactory. This
-  // code supports multiple JobInfos in order to provide a sensible implementation of
-  // Factory.get(JobInfo), which in theory could be called with different JobInfos.
   private static final ConcurrentMap<String, ExecutableStageContext.Factory> jobFactories =
       new ConcurrentHashMap<>();
 
@@ -48,16 +46,41 @@ public ExecutableStageContext get(JobInfo jobInfo) {
     ExecutableStageContext.Factory jobFactory =
         jobFactories.computeIfAbsent(
             jobInfo.jobId(),
-            k -> {
-              return ReferenceCountingExecutableStageContextFactory.create(
-                  DefaultExecutableStageContext::create,
-                  // Clean up context immediately if its class is not loaded on Flink parent
-                  // classloader.
-                  (caller) ->
-                      caller.getClass().getClassLoader()
-                          != StreamExecutionEnvironment.class.getClassLoader());
-            });
-
+            k ->
+                ReferenceCountingExecutableStageContextFactory.create(
+                    info -> new JobScopedExecutableStageContext(info, k),
+                    // Release immediately when the factory is not on the Flink parent classloader
+                    // (so cleanup still works after that classloader is discarded). Otherwise keep
+                    // the default environment cache delay so DoFn @Teardown can finish first.
+                    caller ->
+                        caller.getClass().getClassLoader()
+                            != StreamExecutionEnvironment.class.getClassLoader()));
     return jobFactory.get(jobInfo);
   }
+
+  private static class JobScopedExecutableStageContext implements ExecutableStageContext {
+    private final DefaultExecutableStageContext delegate;
+    private final String jobId;
+
+    private JobScopedExecutableStageContext(JobInfo jobInfo, String jobId) {
+      this.delegate = DefaultExecutableStageContext.create(jobInfo);
+      this.jobId = jobId;
+    }
+
+    @Override
+    public StageBundleFactory getStageBundleFactory(ExecutableStage executableStage) {
+      return delegate.getStageBundleFactory(executableStage);
+    }
+
+    @Override
+    public void close() throws Exception {
+      try {
+        delegate.close();
+      } finally {
+        // Drop the per-job factory entry so a long-lived JVM that runs many jobs does not
+        // accumulate one entry per finished job.
+        jobFactories.remove(jobId);
+      }
+    }
+  }
 }

runners/flink/job-server/flink_job_server.gradle

@@ -129,6 +129,7 @@ runShadow {
     jvmArgs += ["-Dorg.slf4j.simpleLogger.defaultLogLevel=${project.property('logLevel')}"]
 }
 
+
 def portableValidatesRunnerTask(String name, String mode, boolean checkpointing, boolean docker) {
   def pipelineOptions = [
       // Limit resource consumption via parallelism
@@ -257,6 +258,7 @@ def flinkJobServerJvmArgs() {
 [validatesPortableRunnerDocker, validatesPortableRunnerBatchDataSet, validatesPortableRunnerBatch, validatesPortableRunnerStreaming, validatesPortableRunnerStreamingCheckpoint].each { task ->
   task.configure {
     jvmArgs += flinkJobServerJvmArgs()
+    systemProperty "beam.validatesRunner.parallelThreads", "1"
   }
 }
 

runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContextFactory.java

@@ -22,18 +22,16 @@
 import org.apache.beam.runners.fnexecution.control.DefaultExecutableStageContext;
 import org.apache.beam.runners.fnexecution.control.ExecutableStageContext;
 import org.apache.beam.runners.fnexecution.control.ReferenceCountingExecutableStageContextFactory;
+import org.apache.beam.runners.fnexecution.control.StageBundleFactory;
 import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
+import org.apache.beam.sdk.util.construction.graph.ExecutableStage;
 import org.apache.flink.api.java.ExecutionEnvironment;
 
 /** Singleton class that contains one {@link ExecutableStageContext.Factory} per job. */
 public class FlinkExecutableStageContextFactory implements ExecutableStageContext.Factory {
 
   private static final FlinkExecutableStageContextFactory instance =
       new FlinkExecutableStageContextFactory();
-  // This map should only ever have a single element, as each job will have its own
-  // classloader and therefore its own instance of FlinkExecutableStageContextFactory. This
-  // code supports multiple JobInfos in order to provide a sensible implementation of
-  // Factory.get(JobInfo), which in theory could be called with different JobInfos.
   private static final ConcurrentMap<String, ExecutableStageContext.Factory> jobFactories =
       new ConcurrentHashMap<>();
 
@@ -48,16 +46,41 @@ public ExecutableStageContext get(JobInfo jobInfo) {
     ExecutableStageContext.Factory jobFactory =
         jobFactories.computeIfAbsent(
             jobInfo.jobId(),
-            k -> {
-              return ReferenceCountingExecutableStageContextFactory.create(
-                  DefaultExecutableStageContext::create,
-                  // Clean up context immediately if its class is not loaded on Flink parent
-                  // classloader.
-                  (caller) ->
-                      caller.getClass().getClassLoader()
-                          != ExecutionEnvironment.class.getClassLoader());
-            });
-
+            k ->
+                ReferenceCountingExecutableStageContextFactory.create(
+                    info -> new JobScopedExecutableStageContext(info, k),
+                    // Release immediately when the factory is not on the Flink parent classloader
+                    // (so cleanup still works after that classloader is discarded). Otherwise keep
+                    // the default environment cache delay so DoFn @Teardown can finish first.
+                    caller ->
+                        caller.getClass().getClassLoader()
+                            != ExecutionEnvironment.class.getClassLoader()));
     return jobFactory.get(jobInfo);
   }
+
+  private static class JobScopedExecutableStageContext implements ExecutableStageContext {
+    private final DefaultExecutableStageContext delegate;
+    private final String jobId;
+
+    private JobScopedExecutableStageContext(JobInfo jobInfo, String jobId) {
+      this.delegate = DefaultExecutableStageContext.create(jobInfo);
+      this.jobId = jobId;
+    }
+
+    @Override
+    public StageBundleFactory getStageBundleFactory(ExecutableStage executableStage) {
+      return delegate.getStageBundleFactory(executableStage);
+    }
+
+    @Override
+    public void close() throws Exception {
+      try {
+        delegate.close();
+      } finally {
+        // Drop the per-job factory entry so a long-lived JVM that runs many jobs does not
+        // accumulate one entry per finished job.
+        jobFactories.remove(jobId);
+      }
+    }
+  }
 }