Make Java PreSubmit faster: Parallelize test

buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy

@@ -1259,6 +1259,7 @@ class BeamModulePlugin implements Plugin<Project> {
         useJUnit {}
         // default maxHeapSize on gradle 5 is 512m, lets increase to handle more demanding tests
         maxHeapSize = '2g'
+        systemProperty 'beam.test.parallelThreads', project.findProperty('testParallelThreads') ?: '8'
         // Windows OS: Snappy needs an executable temp dir for native lib. Default AppData/Temp
         // failing with Access error without elevated permissions
         if (System.getProperty("os.name").toLowerCase().contains("windows")) {

examples/java/src/test/java/org/apache/beam/examples/WindowedWordCountIT.java

@@ -36,6 +36,7 @@
 import org.apache.beam.sdk.io.fs.ResourceId;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.options.StreamingOptions;
+import org.apache.beam.sdk.testing.BeamParallelJunit4Runner;
 import org.apache.beam.sdk.testing.SerializableMatcher;
 import org.apache.beam.sdk.testing.StreamingIT;
 import org.apache.beam.sdk.testing.TestPipeline;
@@ -59,10 +60,9 @@
 import org.junit.experimental.categories.Category;
 import org.junit.rules.TestName;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 
 /** End-to-end integration test of {@link WindowedWordCount}. */
-@RunWith(JUnit4.class)
+@RunWith(BeamParallelJunit4Runner.class)
 public class WindowedWordCountIT {
 
   @Rule public TestName testName = new TestName();

runners/direct-java/build.gradle

@@ -120,6 +120,7 @@ task needsRunnerTests(type: Test) {
   group = "Verification"
   description = "Runs tests that require a runner to validate that pipelines/transforms work correctly"
 
+  maxParallelForks 4
   testLogging.showStandardStreams = true
 
   String[] pipelineOptions = ["--runner=DirectRunner", "--runnerDeterminedSharding=false"]
@@ -160,6 +161,8 @@ task validatesRunner(type: Test) {
   group = "Verification"
   description "Validates Direct runner"
 
+  maxParallelForks 4
+
   String[] pipelineOptions = ["--runner=DirectRunner", "--runnerDeterminedSharding=false"]
   systemProperty "beamTestPipelineOptions", pipelineOptionsStringCrossPlatformHandling(pipelineOptions)
 

runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/TestDataflowRunnerTest.java

@@ -51,6 +51,7 @@
 import org.apache.beam.sdk.extensions.gcp.util.Transport;
 import org.apache.beam.sdk.io.FileSystems;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.testing.BeamParallelJunit4Runner;
 import org.apache.beam.sdk.testing.PAssert;
 import org.apache.beam.sdk.testing.SerializableMatcher;
 import org.apache.beam.sdk.testing.TestPipeline;
@@ -69,13 +70,12 @@
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 
 /** Tests for {@link TestDataflowRunner}. */
-@RunWith(JUnit4.class)
+@RunWith(BeamParallelJunit4Runner.class)
 public class TestDataflowRunnerTest {
   @Rule public ExpectedException expectedException = ExpectedException.none();
   @Mock private DataflowClient mockClient;

sdks/java/core/build.gradle

@@ -67,6 +67,7 @@ processResources {
 
 // Exclude tests that need a runner
 test {
+  maxParallelForks 4
   systemProperty "beamUseDummyRunner", "true"
   useJUnit {
     excludeCategories "org.apache.beam.sdk.testing.NeedsRunner"

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

@@ -182,6 +182,7 @@ private static boolean computeClassSerial(@Nullable Class<?> clazz) {
     if (clazz == null || clazz == Object.class) {
       return false;
     }
+    // Don't access CLASS_SERIAL_CACHE here, as ConcurrentHashMap does not support reentrancy.
     return clazz.isAnnotationPresent(SerialTest.class)
         || computeClassSerial(clazz.getSuperclass())
         || computeClassSerial(clazz.getEnclosingClass());

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

@@ -32,6 +32,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.Pipeline.PipelineVisitor;
 import org.apache.beam.sdk.PipelineRunner;
@@ -121,10 +122,12 @@ public class PAssert {
   private static final Counter failureCounter =
       Metrics.counter(PAssert.class, PAssert.FAILURE_COUNTER);
 
-  private static int assertCount = 0;
+  // Atomic so that PAssert transforms constructed concurrently (e.g. by tests running under
+  // BeamParallelJunit4Runner) never receive duplicate names within a pipeline.
+  private static final AtomicInteger assertCount = new AtomicInteger(0);
 
   private static String nextAssertionName() {
-    return "PAssert$" + assertCount++;
+    return "PAssert$" + assertCount.getAndIncrement();
   }
 
   // Do not instantiate.
@@ -786,7 +789,7 @@ PCollectionContentsAssert<T> satisfies(
       SerializableFunction<Iterable<T>, Void> checkerFn =
           (SerializableFunction) new MatcherCheckerFn<>(matcher);
       actual.apply(
-          "PAssert$" + assertCount++,
+          nextAssertionName(),
           new GroupThenAssert<>(checkerFn, rewindowingStrategy, paneExtractor, site));
       return this;
     }
@@ -940,7 +943,7 @@ public PCollectionSingletonIterableAssert<T> containsInAnyOrder() {
     public PCollectionSingletonIterableAssert<T> satisfies(
         SerializableFunction<Iterable<T>, Void> checkerFn) {
       actual.apply(
-          "PAssert$" + assertCount++,
+          nextAssertionName(),
           new GroupThenAssertForSingleton<>(checkerFn, rewindowingStrategy, paneExtractor, site));
       return this;
     }
@@ -1034,7 +1037,7 @@ private PCollectionSingletonAssert<T> withPanes(
     @Override
     public PCollectionSingletonAssert<T> satisfies(SerializableFunction<T, Void> checkerFn) {
       actual.apply(
-          "PAssert$" + assertCount++,
+          nextAssertionName(),
           new GroupThenAssertForSingleton<>(checkerFn, rewindowingStrategy, paneExtractor, site));
       return this;
     }
@@ -1173,7 +1176,7 @@ public PCollectionViewAssert<ElemT, ViewT> satisfies(
       actual
           .getPipeline()
           .apply(
-              "PAssert$" + assertCount++,
+              nextAssertionName(),
               new OneSideInputAssert<>(
                   CreateActual.from(actual, rewindowActuals, paneExtractor, view),
                   rewindowActuals.windowDummy(),

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

@@ -56,6 +56,7 @@
 import org.apache.beam.sdk.state.StateSpec;
 import org.apache.beam.sdk.state.StateSpecs;
 import org.apache.beam.sdk.state.ValueState;
+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;
@@ -91,10 +92,9 @@
 import org.junit.rules.TemporaryFolder;
 import org.junit.rules.Timeout;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 
 /** Tests for {@link FileIO}. */
-@RunWith(JUnit4.class)
+@RunWith(BeamParallelJunit4Runner.class)
 public class FileIOTest implements Serializable {
   @Rule public transient TestPipeline p = TestPipeline.create();
 

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

@@ -21,6 +21,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.is;
 
+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;
@@ -40,10 +41,9 @@
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 
 /** Tests for {@link GenerateSequence}. */
-@RunWith(JUnit4.class)
+@RunWith(BeamParallelJunit4Runner.class)
 public class GenerateSequenceTest {
   public static void addCountingAsserts(PCollection<Long> input, long start, long end) {
     // Count == numElements

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

@@ -59,6 +59,7 @@
 import org.apache.beam.sdk.io.FileIO.ReadableFile;
 import org.apache.beam.sdk.io.TFRecordIO.TFRecordCodec;
 import org.apache.beam.sdk.io.fs.MatchResult.Metadata;
+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;
@@ -77,10 +78,9 @@
 import org.junit.rules.ExpectedException;
 import org.junit.rules.TemporaryFolder;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 
 /** Tests for TFRecordIO Read and Write transforms. */
-@RunWith(JUnit4.class)
+@RunWith(BeamParallelJunit4Runner.class)
 public class TFRecordIOTest {
 
   /*

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

@@ -50,6 +50,7 @@
 import org.apache.beam.sdk.schemas.transforms.SchemaTransform;
 import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider;
 import org.apache.beam.sdk.schemas.transforms.providers.ErrorHandling;
+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;
@@ -67,10 +68,9 @@
 import org.junit.rules.ExpectedException;
 import org.junit.rules.TemporaryFolder;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 
 /** Tests for TFRecordIO Read and Write transforms. */
-@RunWith(JUnit4.class)
+@RunWith(BeamParallelJunit4Runner.class)
 public class TFRecordSchemaTransformProviderTest {
 
   /*

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

@@ -73,6 +73,7 @@
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.options.ValueProvider;
+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.SourceTestUtils;
@@ -106,7 +107,6 @@
 import org.junit.experimental.runners.Enclosed;
 import org.junit.rules.TemporaryFolder;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 import org.junit.runners.Parameterized;
 
 /** Tests for {@link TextIO.Read}. */
@@ -633,7 +633,7 @@ public void close() throws IOException {
   }
 
   /** Tests for some basic operations in {@link TextIO.Read}. */
-  @RunWith(JUnit4.class)
+  @RunWith(BeamParallelJunit4Runner.class)
   public static class BasicIOTest {
     @Rule public TemporaryFolder tempFolder = new TemporaryFolder();
     @Rule public TestPipeline p = TestPipeline.create();
@@ -1171,6 +1171,9 @@ public void testReadFilesWithFilename() throws IOException {
 
     @Test
     @Category({NeedsRunner.class, UsesUnboundedSplittableParDo.class})
+    // The watch terminates after 3s without new output; running concurrently with other pipelines
+    // in this JVM can stall the in-pipeline writer long enough to truncate the results.
+    @BeamParallelJunit4Runner.SerialTest
     public void testReadWatchForNewFiles() throws IOException, InterruptedException {
       final Path basePath = tempFolder.getRoot().toPath().resolve("readWatch");
       basePath.toFile().mkdir();
@@ -1202,7 +1205,7 @@ public void testReadWatchForNewFiles() throws IOException, InterruptedException
   }
 
   /** Tests for TextSource class. */
-  @RunWith(JUnit4.class)
+  @RunWith(BeamParallelJunit4Runner.class)
   public static class TextSourceTest {
     @Rule public transient TestPipeline pipeline = TestPipeline.create();
 

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

@@ -67,6 +67,7 @@
 import org.apache.beam.sdk.options.Description;
 import org.apache.beam.sdk.options.PipelineOptionsFactoryTest.TestPipelineOptions;
 import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
+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;
@@ -116,10 +117,9 @@
 import org.junit.rules.ExpectedException;
 import org.junit.rules.TemporaryFolder;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 
 /** Tests for the WriteFiles PTransform. */
-@RunWith(JUnit4.class)
+@RunWith(BeamParallelJunit4Runner.class)
 public class WriteFilesTest {
   @Rule public TemporaryFolder tmpFolder = new TemporaryFolder();
   @Rule public final TestPipeline p = TestPipeline.create();
@@ -633,6 +633,7 @@ public void testWriteEvictWritersWhenFullLruOrderAndReopen() throws IOException
 
   @Test
   @Category(NeedsRunner.class)
+  @BeamParallelJunit4Runner.SerialTest
   public void testWriteEvictWritersWhenFullCloseExceptionCleanup() {
     FailingCloseOnEvictSink.EVICTED_TEMP_FILE.set(null);
     FailingCloseOnEvictSink.THREW_ON_CLOSE.set(false);

sdks/java/core/src/test/java/org/apache/beam/sdk/schemas/AutoValueSchemaTest.java

@@ -40,6 +40,7 @@
 import org.apache.beam.sdk.schemas.annotations.SchemaFieldName;
 import org.apache.beam.sdk.schemas.annotations.SchemaFieldNumber;
 import org.apache.beam.sdk.schemas.utils.SchemaTestUtils;
+import org.apache.beam.sdk.testing.BeamParallelJunit4Runner;
 import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.sdk.util.SerializableUtils;
 import org.apache.beam.sdk.values.Row;
@@ -51,10 +52,9 @@
 import org.joda.time.Instant;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 
 /** Tests for {@link AutoValueSchema}. */
-@RunWith(JUnit4.class)
+@RunWith(BeamParallelJunit4Runner.class)
 public class AutoValueSchemaTest {
   static final DateTime DATE = DateTime.parse("1979-03-14");
   static final byte[] BYTE_ARRAY = "bytearray".getBytes(StandardCharsets.UTF_8);

sdks/java/core/src/test/java/org/apache/beam/sdk/schemas/transforms/ConvertTest.java

@@ -26,6 +26,7 @@
 import org.apache.beam.sdk.schemas.Schema.FieldType;
 import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
 import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes;
+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;
@@ -42,10 +43,9 @@
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 
 /** Tests for the {@link Convert} class. */
-@RunWith(JUnit4.class)
+@RunWith(BeamParallelJunit4Runner.class)
 @Category(UsesSchema.class)
 public class ConvertTest {
   @Rule public final transient TestPipeline pipeline = TestPipeline.create();

sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ApproximateQuantilesTest.java

@@ -33,6 +33,7 @@
 import org.apache.beam.sdk.coders.BigEndianIntegerCoder;
 import org.apache.beam.sdk.coders.KvCoder;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
+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;
@@ -50,14 +51,13 @@
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 import org.junit.runners.Parameterized;
 
 /** Tests for {@link ApproximateQuantiles}. */
 public class ApproximateQuantilesTest {
 
   /** Tests for the overall combiner behavior. */
-  @RunWith(JUnit4.class)
+  @RunWith(BeamParallelJunit4Runner.class)
   public static class CombinerTests {
     static final List<KV<String, Integer>> TABLE =
         Arrays.asList(