added waitForConditionAndForceCancel (#35465)

Author: liferoad

it/common/src/main/java/org/apache/beam/it/common/PipelineLauncher.java

@@ -427,6 +427,17 @@ List<JobMessage> listMessages(
    */
   Job drainJob(String project, String region, String jobId) throws IOException;
 
+  /**
+   * Force-cancels the given job.
+   *
+   * @param project the project that the job is running under
+   * @param region the region that the job was launched in
+   * @param jobId the id of the job to force-cancel
+   * @throws IOException if there is an issue sending the request
+   * @return Updated job instance
+   */
+  Job forceCancelJob(String project, String region, String jobId) throws IOException;
+
   /**
    * Get the specified metric of the given job.
    *

it/common/src/main/java/org/apache/beam/it/common/PipelineOperator.java

@@ -141,6 +141,14 @@ public Result waitForConditionAndCancel(Config config, Supplier<Boolean>... cond
     return waitForConditionAndExecute(config, conditionCheck, this::cancelJobAndFinish);
   }
 
+  /**
+   * Similar to {@link #waitForConditionAndCancel(Config, Supplier...)} but force-cancels the job.
+   */
+  public Result waitForConditionAndForceCancel(Config config, Supplier<Boolean>... conditionCheck)
+      throws IOException {
+    return waitForConditionAndExecute(config, conditionCheck, this::forceCancelJobAndFinish);
+  }
+
   private Result waitForConditionAndExecute(
       Config config,
       Supplier<Boolean>[] conditionCheck,
@@ -171,6 +179,12 @@ public Result cancelJobAndFinish(Config config) throws IOException {
     return waitUntilDone(config);
   }
 
+  /** Similar to {@link #cancelJobAndFinish(Config)} but force-cancels the job. */
+  public Result forceCancelJobAndFinish(Config config) throws IOException {
+    client.forceCancelJob(config.project(), config.region(), config.jobId());
+    return waitUntilDone(config);
+  }
+
   private static Result finishOrTimeout(
       Config config, Supplier<Boolean>[] conditionCheck, Supplier<Boolean>... stopChecking) {
     Instant start = Instant.now();

it/common/src/test/java/org/apache/beam/it/common/PipelineOperatorTest.java

@@ -199,4 +199,38 @@ public void testFinishAfterConditionTimeout() throws IOException {
     verify(client).drainJob(any(), any(), any());
     assertThat(result).isEqualTo(Result.TIMEOUT);
   }
+
+  @Test
+  public void testForceCancelAfterCondition() throws IOException {
+    // Arrange
+    AtomicInteger callCount = new AtomicInteger();
+    int totalCalls = 3;
+    Supplier<Boolean> checker = () -> callCount.incrementAndGet() >= totalCalls;
+
+    when(client.getJobStatus(any(), any(), any()))
+        .thenReturn(JobState.RUNNING)
+        .thenThrow(new IOException())
+        .thenReturn(JobState.RUNNING)
+        .thenReturn(JobState.CANCELLING)
+        .thenReturn(JobState.CANCELLED);
+
+    // Act
+    Result result =
+        new PipelineOperator(client).waitForConditionAndForceCancel(DEFAULT_CONFIG, checker);
+
+    // Assert
+    verify(client, atLeast(totalCalls))
+        .getJobStatus(projectCaptor.capture(), regionCaptor.capture(), jobIdCaptor.capture());
+    verify(client)
+        .forceCancelJob(projectCaptor.capture(), regionCaptor.capture(), jobIdCaptor.capture());
+
+    Set<String> allProjects = new HashSet<>(projectCaptor.getAllValues());
+    Set<String> allRegions = new HashSet<>(regionCaptor.getAllValues());
+    Set<String> allJobIds = new HashSet<>(jobIdCaptor.getAllValues());
+
+    assertThat(allProjects).containsExactly(PROJECT);
+    assertThat(allRegions).containsExactly(REGION);
+    assertThat(allJobIds).containsExactly(JOB_ID);
+    assertThat(result).isEqualTo(Result.CONDITION_MET);
+  }
 }

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/AbstractPipelineLauncher.java

@@ -140,6 +140,20 @@ public Job drainJob(String project, String region, String jobId) {
                 client.projects().locations().jobs().update(project, region, jobId, job).execute());
   }
 
+  @Override
+  public Job forceCancelJob(String project, String region, String jobId) throws IOException {
+    LOG.info("Force-cancelling {} under {}", jobId, project);
+    Job job =
+        new Job()
+            .setRequestedState(JobState.CANCELLED.toString())
+            .setLabels(ImmutableMap.of("force_cancel_job", "true"));
+    LOG.info("Sending job to update {}:\n{}", jobId, formatForLogging(job));
+    return Failsafe.with(clientRetryPolicy())
+        .get(
+            () ->
+                client.projects().locations().jobs().update(project, region, jobId, job).execute());
+  }
+
   @Override
   public @Nullable Double getMetric(String project, String region, String jobId, String metricName)
       throws IOException {

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/DirectRunnerClient.java

@@ -135,6 +135,12 @@ public Job drainJob(String project, String region, String jobId) {
     return cancelJob(project, region, jobId);
   }
 
+  @Override
+  public Job forceCancelJob(String project, String region, String jobId) throws IOException {
+    LOG.warn("Cannot force-cancel a direct runner job. Cancelling the job instead.");
+    return cancelJob(project, region, jobId);
+  }
+
   @Override
   public Double getMetric(String project, String region, String jobId, String metricName) {
     return null;