feat(core): handle timeout on Pause task (#1040)

close #1039
kestra-io · Mar 3, 2023 · c8d32f0 · c8d32f0
1 parent 10c5d65
commit c8d32f0
Show file tree

Hide file tree

Showing 10 changed files with 80 additions and 20 deletions.
diff --git a/core/src/main/java/io/kestra/core/runners/ExecutionDelay.java b/core/src/main/java/io/kestra/core/runners/ExecutionDelay.java
@@ -1,5 +1,6 @@
 package io.kestra.core.runners;
 
+import io.kestra.core.models.flows.State;
 import lombok.AllArgsConstructor;
 import lombok.Builder;
 import lombok.Value;
@@ -19,4 +20,6 @@ public class ExecutionDelay {
 
     @NotNull
     Instant date;
+
+    @NotNull State.Type state;
 }
diff --git a/core/src/main/java/io/kestra/core/runners/ExecutorService.java b/core/src/main/java/io/kestra/core/runners/ExecutorService.java
@@ -386,11 +386,12 @@ private Executor handlePausedDelay(Executor executor, List<WorkerTaskResult> wor
                 if (task instanceof Pause) {
                     Pause pauseTask = (Pause) task;
 
-                    if (pauseTask.getDelay() != null) {
+                    if (pauseTask.getDelay() != null || pauseTask.getTimeout() != null) {
                         return ExecutionDelay.builder()
                             .taskRunId(workerTaskResult.getTaskRun().getId())
                             .executionId(executor.getExecution().getId())
-                            .date(workerTaskResult.getTaskRun().getState().maxDate().plus(pauseTask.getDelay()))
+                            .date(workerTaskResult.getTaskRun().getState().maxDate().plus(pauseTask.getDelay() != null ? pauseTask.getDelay() : pauseTask.getTimeout()))
+                            .state(pauseTask.getDelay() != null ? State.Type.RUNNING : State.Type.FAILED)
                             .build();
                     }
                 }

diff --git a/core/src/main/java/io/kestra/core/tasks/flows/Pause.java b/core/src/main/java/io/kestra/core/tasks/flows/Pause.java
@@ -58,11 +58,18 @@
 public class Pause extends Sequential implements FlowableTask<VoidOutput> {
     @Schema(
         title = "Duration of the pause.",
-        description = "If null, a manual approval is need, if not, the delay before automatically continue the execution"
+        description = "If null and no timeout, a manual approval is needed, if not, the delay before continuing the execution"
     )
     @PluginProperty
     private Duration delay;
 
+    @Schema(
+        title = "Timeout of the pause.",
+        description = "If null and no delay, a manual approval is needed, else a manual approval is needed before the timeout or the task will fail"
+    )
+    @PluginProperty
+    private Duration timeout;
+
     @Override
     public List<NextTaskRun> resolveNexts(RunContext runContext, Execution execution, TaskRun parentTaskRun) throws IllegalVariableEvaluationException {
         if (this.needPause(parentTaskRun) || parentTaskRun.getState().getCurrent() == State.Type.PAUSED) {

diff --git a/core/src/test/java/io/kestra/core/Helpers.java b/core/src/test/java/io/kestra/core/Helpers.java
@@ -19,7 +19,7 @@
 import java.util.function.Consumer;
 
 public class Helpers {
-    public static long FLOWS_COUNT = 52;
+    public static long FLOWS_COUNT = 53;
 
     public static ApplicationContext applicationContext() throws URISyntaxException {
         return applicationContext(

diff --git a/core/src/test/java/io/kestra/core/tasks/flows/PauseTest.java b/core/src/test/java/io/kestra/core/tasks/flows/PauseTest.java
@@ -31,10 +31,15 @@ void run() throws Exception {
     }
 
     @Test
-    void failed() throws Exception {
+    void delay() throws Exception {
         suite.runDelay(runnerUtils);
     }
 
+    @Test
+    void timeout() throws Exception {
+        suite.runTimeout(runnerUtils);
+    }
+
     @Singleton
     public static class Suite {
         @Inject
@@ -87,6 +92,25 @@ public void runDelay(RunnerUtils runnerUtils) throws Exception {
             assertThat(execution.getTaskRunList().get(0).getState().getHistories().stream().filter(history -> history.getState() == State.Type.RUNNING).count(), is(2L));
             assertThat(execution.getTaskRunList(), hasSize(3));
         }
-    }
 
+
+        public void runTimeout(RunnerUtils runnerUtils) throws Exception {
+            Execution execution = runnerUtils.runOne("io.kestra.tests", "pause-timeout");
+
+            assertThat(execution.getState().getCurrent(), is(State.Type.PAUSED));
+            assertThat(execution.getTaskRunList().get(0).getState().getCurrent(), is(State.Type.PAUSED));
+            assertThat(execution.getTaskRunList(), hasSize(1));
+
+            execution = runnerUtils.awaitExecution(
+                e -> e.getState().getCurrent() == State.Type.FAILED,
+                () -> {},
+                Duration.ofSeconds(30)
+            );
+
+            assertThat(execution.getTaskRunList().get(0).getState().getHistories().stream().filter(history -> history.getState() == State.Type.PAUSED).count(), is(1L));
+            assertThat(execution.getTaskRunList().get(0).getState().getHistories().stream().filter(history -> history.getState() == State.Type.RUNNING).count(), is(1L));
+            assertThat(execution.getTaskRunList().get(0).getState().getHistories().stream().filter(history -> history.getState() == State.Type.FAILED).count(), is(1L));
+            assertThat(execution.getTaskRunList(), hasSize(1));
+        }
+    }
 }
diff --git a/core/src/test/resources/flows/valids/pause-timeout.yaml b/core/src/test/resources/flows/valids/pause-timeout.yaml
@@ -0,0 +1,15 @@
+id: pause-timeout
+namespace: io.kestra.tests
+
+tasks:
+  - id: pause
+    type: io.kestra.core.tasks.flows.Pause
+    timeout: PT1S
+    tasks:
+      - id: ko
+        type: io.kestra.core.tasks.scripts.Bash
+        commands:
+          - echo "trigger 1 seconds pause"
+  - id: last
+    type: io.kestra.core.tasks.debugs.Return
+    format: "{{task.id}} > {{taskrun.startDate}}"
diff --git a/jdbc/src/main/java/io/kestra/jdbc/runner/JdbcExecutor.java b/jdbc/src/main/java/io/kestra/jdbc/runner/JdbcExecutor.java
@@ -443,13 +443,16 @@ private void executionDelaySend() {
                 Executor executor = new Executor(pair.getLeft(), null);
 
                 try {
-                    Execution markAsExecution = executionService.markAs(
-                        pair.getKey(),
-                        executionDelay.getTaskRunId(),
-                        State.Type.RUNNING
-                    );
+                    if (executor.getExecution().findTaskRunByTaskRunId(executionDelay.getTaskRunId()).getState().getCurrent() == State.Type.PAUSED) {
+
+                        Execution markAsExecution = executionService.markAs(
+                            pair.getKey(),
+                            executionDelay.getTaskRunId(),
+                            executionDelay.getState()
+                        );
 
-                    executor = executor.withExecution(markAsExecution, "pausedRestart");
+                        executor = executor.withExecution(markAsExecution, "pausedRestart");
+                    }
                 } catch (Exception e) {
                     executor = handleFailedExecutionFromExecutor(executor, e);
                 }

diff --git a/jdbc/src/test/java/io/kestra/jdbc/runner/JdbcRunnerTest.java b/jdbc/src/test/java/io/kestra/jdbc/runner/JdbcRunnerTest.java
@@ -227,4 +227,9 @@ public void pauseRun() throws Exception {
     public void pauseRunDelay() throws Exception {
         pauseTest.runDelay(runnerUtils);
     }
+
+    @Test
+    public void pauseRunTimeout() throws Exception {
+        pauseTest.runTimeout(runnerUtils);
+    }
 }
diff --git a/runner-memory/src/main/java/io/kestra/runner/memory/MemoryExecutor.java b/runner-memory/src/main/java/io/kestra/runner/memory/MemoryExecutor.java
@@ -184,13 +184,15 @@ private void handleExecution(ExecutionState state) {
                                 try {
                                     ExecutionState executionState = EXECUTIONS.get(workerTaskResultDelay.getExecutionId());
 
-                                    Execution markAsExecution = executionService.markAs(
-                                        executionState.execution,
-                                        workerTaskResultDelay.getTaskRunId(),
-                                        State.Type.RUNNING
-                                    );
-
-                                    executionQueue.emit(markAsExecution);
+                                    if (executionState.execution.findTaskRunByTaskRunId(workerTaskResultDelay.getTaskRunId()).getState().getCurrent() == State.Type.PAUSED) {
+                                        Execution markAsExecution = executionService.markAs(
+                                            executionState.execution,
+                                            workerTaskResultDelay.getTaskRunId(),
+                                            workerTaskResultDelay.getState()
+                                        );
+
+                                        executionQueue.emit(markAsExecution);
+                                    }
                                 } catch (Exception e) {
                                     throw new RuntimeException(e);
                                 }

diff --git a/webserver/src/test/java/io/kestra/webserver/controllers/FlowControllerTest.java b/webserver/src/test/java/io/kestra/webserver/controllers/FlowControllerTest.java
@@ -481,7 +481,7 @@ void exportByQuery() throws IOException {
         Files.write(file.toPath(), zip);
 
         try (ZipFile zipFile = new ZipFile(file)) {
-            assertThat(zipFile.stream().count(), is(52L));
+            assertThat(zipFile.stream().count(), is(Helpers.FLOWS_COUNT));
         }
 
         file.delete();