[release-2.77] Cherry-pick #40259: Speed up Beam presubmits

.github/workflows/beam_PreCommit_Java_GCP_IO_Direct.yml

@@ -82,9 +82,26 @@ jobs:
   beam_PreCommit_Java_GCP_IO_Direct:
     name: ${{ matrix.job_name }} (${{ matrix.job_phrase }})
     strategy:
+      fail-fast: false
       matrix:
-        job_name: ["beam_PreCommit_Java_GCP_IO_Direct"]
-        job_phrase: ["Run Java_GCP_IO_Direct PreCommit"]
+        include:
+          - job_name: "beam_PreCommit_Java_GCP_IO_Direct"
+            job_phrase: "Run Java_GCP_IO_Direct PreCommit"
+            gradle_command: |
+              :sdks:java:io:google-cloud-platform:build \
+              :sdks:java:io:google-cloud-platform:expansion-service:build \
+            extra_arguments: ""
+            is_build_job: true
+          - job_name: "beam_PreCommit_Java_GCP_IO_Direct_BigQuery"
+            job_phrase: "Run Java_GCP_IO_Direct PreCommit"
+            gradle_command: ":sdks:java:io:google-cloud-platform:postCommitBigQuery"
+            extra_arguments: "-x test"
+            is_build_job: false
+          - job_name: "beam_PreCommit_Java_GCP_IO_Direct_OtherGCP"
+            job_phrase: "Run Java_GCP_IO_Direct PreCommit"
+            gradle_command: ":sdks:java:io:google-cloud-platform:postCommitOtherGcp"
+            extra_arguments: "-x test"
+            is_build_job: false
     timeout-minutes: 180
     if: |
       github.event_name == 'push' ||
@@ -108,19 +125,16 @@ jobs:
       - name: run PreCommit Java GCP IO Direct script
         uses: ./.github/actions/gradle-command-self-hosted-action
         with:
-          gradle-command: |
-            :sdks:java:io:google-cloud-platform:build \
-            :sdks:java:io:google-cloud-platform:expansion-service:build \
-            :sdks:java:io:google-cloud-platform:postCommit \
+          gradle-command: ${{ matrix.gradle_command }}
           arguments: |
             -PdisableSpotlessCheck=true \
             -PdisableCheckStyle=true \
-            -PenableJacocoReport \
+            ${{ matrix.extra_arguments }}
       - name: Archive JUnit Test Results
         uses: actions/upload-artifact@v7
         if: ${{ !success() }}
         with:
-          name: JUnit Test Results
+          name: JUnit Test Results (${{ matrix.job_name }})
           path: "**/build/reports/tests/"
       - name: Publish JUnit Test Results
         uses: EnricoMi/publish-unit-test-result-action@v2
@@ -132,27 +146,13 @@ jobs:
           large_files: true
       - name: Archive SpotBugs Results
         uses: actions/upload-artifact@v7
-        if: always()
+        if: ${{ always() && matrix.is_build_job }}
         with:
           name: SpotBugs Results
           path: '**/build/reports/spotbugs/*.html'
       - name: Publish SpotBugs Results
         uses: jwgmeligmeyling/spotbugs-github-action@v1.2
-        if: always()
+        if: ${{ always() && matrix.is_build_job }}
         with:
           name: Publish SpotBugs
           path: '**/build/reports/spotbugs/*.html'
-      - name: Set Jacoco report path
-        id: jacoco_report_path
-        run: |
-          echo "path=sdks/java/io/google-cloud-platform/build/reports/jacoco/test/jacocoTestReport.xml" >> $GITHUB_OUTPUT
-      - name: Archive Jacoco Results
-        uses: actions/upload-artifact@v7
-        with:
-          name: Jacoco Results
-          path: '**/build/jacoco/report/**'
-      - name: Upload coverage reports to Codecov
-        uses: codecov/codecov-action@v3
-        with:
-          file: ${{ steps.jacoco_report_path.outputs.path }}
-          flags: java

sdks/java/io/google-cloud-platform/build.gradle

@@ -33,6 +33,14 @@ ext.summary = "IO library to read and write Google Cloud Platform systems from B
 tasks.withType(Test).configureEach {
   // Open java.nio for Apache Arrow direct memory access under Java 17+
   jvmArgs '--add-opens=java.base/java.nio=ALL-UNNAMED'
+  maxParallelForks 4
+  jacoco.enabled = false
+}
+
+gradle.taskGraph.whenReady {
+  tasks.withType(Test).configureEach {
+    jacoco.enabled = false
+  }
 }
 
 dependencies {
@@ -328,62 +336,86 @@ task bigQueryEarlyRolloutIntegrationTest(type: Test, dependsOn: processTestResou
   testClassesDirs = sourceSets.test.output.classesDirs
 }
 
-// path(s) for Cloud Spanner related classes
-def spannerIncludes = [
-        '**/org/apache/beam/sdk/io/gcp/spanner/**',
-]
-
-// exclude auto-generated classes and integration tests
-def jacocoExcludes = [
-        '**/AutoValue_*',
-        '**/*IT*',
-]
-
-task spannerCodeCoverageReport(type: JacocoReport, dependsOn: test) {
-  group = "Reporting"
-  description = "Generates code coverage report for Cloud Spanner related classes"
-  classDirectories.setFrom(files(files(project.sourceSets.main.output).collect {
-            project.fileTree(
-                    dir: it,
-                    includes: spannerIncludes,
-                    excludes: jacocoExcludes)
-  }))
-  sourceDirectories.setFrom(files(project.sourceSets.main.allSource.srcDirs))
-  executionData.setFrom(file("${buildDir}/jacoco/test.exec"))
-  reports {
-    html.getRequired().set(true)
-    html.getOutputLocation().set(file("${buildDir}/reports/jacoco/spanner/"))
+task integrationTestBigQuery(type: Test, dependsOn: processTestResources) {
+  group = "Verification"
+  systemProperty "beamTestPipelineOptions", integrationTest.systemProperties["beamTestPipelineOptions"]
+  outputs.upToDateWhen { false }
+  include '**/bigquery/**/*IT.class'
+  exclude '**/BigQueryIOIcebergManagedTableIT.class'
+  exclude '**/BigQueryIOReadIT.class'
+  exclude '**/BigQueryIOStorageQueryIT.class'
+  exclude '**/BigQueryIOStorageReadIT.class'
+  exclude '**/BigQueryIOStorageWriteIT.class'
+  exclude '**/BigQueryIODynamicQueryIT.class'
+  exclude '**/BigQueryIODynamicReadIT.class'
+  exclude '**/BigQueryToTableIT.class'
+  maxParallelForks 4
+  classpath = sourceSets.test.runtimeClasspath
+  testClassesDirs = sourceSets.test.output.classesDirs
+  useJUnit {
+    excludeCategories "org.apache.beam.sdk.testing.UsesKms"
+  }
+}
+
+task integrationTestBigQueryKms(type: Test, dependsOn: processTestResources) {
+  group = "Verification"
+  systemProperty "beamTestPipelineOptions", integrationTestKms.systemProperties["beamTestPipelineOptions"]
+  outputs.upToDateWhen { false }
+  include '**/bigquery/**/*IT.class'
+  maxParallelForks 4
+  classpath = sourceSets.test.runtimeClasspath
+  testClassesDirs = sourceSets.test.output.classesDirs
+  useJUnit {
+    includeCategories "org.apache.beam.sdk.testing.UsesKms"
   }
 }
 
-task spannerCodeCoverageVerification(type: JacocoCoverageVerification, dependsOn: spannerCodeCoverageReport) {
+task integrationTestOtherGcp(type: Test, dependsOn: processTestResources) {
   group = "Verification"
-  description = "Enforces code coverage verification for Cloud Spanner related classes"
-  classDirectories.setFrom(files(files(project.sourceSets.main.output).collect {
-    project.fileTree(
-            dir: it,
-            includes: spannerIncludes,
-            excludes: jacocoExcludes)
-  }))
-  sourceDirectories.setFrom(files(project.sourceSets.main.allSource.srcDirs))
-  executionData.setFrom(file("${buildDir}/jacoco/test.exec"))
-  violationRules {
-    failOnViolation = true
-    rule {
-      element = 'BUNDLE'
-
-      limit {
-        value = 'COVEREDRATIO'
-        counter = 'INSTRUCTION'
-        minimum = 0.60
-      }
+  systemProperty "beamTestPipelineOptions", integrationTest.systemProperties["beamTestPipelineOptions"]
+  outputs.upToDateWhen { false }
+  include '**/*IT.class'
+  exclude '**/bigquery/**/*IT.class'
+  maxParallelForks 4
+  classpath = sourceSets.test.runtimeClasspath
+  testClassesDirs = sourceSets.test.output.classesDirs
+  useJUnit {
+    excludeCategories "org.apache.beam.sdk.testing.UsesKms"
+    filter {
+      // https://github.com/apache/beam/issues/32071
+      excludeTestsMatching 'org.apache.beam.sdk.io.gcp.bigtable.BigtableReadIT.testE2EBigtableSegmentRead'
     }
   }
 }
 
-// make the check task depend on spannerCodeCoverageVerification so that the build
-// fails when the code coverage threshold is violated.
-project.check.dependsOn "spannerCodeCoverageVerification"
+task integrationTestOtherGcpKms(type: Test, dependsOn: processTestResources) {
+  group = "Verification"
+  systemProperty "beamTestPipelineOptions", integrationTestKms.systemProperties["beamTestPipelineOptions"]
+  outputs.upToDateWhen { false }
+  include '**/*IT.class'
+  exclude '**/bigquery/**/*IT.class'
+  maxParallelForks 4
+  classpath = sourceSets.test.runtimeClasspath
+  testClassesDirs = sourceSets.test.output.classesDirs
+  useJUnit {
+    includeCategories "org.apache.beam.sdk.testing.UsesKms"
+  }
+}
+
+task postCommitBigQuery {
+  group = "Verification"
+  description = "Integration tests of BigQuery connectors using the DirectRunner."
+  dependsOn integrationTestBigQuery
+  dependsOn integrationTestBigQueryKms
+  dependsOn bigQueryManagedIcebergIntegrationTest
+}
+
+task postCommitOtherGcp {
+  group = "Verification"
+  description = "Integration tests of non-BigQuery GCP connectors using the DirectRunner."
+  dependsOn integrationTestOtherGcp
+  dependsOn integrationTestOtherGcpKms
+}
 
 task postCommit {
   group = "Verification"

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteStorageApiBatchTest.java

@@ -0,0 +1,35 @@
+/*
+ * 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.io.gcp.bigquery;
+
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameters;
+
+/**
+ * Runs {@link BigQueryIOWriteTest} for Storage Write API batch mode in a separate test class so
+ * Gradle {@code maxParallelForks} can execute it in parallel with other modes.
+ */
+@RunWith(Parameterized.class)
+public class BigQueryIOWriteStorageApiBatchTest extends BigQueryIOWriteTest {
+  @Parameters(name = "useStorageApi={0}, useStorageApiApproximate={1}, useStreaming={2}")
+  public static Iterable<Object[]> data() {
+    return ImmutableList.of(new Object[] {true, false, false});
+  }
+}

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteStorageApiStreamTest.java

@@ -0,0 +1,35 @@
+/*
+ * 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.io.gcp.bigquery;
+
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameters;
+
+/**
+ * Runs {@link BigQueryIOWriteTest} for Storage Write API streaming modes in a separate test class
+ * so Gradle {@code maxParallelForks} can execute them in parallel with other modes.
+ */
+@RunWith(Parameterized.class)
+public class BigQueryIOWriteStorageApiStreamTest extends BigQueryIOWriteTest {
+  @Parameters(name = "useStorageApi={0}, useStorageApiApproximate={1}, useStreaming={2}")
+  public static Iterable<Object[]> data() {
+    return ImmutableList.of(new Object[] {true, false, true}, new Object[] {true, true, true});
+  }
+}

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java

@@ -228,14 +228,9 @@ public class BigQueryIOWriteTest implements Serializable {
   private transient TemporaryFolder testFolder = new TemporaryFolder();
   private transient TestPipeline p;
 
-  @Parameters
+  @Parameters(name = "useStorageApi={0}, useStorageApiApproximate={1}, useStreaming={2}")
   public static Iterable<Object[]> data() {
-    return ImmutableList.of(
-        new Object[] {false, false, false},
-        new Object[] {false, false, true},
-        new Object[] {true, false, false},
-        new Object[] {true, false, true},
-        new Object[] {true, true, true});
+    return ImmutableList.of(new Object[] {false, false, false}, new Object[] {false, false, true});
   }
 
   @Parameter(0)

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFnTest.java

@@ -103,8 +103,13 @@ public void setup() throws IOException {
 
   @Test
   public void testProcessElementAndGetSize() throws IOException, InterruptedException {
+    Instant now = Instant.now();
+    doFn =
+        new ReadChangeStreamPartitionDoFn(
+            daoFactory, actionFactory, metrics, Duration.ZERO, () -> now);
+    doFn.setSizeEstimator(sizeEstimator);
     long watermarkLag = 10;
-    Instant tenSecondsAgo = Instant.now().minus(Duration.standardSeconds(watermarkLag));
+    Instant tenSecondsAgo = now.minus(Duration.standardSeconds(watermarkLag));
     Range.ByteStringRange partitionRange = Range.ByteStringRange.create("", "");
     ChangeStreamContinuationToken testToken =
         ChangeStreamContinuationToken.create(partitionRange, "test");
@@ -115,7 +120,7 @@ public void testProcessElementAndGetSize() throws IOException, InterruptedExcept
             "uid-a",
             tenSecondsAgo,
             Collections.emptyList(),
-            Instant.now().plus(Duration.standardSeconds(60)));
+            now.plus(Duration.standardSeconds(60)));
     long mutationSize = 100L;
     when(sizeEstimator.sizeOf(any())).thenReturn(mutationSize);
     ReadChangeStreamPartitionProgressTracker restrictionTracker =
@@ -149,11 +154,10 @@ public void testProcessElementAndGetSize() throws IOException, InterruptedExcept
     doFn.processElement(partition, restrictionTracker, receiver, watermarkEstimator);
     double sizeEstimate =
         doFn.getSize(
-            new StreamProgress(
-                testToken, tenSecondsAgo, BigDecimal.valueOf(20), Instant.now(), false));
+            new StreamProgress(testToken, tenSecondsAgo, BigDecimal.valueOf(20), now, false));
     // we should have output 2 100B mutations in the past 10s
     long bytesPerSecond = (mutationSize * 2) / 10;
-    assertEquals(sizeEstimate, (double) (bytesPerSecond * watermarkLag), 10);
+    assertEquals((double) (bytesPerSecond * watermarkLag), sizeEstimate, 10);
     verify(receiver, times(2)).outputWithTimestamp(KV.of(rowKey, mockMutation), Instant.EPOCH);
   }