feat(runner): add new completionCheckInterval property

related-to: kestra-io/kestra#3701
kestra-io · May 15, 2024 · ee38e2f · ee38e2f
1 parent e2fea2a
commit ee38e2f
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 4 deletions.
diff --git a/src/main/java/io/kestra/plugin/gcp/runner/GcpBatchTaskRunner.java b/src/main/java/io/kestra/plugin/gcp/runner/GcpBatchTaskRunner.java
@@ -179,6 +179,13 @@ public class GcpBatchTaskRunner extends TaskRunner implements GcpInterface, Remo
     @PluginProperty
     private final Boolean resume = true;
 
+    @Schema(
+        title = "The frequency with which the TaskRunner checks whether the job is completed."
+    )
+    @Builder.Default
+    @PluginProperty
+    private final Duration completionCheckInterval = Duration.ofSeconds(1);
+
     @Override
     public RunnerResult run(RunContext runContext, TaskCommands taskCommands, List<String> filesToUpload, List<String> filesToDownload) throws Exception {
         String renderedBucket = runContext.render(this.bucket);
@@ -336,7 +343,7 @@ public RunnerResult run(RunContext runContext, TaskCommands taskCommands, List<S
             LogEntryServerStream stream = logging.tailLogEntries(Logging.TailOption.filter(logFilter));
             try (LogTail ignored = new LogTail(stream, taskCommands.getLogConsumer())) {
                 // Wait for the job termination
-                result = waitFormTerminated(batchServiceClient, result, waitDuration);
+                result = waitForTerminated(batchServiceClient, result, waitDuration);
                 if (result == null) {
                     throw new TimeoutException();
                 }
@@ -420,7 +427,7 @@ private Runnable.Container mainContainer(RunContext runContext, TaskCommands tas
         return builder.build();
     }
 
-    private Job waitFormTerminated(BatchServiceClient batchServiceClient, Job result, Duration waitDuration) throws TimeoutException {
+    private Job waitForTerminated(BatchServiceClient batchServiceClient, Job result, Duration waitDuration) throws TimeoutException {
         return Await.until(
             () -> {
                 Job terminated = batchServiceClient.getJob(result.getName());
@@ -429,7 +436,7 @@ private Job waitFormTerminated(BatchServiceClient batchServiceClient, Job result
                 }
                 return null;
             },
-            Duration.ofMillis(500),
+            completionCheckInterval,
             waitDuration
         );
     }

diff --git a/src/main/java/io/kestra/plugin/gcp/runner/GcpCloudRunTaskRunner.java b/src/main/java/io/kestra/plugin/gcp/runner/GcpCloudRunTaskRunner.java
@@ -171,6 +171,13 @@ public class GcpCloudRunTaskRunner extends TaskRunner implements GcpInterface, R
     @PluginProperty
     private final Boolean delete = true;
 
+    @Schema(
+        title = "The frequency with which the TaskRunner checks whether the job is completed."
+    )
+    @Builder.Default
+    @PluginProperty
+    private final Duration completionCheckInterval = Duration.ofSeconds(1);
+
     @Override
     public RunnerResult run(RunContext runContext, TaskCommands taskCommands, List<String> filesToUpload, List<String> filesToDownload) throws Exception {
 
@@ -394,7 +401,7 @@ private Execution awaitJobExecutionTermination(final ExecutionsClient executions
                 }
                 return null;
             },
-            Duration.ofMillis(500),
+            completionCheckInterval,
             timeout
         );
     }

diff --git a/src/test/java/io/kestra/plugin/gcp/runner/GcpBatchTaskRunnerTest.java b/src/test/java/io/kestra/plugin/gcp/runner/GcpBatchTaskRunnerTest.java
@@ -5,6 +5,7 @@
 import io.micronaut.context.annotation.Value;
 import org.junit.jupiter.api.Disabled;
 
+import java.time.Duration;
 import java.util.List;
 
 @Disabled("Need complex CI setup still needed to be done")
@@ -30,6 +31,7 @@ protected TaskRunner taskRunner() {
             .bucket(bucket)
             .networkInterfaces(List.of(GcpBatchTaskRunner.NetworkInterface.builder().network(network).subnetwork(subnetwork).build()))
             .delete(false)
+            .completionCheckInterval(Duration.ofMillis(100))
             .build();
     }
 }
diff --git a/src/test/java/io/kestra/plugin/gcp/runner/GcpCloudRunTaskRunnerTest.java b/src/test/java/io/kestra/plugin/gcp/runner/GcpCloudRunTaskRunnerTest.java
@@ -5,6 +5,7 @@
 import io.micronaut.context.annotation.Value;
 import org.junit.jupiter.api.Disabled;
 
+import java.time.Duration;
 import java.util.List;
 
 @Disabled("Need complex CI setup still needed to be done")
@@ -23,6 +24,7 @@ protected TaskRunner taskRunner() {
             .region("us-central1")
             .bucket(bucket)
             .delete(false)
+            .completionCheckInterval(Duration.ofMillis(100))
             .build();
     }
 }