diff --git a/.github/trigger_files/beam_PostCommit_Java_PVR_Flink_Batch.json b/.github/trigger_files/beam_PostCommit_Java_PVR_Flink_Batch.json index 6f486a1f8c14..73e12ca8961b 100644 --- a/.github/trigger_files/beam_PostCommit_Java_PVR_Flink_Batch.json +++ b/.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" } diff --git a/.github/trigger_files/beam_PostCommit_Java_PVR_Flink_Streaming.json b/.github/trigger_files/beam_PostCommit_Java_PVR_Flink_Streaming.json index 626e4fea5c8b..ebd004055319 100644 --- a/.github/trigger_files/beam_PostCommit_Java_PVR_Flink_Streaming.json +++ b/.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" } diff --git a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContextFactory.java b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContextFactory.java index 3f42eb93e4e6..afcfb9d67300 100644 --- a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContextFactory.java +++ b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContextFactory.java @@ -22,7 +22,9 @@ 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. */ @@ -30,10 +32,6 @@ public class FlinkExecutableStageContextFactory implements ExecutableStageContex 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 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); + } + } + } } diff --git a/runners/flink/job-server/flink_job_server.gradle b/runners/flink/job-server/flink_job_server.gradle index 6f39b67abf74..edf55807503f 100644 --- a/runners/flink/job-server/flink_job_server.gradle +++ b/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" } } diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContextFactory.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContextFactory.java index a2e3bb3142bf..5a11336add14 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContextFactory.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContextFactory.java @@ -22,7 +22,9 @@ 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. */ @@ -30,10 +32,6 @@ public class FlinkExecutableStageContextFactory implements ExecutableStageContex 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 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); + } + } + } }