Improve ValidatesRunner performance. Parallelize runs

.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json

@@ -1,5 +1,5 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run!",
-  "modification": 2,
+  "modification": 7,
 }
 

runners/google-cloud-dataflow-java/build.gradle

@@ -236,16 +236,31 @@ def commonRunnerV2ExcludeCategories = [
   'org.apache.beam.sdk.testing.UsesBoundedTrieMetrics', // Dataflow QM as of now does not support returning back BoundedTrie in metric result.
 ]
 
+
+def validatesRunnerParallelThreads = project.findProperty('parallelThreads') ?: '6'
+
 def isValidatesRunnerTestClass = { FileTreeElement element ->
   if (element.isDirectory()) {
     return true
   }
   if (!element.name.endsWith('.class') ||
-      element.name.startsWith('ValidateRunnerXlangTest')) {
+      element.name.startsWith('ValidateRunnerXlangTest') ||
+      element.name.startsWith('BeamParallelJunit4RunnerTest')) {
     return false
   }
-  return new String(element.file.bytes, java.nio.charset.StandardCharsets.ISO_8859_1)
-    .contains('Lorg/apache/beam/sdk/testing/ValidatesRunner;')
+  def content = new String(element.file.bytes, java.nio.charset.StandardCharsets.ISO_8859_1)
+  if (content.contains('Lorg/apache/beam/sdk/testing/ValidatesRunner;')) {
+    return true
+  }
+  if (content.contains('Lorg/junit/experimental/runners/Enclosed;')) {
+    def prefix = element.name.substring(0, element.name.length() - 6) + '$'
+    return element.file.parentFile.listFiles()?.any { File f ->
+      f.name.startsWith(prefix) && f.name.endsWith('.class') &&
+        new String(f.bytes, java.nio.charset.StandardCharsets.ISO_8859_1)
+          .contains('Lorg/apache/beam/sdk/testing/ValidatesRunner;')
+    } ?: false
+  }
+  return false
 }
 
 def createLegacyWorkerValidatesRunnerTest = { Map args ->
@@ -258,10 +273,12 @@ def createLegacyWorkerValidatesRunnerTest = { Map args ->
     dependsOn ":runners:google-cloud-dataflow-java:worker:shadowJar"
 
     systemProperty "beamTestPipelineOptions", JsonOutput.toJson(pipelineOptions)
+    systemProperty "beam.validatesRunner.parallelThreads", validatesRunnerParallelThreads
 
     // Increase test parallelism up to the number of Gradle workers. By default this is equal
     // to the number of CPU cores, but can be increased by setting --max-workers=N.
     maxParallelForks Integer.MAX_VALUE
+    maxHeapSize = '2g'
     classpath = configurations.validatesRunner
     testClassesDirs = files(project(":sdks:java:core").sourceSets.test.output.classesDirs) +
       files(project(project.path).sourceSets.test.output.classesDirs)
@@ -292,10 +309,12 @@ def createRunnerV2ValidatesRunnerTest = { Map args ->
   return tasks.create(name: name, type: Test, group: "Verification") {
     dependsOn buildAndPushDockerJavaContainer
     systemProperty "beamTestPipelineOptions", JsonOutput.toJson(pipelineOptions)
+    systemProperty "beam.validatesRunner.parallelThreads", validatesRunnerParallelThreads
 
     // Increase test parallelism up to the number of Gradle workers. By default this is equal
     // to the number of CPU cores, but can be increased by setting --max-workers=N.
     maxParallelForks Integer.MAX_VALUE
+    maxHeapSize = '768m'
     classpath = configurations.validatesRunner
     testClassesDirs = files(project(":sdks:java:core").sourceSets.test.output.classesDirs) +
       files(project(project.path).sourceSets.test.output.classesDirs)

sdks/java/core/src/main/java/org/apache/beam/sdk/testing/BeamParallelJunit4Runner.java

@@ -0,0 +1,189 @@
+/*
+ * 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.sdk.testing;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.apache.beam.sdk.annotations.Internal;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.notification.RunNotifier;
+import org.junit.runners.BlockJUnit4ClassRunner;
+import org.junit.runners.model.FrameworkMethod;
+import org.junit.runners.model.InitializationError;
+import org.junit.runners.model.RunnerScheduler;
+
+/**
+ * A JUnit 4 runner that supports concurrent execution of {@code @Test} methods within a test class,
+ * with independent parallelism settings for {@link ValidatesRunner} tests versus other tests:
+ *
+ * <ul>
+ *   <li>{@code -Dbeam.validatesRunner.parallelThreads=N}: number of concurrent threads per test
+ *       worker JVM for {@code @Test} methods annotated with
+ *       {@code @Category(ValidatesRunner.class)} (defaults to {@code beam.test.parallelThreads}).
+ *   <li>{@code -Dbeam.test.parallelThreads=N}: number of concurrent threads per test worker JVM for
+ *       all other {@code @Test} methods (defaults to {@code 1}, i.e., sequential execution).
+ * </ul>
+ *
+ * <p>Classes or methods annotated with {@link SerialTest} are always executed sequentially on the
+ * calling thread after draining any in-flight parallel test methods in the class.
+ */
+@Internal
+public final class BeamParallelJunit4Runner extends BlockJUnit4ClassRunner {
+
+  /**
+   * Marks a test class or {@code @Test} method as requiring serial (non-parallel) execution even
+   * when parallel test execution is enabled.
+   */
+  @Retention(RetentionPolicy.RUNTIME)
+  @Target({ElementType.TYPE, ElementType.METHOD})
+  public @interface SerialTest {}
+
+  public static final String VALIDATES_RUNNER_THREADS_PROPERTY =
+      "beam.validatesRunner.parallelThreads";
+  public static final String DEFAULT_TEST_THREADS_PROPERTY = "beam.test.parallelThreads";
+
+  private static final ConcurrentHashMap<String, ExecutorService> EXECUTORS =
+      new ConcurrentHashMap<>();
+
+  private static final ConcurrentHashMap<Class<?>, Boolean> CLASS_SERIAL_CACHE =
+      new ConcurrentHashMap<>();
+
+  private CompletableFuture<Void> pendingFutures = CompletableFuture.allOf();
+
+  private static @Nullable ExecutorService getOrCreateExecutor(String poolKey, int threads) {
+    if (threads <= 1) {
+      return null;
+    }
+    return EXECUTORS.computeIfAbsent(
+        poolKey + ":" + threads,
+        k -> {
+          AtomicInteger counter = new AtomicInteger(1);
+          return Executors.newFixedThreadPool(
+              threads,
+              runnable -> {
+                Thread thread = new Thread(runnable);
+                thread.setDaemon(true);
+                thread.setName(poolKey + "-" + counter.getAndIncrement());
+                return thread;
+              });
+        });
+  }
+
+  public BeamParallelJunit4Runner(Class<?> klass) throws InitializationError {
+    super(klass);
+    setScheduler(
+        new RunnerScheduler() {
+          @Override
+          public void schedule(Runnable childStatement) {
+            childStatement.run();
+          }
+
+          @Override
+          public void finished() {
+            awaitPendingFutures();
+          }
+        });
+  }
+
+  private void awaitPendingFutures() {
+    try {
+      pendingFutures.join();
+    } finally {
+      pendingFutures = CompletableFuture.allOf();
+    }
+  }
+
+  @Override
+  protected void runChild(final FrameworkMethod method, final RunNotifier notifier) {
+    if (isIgnored(method)) {
+      super.runChild(method, notifier);
+      return;
+    }
+    ExecutorService executor = selectExecutor(method);
+    if (executor == null) {
+      super.runChild(method, notifier);
+      return;
+    }
+    if (isClassMarkedSerial(getTestClass().getJavaClass()) || isMethodMarkedSerial(method)) {
+      // This is a serial test. Make sure to wait for any in-flight tests to complete, and then run
+      // this test
+      // serially.
+      awaitPendingFutures();
+      super.runChild(method, notifier);
+      return;
+    }
+    pendingFutures =
+        CompletableFuture.allOf(
+            pendingFutures,
+            CompletableFuture.runAsync(() -> super.runChild(method, notifier), executor));
+  }
+
+  private @Nullable ExecutorService selectExecutor(FrameworkMethod method) {
+    int defaultThreads = Integer.getInteger(DEFAULT_TEST_THREADS_PROPERTY, 1);
+    if (isValidatesRunnerMethod(method)) {
+      int vrThreads = Integer.getInteger(VALIDATES_RUNNER_THREADS_PROPERTY, defaultThreads);
+      return getOrCreateExecutor("beam-vr-worker", vrThreads);
+    }
+    return getOrCreateExecutor("beam-test-worker", defaultThreads);
+  }
+
+  private boolean isValidatesRunnerMethod(FrameworkMethod method) {
+    return hasValidatesRunnerCategory(method.getAnnotation(Category.class))
+        || hasValidatesRunnerCategory(getTestClass().getAnnotation(Category.class));
+  }
+
+  private static boolean hasValidatesRunnerCategory(@Nullable Category category) {
+    if (category == null) {
+      return false;
+    }
+    for (Class<?> c : category.value()) {
+      if (ValidatesRunner.class.isAssignableFrom(c)) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  private static boolean isMethodMarkedSerial(FrameworkMethod method) {
+    return method.getAnnotation(SerialTest.class) != null;
+  }
+
+  private static boolean isClassMarkedSerial(@Nullable Class<?> clazz) {
+    if (clazz == null || clazz == Object.class) {
+      return false;
+    }
+    return CLASS_SERIAL_CACHE.computeIfAbsent(clazz, BeamParallelJunit4Runner::computeClassSerial);
+  }
+
+  private static boolean computeClassSerial(@Nullable Class<?> clazz) {
+    if (clazz == null || clazz == Object.class) {
+      return false;
+    }
+    return clazz.isAnnotationPresent(SerialTest.class)
+        || computeClassSerial(clazz.getSuperclass())
+        || computeClassSerial(clazz.getEnclosingClass());
+  }
+}

sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/resources/ClasspathScanningResourcesDetector.java

@@ -82,7 +82,6 @@ public List<String> detect(@Nullable ClassLoader classLoader) {
       if (snapshot != null && snapshot.matches(classLoader, currentJavaClassPath)) {
         return new ArrayList<>(snapshot.files);
       }
-
       List<File> classpathContents;
       if (classLoader != null) {
         classpathContents =

sdks/java/core/src/test/java/org/apache/beam/sdk/PipelineTest.java

@@ -42,6 +42,7 @@
 import org.apache.beam.sdk.runners.PTransformOverride;
 import org.apache.beam.sdk.runners.PTransformOverrideFactory;
 import org.apache.beam.sdk.runners.TransformHierarchy.Node;
+import org.apache.beam.sdk.testing.BeamParallelJunit4Runner;
 import org.apache.beam.sdk.testing.CrashingRunner;
 import org.apache.beam.sdk.testing.ExpectedLogs;
 import org.apache.beam.sdk.testing.NeedsRunner;
@@ -77,10 +78,9 @@
 import org.junit.rules.ExpectedException;
 import org.junit.rules.Timeout;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 
 /** Tests for Pipeline. */
-@RunWith(JUnit4.class)
+@RunWith(BeamParallelJunit4Runner.class)
 @SuppressWarnings({
   "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
 })
@@ -273,6 +273,7 @@ public void testToString() {
   }
 
   @Test
+  @BeamParallelJunit4Runner.SerialTest
   public void testStableUniqueNameOff() {
     pipeline.enableAbandonedNodeEnforcement(false);
 
@@ -285,6 +286,7 @@ public void testStableUniqueNameOff() {
   }
 
   @Test
+  @BeamParallelJunit4Runner.SerialTest
   public void testStableUniqueNameWarning() {
     pipeline.enableAbandonedNodeEnforcement(false);
 
@@ -297,6 +299,7 @@ public void testStableUniqueNameWarning() {
   }
 
   @Test
+  @BeamParallelJunit4Runner.SerialTest
   public void testStableUniqueNameError() {
     pipeline.getOptions().setStableUniqueNames(CheckEnabled.ERROR);
 

sdks/java/core/src/test/java/org/apache/beam/sdk/SdkHarnessEnvironmentTest.java

@@ -33,6 +33,7 @@
 import org.apache.beam.sdk.options.SdkHarnessOptions;
 import org.apache.beam.sdk.options.SdkHarnessOptions.LogLevel;
 import org.apache.beam.sdk.options.SdkHarnessOptions.SdkHarnessLogLevelOverrides;
+import org.apache.beam.sdk.testing.BeamParallelJunit4Runner;
 import org.apache.beam.sdk.testing.ExpectedLogs;
 import org.apache.beam.sdk.testing.ExpectedLogs.LogSaver;
 import org.apache.beam.sdk.testing.PAssert;
@@ -49,10 +50,9 @@
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 
 /** Tests that validate the SDK harness is configured correctly for a runner. */
-@RunWith(JUnit4.class)
+@RunWith(BeamParallelJunit4Runner.class)
 public class SdkHarnessEnvironmentTest {
 
   @Rule public final TestPipeline p = TestPipeline.create();

sdks/java/core/src/test/java/org/apache/beam/sdk/io/CountingSourceTest.java

@@ -31,6 +31,7 @@
 import org.apache.beam.sdk.io.CountingSource.UnboundedCountingSource;
 import org.apache.beam.sdk.io.UnboundedSource.UnboundedReader;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.testing.BeamParallelJunit4Runner;
 import org.apache.beam.sdk.testing.NeedsRunner;
 import org.apache.beam.sdk.testing.PAssert;
 import org.apache.beam.sdk.testing.TestPipeline;
@@ -53,10 +54,9 @@
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 
 /** Tests of {@link CountingSource}. */
-@RunWith(JUnit4.class)
+@RunWith(BeamParallelJunit4Runner.class)
 public class CountingSourceTest {
 
   public static void addCountingAsserts(PCollection<Long> input, long numElements) {