[CI] Stabilize PostCommit Java ValidatesRunner Dataflow Streaming workflow

.github/workflows/beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming.yml

@@ -79,7 +79,7 @@ jobs:
       - name: run validatesRunnerStreaming script
         uses: ./.github/actions/gradle-command-self-hosted-action
         with:
-          gradle-command: :runners:google-cloud-dataflow-java:validatesRunnerStreaming
+          gradle-command: :runners:google-cloud-dataflow-java:validatesRunnerStreamingEngine
           max-workers: 12
       - name: Archive JUnit Test Results
         uses: actions/upload-artifact@v7

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

@@ -232,9 +232,9 @@ def createLegacyWorkerValidatesRunnerTest = { Map args ->
 
     systemProperty "beamTestPipelineOptions", JsonOutput.toJson(pipelineOptions)
 
-    // 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
+    // By default throttle parallelism to 4 to avoid GHA and GCP quota exhaustion.
+    // Can be overridden via -PmaxParallelForks=N.
+    maxParallelForks project.findProperty('maxParallelForks') ? (project.findProperty('maxParallelForks') as Integer) : 4
     classpath = configurations.validatesRunner
     testClassesDirs = files(project(":sdks:java:core").sourceSets.test.output.classesDirs) +
       files(project(project.path).sourceSets.test.output.classesDirs)
@@ -265,9 +265,9 @@ def createRunnerV2ValidatesRunnerTest = { Map args ->
     dependsOn buildAndPushDockerJavaContainer
     systemProperty "beamTestPipelineOptions", JsonOutput.toJson(pipelineOptions)
 
-    // 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
+    // By default throttle parallelism to 4 to avoid GHA and GCP quota exhaustion.
+    // Can be overridden via -PmaxParallelForks=N.
+    maxParallelForks project.findProperty('maxParallelForks') ? (project.findProperty('maxParallelForks') as Integer) : 4
     classpath = configurations.validatesRunner
     testClassesDirs = files(project(":sdks:java:core").sourceSets.test.output.classesDirs) +
       files(project(project.path).sourceSets.test.output.classesDirs)

runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/TestDataflowRunner.java

@@ -121,19 +121,30 @@ DataflowPipelineJob run(Pipeline pipeline, DataflowRunner runner) {
     ErrorMonitorMessagesHandler messageHandler =
         new ErrorMonitorMessagesHandler(job, new MonitoringUtil.LoggingHandler());
 
+    java.util.concurrent.atomic.AtomicReference<Optional<Boolean>> assertionsPassedRef =
+        new java.util.concurrent.atomic.AtomicReference<>(Optional.absent());
+
     if (options.isStreaming()) {
       if (options.isBlockOnRun()) {
-        jobSuccess = waitForStreamingJobTermination(job, messageHandler);
+        jobSuccess = waitForStreamingJobTermination(job, messageHandler, assertionsPassedRef);
       } else {
         jobSuccess = true;
       }
-      // No metrics in streaming
-      allAssertionsPassed = Optional.absent();
+      allAssertionsPassed = assertionsPassedRef.get();
+      if (!allAssertionsPassed.isPresent()) {
+        allAssertionsPassed = checkForPAssertSuccess(job);
+      }
     } else {
       jobSuccess = waitForBatchJobTermination(job, messageHandler);
       allAssertionsPassed = checkForPAssertSuccess(job);
     }
 
+    if (allAssertionsPassed.isPresent() && allAssertionsPassed.get()) {
+      if (job.getState() != State.FAILED && !messageHandler.hasSeenError()) {
+        jobSuccess = true;
+      }
+    }
+
     // If there is a certain assertion failure, throw the most precise exception we can.
     // There are situations where the metric will not be available, but as long as we recover
     // the actionable message from the logs it is acceptable.
@@ -160,11 +171,16 @@ DataflowPipelineJob run(Pipeline pipeline, DataflowRunner runner) {
   @SuppressWarnings("FutureReturnValueIgnored") // Job status checked via job.waitUntilFinish
   @SuppressFBWarnings("RV_RETURN_VALUE_IGNORED_BAD_PRACTICE")
   private boolean waitForStreamingJobTermination(
-      final DataflowPipelineJob job, ErrorMonitorMessagesHandler messageHandler) {
+      final DataflowPipelineJob job,
+      ErrorMonitorMessagesHandler messageHandler,
+      java.util.concurrent.atomic.AtomicReference<Optional<Boolean>> assertionsPassedRef) {
     // In streaming, there are infinite retries, so rather than timeout
     // we try to terminate early by polling and canceling if we see
-    // an error message
-    options.getExecutorService().submit(new CancelOnError(job, messageHandler));
+    // an error message or when all assertions have succeeded
+    java.util.concurrent.Future<Void> monitorFuture =
+        options
+            .getExecutorService()
+            .submit(new CancelOnError(job, messageHandler, this, assertionsPassedRef));
 
     // Whether we canceled or not, this gets the final state of the job or times out
     State finalState;
@@ -177,6 +193,8 @@ private boolean waitForStreamingJobTermination(
     } catch (InterruptedException e) {
       Thread.interrupted();
       return false;
+    } finally {
+      monitorFuture.cancel(true);
     }
 
     // Getting the final state may have timed out; it may not indicate a failure.
@@ -373,29 +391,91 @@ private static class CancelOnError implements Callable<Void> {
 
     private final DataflowPipelineJob job;
     private final ErrorMonitorMessagesHandler messageHandler;
-
-    public CancelOnError(DataflowPipelineJob job, ErrorMonitorMessagesHandler messageHandler) {
+    private final TestDataflowRunner runner;
+    private final java.util.concurrent.atomic.AtomicReference<Optional<Boolean>>
+        assertionsPassedRef;
+
+    public CancelOnError(
+        DataflowPipelineJob job,
+        ErrorMonitorMessagesHandler messageHandler,
+        TestDataflowRunner runner,
+        java.util.concurrent.atomic.AtomicReference<Optional<Boolean>> assertionsPassedRef) {
       this.job = job;
       this.messageHandler = messageHandler;
+      this.runner = runner;
+      this.assertionsPassedRef = assertionsPassedRef;
     }
 
     @Override
     public Void call() throws Exception {
+      int checkMetricsIntervalSteps = 5; // Check metrics every 15 seconds (5 * 3s)
+      int steps = 0;
+      boolean cancellationPending = false;
       while (true) {
-        State jobState = job.getState();
-
-        // If we see an error, cancel and note failure
-        if (messageHandler.hasSeenError() && !job.getState().isTerminal()) {
-          job.cancel();
-          LOG.info("Cancelling Dataflow job {}", job.getJobId());
-          return null;
-        }
-
-        if (jobState.isTerminal()) {
-          return null;
+        try {
+          State jobState = job.getState();
+
+          if (jobState.isTerminal()) {
+            return null;
+          }
+
+          // Check if we should initiate cancellation based on metrics (only if assertion state is
+          // not yet known)
+          if (!assertionsPassedRef.get().isPresent() && !cancellationPending) {
+            if (runner.expectedNumberOfAssertions > 0 && steps > 0 && steps % checkMetricsIntervalSteps == 0) {
+              try {
+                Optional<Boolean> assertionsPassed = runner.checkForPAssertSuccess(job);
+                if (assertionsPassed.isPresent()) {
+                  assertionsPassedRef.set(assertionsPassed);
+                  cancellationPending = true;
+                  if (assertionsPassed.get()) {
+                    LOG.info(
+                        "All assertions passed for streaming job {}, cancelling job.",
+                        job.getJobId());
+                  } else {
+                    LOG.info(
+                        "Found failed assertion for streaming job {}, cancelling job.",
+                        job.getJobId());
+                  }
+                }
+              } catch (Exception e) {
+                LOG.warn("Transient error polling metrics for job {}", job.getJobId(), e);
+              }
+            }
+          }
+
+          // Check if we should initiate cancellation based on error logs (only if not already
+          // cancellationPending)
+          if (!cancellationPending) {
+            long runningTimeMillis = System.currentTimeMillis() - startTimeMillis;
+            if (messageHandler.hasSeenError()
+                && (runningTimeMillis > 300000L || runner.expectedNumberOfAssertions == 0)) {
+              LOG.info(
+                  "Cancelling Dataflow job due to error messages seen: {}",
+                  messageHandler.getErrorMessage());
+              cancellationPending = true;
+            }
+          }
+
+          // Perform or retry cancellation if cancellation is pending
+          if (cancellationPending) {
+            try {
+              job.cancel();
+              return null; // Successful cancellation
+            } catch (Exception e) {
+              LOG.warn(
+                  "Failed to cancel Dataflow job {}. Will retry on next iteration. Error",
+                  job.getJobId(),
+                  e);
+            }
+          }
+
+        } catch (Exception e) {
+          LOG.warn("Exception in streaming job monitor loop for job {}", job.getJobId(), e);
         }
 
         Thread.sleep(3000L);
+        steps++;
       }
     }
   }

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

@@ -23,6 +23,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
@@ -211,16 +212,9 @@ public void testBatchPipelineFailsIfException() throws Exception {
     when(mockClient.getJobMetrics(anyString()))
         .thenReturn(generateMockMetricResponse(false /* success */, true /* tentative */));
     TestDataflowRunner runner = TestDataflowRunner.fromOptionsAndClient(options, mockClient);
-    try {
-      runner.run(p, mockRunner);
-    } catch (AssertionError expected) {
-      assertThat(expected.getMessage(), containsString("FooException"));
-      verify(mockJob, never()).cancel();
-      return;
-    }
-    // Note that fail throws an AssertionError which is why it is placed out here
-    // instead of inside the try-catch block.
-    fail("AssertionError expected");
+    AssertionError expected = assertThrows(AssertionError.class, () -> runner.run(p, mockRunner));
+    assertThat(expected.getMessage(), containsString("FooException"));
+    verify(mockJob, never()).cancel();
   }
 
   /** A streaming job that terminates with no error messages is a success. */
@@ -414,7 +408,7 @@ public void testStreamingPipelineFailsIfException() throws Exception {
     when(mockRunner.run(any(Pipeline.class))).thenReturn(mockJob);
 
     when(mockClient.getJobMetrics(anyString()))
-        .thenReturn(generateMockMetricResponse(false /* success */, true /* tentative */));
+        .thenReturn(buildJobMetrics(java.util.Collections.emptyList()));
     TestDataflowRunner runner = TestDataflowRunner.fromOptionsAndClient(options, mockClient);
 
     expectedException.expect(RuntimeException.class);
@@ -568,14 +562,9 @@ public void testBatchOnSuccessMatcherWhenPipelineFails() throws Exception {
 
     when(mockClient.getJobMetrics(anyString()))
         .thenReturn(generateMockMetricResponse(false /* success */, true /* tentative */));
-    try {
-      runner.run(p, mockRunner);
-    } catch (AssertionError expected) {
-      verify(mockJob, Mockito.times(1))
-          .waitUntilFinish(any(Duration.class), any(JobMessagesHandler.class));
-      return;
-    }
-    fail("Expected an exception on pipeline failure.");
+    assertThrows(AssertionError.class, () -> runner.run(p, mockRunner));
+    verify(mockJob, Mockito.times(1))
+        .waitUntilFinish(any(Duration.class), any(JobMessagesHandler.class));
   }
 
   /**
@@ -609,6 +598,93 @@ public void testStreamingOnSuccessMatcherWhenPipelineFails() throws Exception {
     // If the onSuccessMatcher were invoked, it would have crashed here with AssertionError
   }
 
+  @Test
+  public void testRunStreamingJobEarlySuccess() throws Exception {
+    options.setStreaming(true);
+    Pipeline p = TestPipeline.create(options);
+    PCollection<Integer> pc = p.apply(Create.of(1, 2, 3));
+    PAssert.that(pc).containsInAnyOrder(1, 2, 3);
+
+    DataflowPipelineJob mockJob = Mockito.mock(DataflowPipelineJob.class);
+    when(mockJob.getState()).thenReturn(State.RUNNING);
+    java.util.concurrent.CountDownLatch cancelLatch = new java.util.concurrent.CountDownLatch(1);
+    try {
+      Mockito.doAnswer(
+              invocation -> {
+                cancelLatch.countDown();
+                return null;
+              })
+          .when(mockJob)
+          .cancel();
+    } catch (Exception e) {
+      // Ignore
+    }
+    when(mockJob.waitUntilFinish(any(Duration.class), any(JobMessagesHandler.class)))
+        .thenAnswer(
+            invocation -> {
+              if (!cancelLatch.await(10, java.util.concurrent.TimeUnit.SECONDS)) {
+                throw new RuntimeException("Timeout waiting for job cancellation");
+              }
+              return State.CANCELLED;
+            });
+    when(mockJob.getProjectId()).thenReturn("test-project");
+    when(mockJob.getJobId()).thenReturn("test-job");
+
+    DataflowRunner mockRunner = Mockito.mock(DataflowRunner.class);
+    when(mockRunner.run(any(Pipeline.class))).thenReturn(mockJob);
+
+    when(mockClient.getJobMetrics(anyString()))
+        .thenReturn(generateMockMetricResponse(true /* success */, true /* tentative */));
+    TestDataflowRunner runner = TestDataflowRunner.fromOptionsAndClient(options, mockClient);
+    runner.run(p, mockRunner);
+
+    Mockito.verify(mockJob, Mockito.timeout(5000)).cancel();
+  }
+
+  @Test
+  public void testRunStreamingJobEarlyFailure() throws Exception {
+    options.setStreaming(true);
+    Pipeline p = TestPipeline.create(options);
+    PCollection<Integer> pc = p.apply(Create.of(1, 2, 3));
+    PAssert.that(pc).containsInAnyOrder(1, 2, 3);
+
+    DataflowPipelineJob mockJob = Mockito.mock(DataflowPipelineJob.class);
+    when(mockJob.getState()).thenReturn(State.RUNNING);
+    java.util.concurrent.CountDownLatch cancelLatch = new java.util.concurrent.CountDownLatch(1);
+    try {
+      Mockito.doAnswer(
+              invocation -> {
+                cancelLatch.countDown();
+                return null;
+              })
+          .when(mockJob)
+          .cancel();
+    } catch (Exception e) {
+      // Ignore
+    }
+    when(mockJob.waitUntilFinish(any(Duration.class), any(JobMessagesHandler.class)))
+        .thenAnswer(
+            invocation -> {
+              if (!cancelLatch.await(10, java.util.concurrent.TimeUnit.SECONDS)) {
+                throw new RuntimeException("Timeout waiting for job cancellation");
+              }
+              return State.CANCELLED;
+            });
+    when(mockJob.getProjectId()).thenReturn("test-project");
+    when(mockJob.getJobId()).thenReturn("test-job");
+
+    DataflowRunner mockRunner = Mockito.mock(DataflowRunner.class);
+    when(mockRunner.run(any(Pipeline.class))).thenReturn(mockJob);
+
+    when(mockClient.getJobMetrics(anyString()))
+        .thenReturn(generateMockMetricResponse(false /* success */, true /* tentative */));
+    TestDataflowRunner runner = TestDataflowRunner.fromOptionsAndClient(options, mockClient);
+
+    assertThrows(AssertionError.class, () -> runner.run(p, mockRunner));
+
+    Mockito.verify(mockJob, Mockito.timeout(5000)).cancel();
+  }
+
   static class TestSuccessMatcher extends BaseMatcher<PipelineResult>
       implements SerializableMatcher<PipelineResult> {
     private final transient DataflowPipelineJob mockJob;

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

@@ -291,23 +291,23 @@ public void testMatchWatchForNewFiles() throws IOException, InterruptedException
                 .filepattern(watchPath.resolve("*").toString())
                 .continuously(
                     Duration.millis(100),
-                    Watch.Growth.afterTimeSinceNewOutput(Duration.standardSeconds(1))));
+                    Watch.Growth.afterTimeSinceNewOutput(Duration.standardSeconds(3))));
     PCollection<MatchResult.Metadata> matchAllMetadata =
         p.apply("create for matchAll new files", Create.of(watchPath.resolve("*").toString()))
             .apply(
                 "match filename through matchAll",
                 FileIO.matchAll()
                     .continuously(
                         Duration.millis(100),
-                        Watch.Growth.afterTimeSinceNewOutput(Duration.standardSeconds(1))));
+                        Watch.Growth.afterTimeSinceNewOutput(Duration.standardSeconds(3))));
     PCollection<MatchResult.Metadata> matchUpdatedMetadata =
         p.apply(
             "match updated",
             FileIO.match()
                 .filepattern(watchPath.resolve("first").toString())
                 .continuously(
                     Duration.millis(100),
-                    Watch.Growth.afterTimeSinceNewOutput(Duration.standardSeconds(1)),
+                    Watch.Growth.afterTimeSinceNewOutput(Duration.standardSeconds(3)),
                     true));
     PCollection<MatchResult.Metadata> matchAllUpdatedMetadata =
         p.apply("create for matchAll updated files", Create.of(watchPath.resolve("*").toString()))
@@ -316,7 +316,7 @@ public void testMatchWatchForNewFiles() throws IOException, InterruptedException
                 FileIO.matchAll()
                     .continuously(
                         Duration.millis(100),
-                        Watch.Growth.afterTimeSinceNewOutput(Duration.standardSeconds(1)),
+                        Watch.Growth.afterTimeSinceNewOutput(Duration.standardSeconds(3)),
                         true));
 
     // write one file at the beginning. This will trigger the first output for matchAll