Repository: kestra-io/kestra (java)

Issue:
Inconsistency in after SLA violation execution state
### Describe the issue

Since the killing is asynchronous, the state after kill in case of SLA is random between the SLA specified state and `KILLED`.

### Environment

- Kestra Version: develop


Pull request:
fix(system): force a state after kill
closes #8936


Patch:
diff --git a/core/src/main/java/io/kestra/core/models/executions/ExecutionKilledExecution.java b/core/src/main/java/io/kestra/core/models/executions/ExecutionKilledExecution.java
index bd6daca31..05afaed76 100644
--- a/core/src/main/java/io/kestra/core/models/executions/ExecutionKilledExecution.java
+++ b/core/src/main/java/io/kestra/core/models/executions/ExecutionKilledExecution.java
@@ -2,7 +2,6 @@ package io.kestra.core.models.executions;
 
 import com.fasterxml.jackson.annotation.JsonInclude;
 import io.kestra.core.models.TenantInterface;
-import io.kestra.core.models.triggers.TriggerContext;
 import io.kestra.core.runners.WorkerTask;
 import jakarta.validation.constraints.NotNull;
 import lombok.*;
@@ -25,6 +24,11 @@ public class ExecutionKilledExecution extends ExecutionKilled implements TenantI
     @NotNull
     String executionId;
 
+    /**
+     * The state to move the execution to after kill.
+     */
+    io.kestra.core.models.flows.State.Type executionState;
+
     /**
      * Specifies whether killing the execution, also kill all sub-flow executions.
      */
diff --git a/core/src/main/java/io/kestra/core/runners/ExecutorService.java b/core/src/main/java/io/kestra/core/runners/ExecutorService.java
index 066c86e50..72b8ec7d6 100644
--- a/core/src/main/java/io/kestra/core/runners/ExecutorService.java
+++ b/core/src/main/java/io/kestra/core/runners/ExecutorService.java
@@ -14,14 +14,14 @@ import io.kestra.core.queues.QueueException;
 import io.kestra.core.queues.QueueFactoryInterface;
 import io.kestra.core.queues.QueueInterface;
 import io.kestra.core.services.*;
-import io.kestra.core.test.flow.TaskFixture;
 import io.kestra.core.storages.StorageContext;
+import io.kestra.core.test.flow.TaskFixture;
 import io.kestra.core.trace.propagation.RunContextTextMapSetter;
 import io.kestra.core.utils.ListUtils;
 import io.kestra.core.utils.TruthUtils;
+import io.kestra.plugin.core.flow.LoopUntil;
 import io.kestra.plugin.core.flow.Pause;
 import io.kestra.plugin.core.flow.Subflow;
-import io.kestra.plugin.core.flow.LoopUntil;
 import io.kestra.plugin.core.flow.WorkingDirectory;
 import io.micronaut.context.ApplicationContext;
 import io.opentelemetry.api.OpenTelemetry;
@@ -1263,6 +1263,7 @@ public class ExecutorService {
         killQueue.emit(ExecutionKilledExecution
             .builder()
             .state(ExecutionKilled.State.REQUESTED)
+            .executionState(state)
             .executionId(execution.getId())
             .isOnKillCascade(false) // TODO we may offer the choice to the user here
             .tenantId(execution.getTenantId())
diff --git a/core/src/main/java/io/kestra/core/services/ExecutionService.java b/core/src/main/java/io/kestra/core/services/ExecutionService.java
index d1028bd24..921403182 100644
--- a/core/src/main/java/io/kestra/core/services/ExecutionService.java
+++ b/core/src/main/java/io/kestra/core/services/ExecutionService.java
@@ -35,8 +35,8 @@ import io.kestra.core.utils.ListUtils;
 import io.kestra.plugin.core.flow.Pause;
 import io.kestra.plugin.core.flow.WorkingDirectory;
 import io.micronaut.context.event.ApplicationEventPublisher;
-import io.micronaut.core.annotation.Nullable;
 import io.micronaut.http.multipart.CompletedPart;
+import jakarta.annotation.Nullable;
 import jakarta.inject.Inject;
 import jakarta.inject.Named;
 import jakarta.inject.Singleton;
@@ -656,29 +656,35 @@ public class ExecutionService {
      *
      * @return the execution in a KILLING state if not already terminated
      */
-    public Execution kill(Execution execution, FlowInterface flow) {
-        if (execution.getState().getCurrent() == State.Type.KILLING || execution.getState().isTerminated()) {
+    public Execution kill(Execution execution, FlowInterface flow, Optional<State.Type> afterKillState) {
+        // We afford the double kill potential (KILLING & afterKillState != null) to ensure we put the afterKillState
+        if ((execution.getState().getCurrent() == State.Type.KILLING && afterKillState.isEmpty()) || execution.getState().isTerminated()) {
             return execution;
         }
 
         Execution newExecution;
+        State.Type killingOrAfterKillState = afterKillState.orElse(State.Type.KILLING);
         if (execution.getState().isPaused()) {
             // Must be resumed and killed, no need to send killing event to the worker as the execution is not executing anything in it.
             // An edge case can exist where the execution is resumed automatically before we resume it with a killing.
             try {
                 newExecution = this.resume(execution, flow, State.Type.KILLING);
+                newExecution = newExecution.withState(afterKillState.orElse(newExecution.getState().getCurrent()));
             } catch (Exception e) {
                 // if we cannot resume, we set it anyway to killing, so we don't throw
                 log.warn("Unable to resume a paused execution before killing it", e);
-                newExecution = execution.withState(State.Type.KILLING);
+                newExecution = execution.withState(killingOrAfterKillState);
             }
         } else {
-            newExecution = execution.withState(State.Type.KILLING);
+            newExecution = execution.withState(killingOrAfterKillState);
         }
 
         eventPublisher.publishEvent(new CrudEvent<>(newExecution, execution, CrudEventType.UPDATE));
         return newExecution;
     }
+    public Execution kill(Execution execution, FlowInterface flow) {
+        return this.kill(execution, flow, Optional.empty());
+    }
 
     /**
      * Climb up the hierarchy of parent taskruns and kill them all.
diff --git a/jdbc/src/main/java/io/kestra/jdbc/runner/JdbcExecutor.java b/jdbc/src/main/java/io/kestra/jdbc/runner/JdbcExecutor.java
index ae1c89f78..425e03b29 100644
--- a/jdbc/src/main/java/io/kestra/jdbc/runner/JdbcExecutor.java
+++ b/jdbc/src/main/java/io/kestra/jdbc/runner/JdbcExecutor.java
@@ -7,7 +7,6 @@ import io.kestra.core.exceptions.InternalException;
 import io.kestra.core.metrics.MetricRegistry;
 import io.kestra.core.models.executions.*;
 import io.kestra.core.models.executions.statistics.ExecutionCount;
-import io.kestra.core.models.flows.Flow;
 import io.kestra.core.models.flows.*;
 import io.kestra.core.models.flows.sla.*;
 import io.kestra.core.models.tasks.ExecutableTask;
@@ -19,9 +18,9 @@ import io.kestra.core.queues.QueueException;
 import io.kestra.core.queues.QueueFactoryInterface;
 import io.kestra.core.queues.QueueInterface;
 import io.kestra.core.repositories.TriggerRepositoryInterface;
+import io.kestra.core.runners.*;
 import io.kestra.core.runners.Executor;
 import io.kestra.core.runners.ExecutorService;
-import io.kestra.core.runners.*;
 import io.kestra.core.schedulers.SchedulerTriggerStateInterface;
 import io.kestra.core.server.ClusterEvent;
 import io.kestra.core.server.Service;
@@ -938,7 +937,7 @@ public class JdbcExecutor implements ExecutorInterface, Service {
             log.error("Unable to kill the execution {}", killedExecution.getExecutionId(), e);
         }
 
-        Executor executor = mayTransitExecutionToKillingStateAndGet(killedExecution.getExecutionId());
+        Executor executor = killingOrAfterKillState(killedExecution.getExecutionId(), Optional.ofNullable(killedExecution.getExecutionState()));
 
         // Check whether kill event should be propagated to downstream executions.
         // By default, always propagate the ExecutionKill to sub-flows (for backward compatibility).
@@ -964,12 +963,12 @@ public class JdbcExecutor implements ExecutorInterface, Service {
         }
     }
 
-    private Executor mayTransitExecutionToKillingStateAndGet(final String executionId) {
+    private Executor killingOrAfterKillState(final String executionId, Optional<State.Type> afterKillState) {
         return executionRepository.lock(executionId, pair -> {
             Execution currentExecution = pair.getLeft();
             FlowInterface flow = flowMetaStore.findByExecution(currentExecution).orElseThrow();
 
-            Execution killing = executionService.kill(currentExecution, flow);
+            Execution killing = executionService.kill(currentExecution, flow, afterKillState);
             Executor current = new Executor(currentExecution, null)
                 .withExecution(killing, "joinKillingExecution");
             return Pair.of(current, pair.getRight());


Test patch:
diff --git a/core/src/test/java/io/kestra/core/runners/ExecutionServiceTest.java b/core/src/test/java/io/kestra/core/runners/ExecutionServiceTest.java
index fd28d74a1..cf603a432 100644
--- a/core/src/test/java/io/kestra/core/runners/ExecutionServiceTest.java
+++ b/core/src/test/java/io/kestra/core/runners/ExecutionServiceTest.java
@@ -8,8 +8,8 @@ import io.kestra.core.models.Label;
 import io.kestra.core.models.executions.Execution;
 import io.kestra.core.models.flows.Flow;
 import io.kestra.core.models.flows.FlowWithSource;
-import io.kestra.core.models.flows.State;
 import io.kestra.core.models.flows.GenericFlow;
+import io.kestra.core.models.flows.State;
 import io.kestra.core.models.property.Property;
 import io.kestra.core.repositories.ExecutionRepositoryInterface;
 import io.kestra.core.repositories.FlowRepositoryInterface;
@@ -32,8 +32,8 @@ import java.util.concurrent.TimeoutException;
 import static io.kestra.core.tenant.TenantService.MAIN_TENANT;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.hamcrest.MatcherAssert.assertThat;
-import static org.hamcrest.Matchers.not;
 import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.not;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 @Slf4j
@@ -428,4 +428,17 @@ class ExecutionServiceTest {
         Execution markedAs = executionService.markAs(execution, flow, execution.getTaskRunList().getFirst().getId(), State.Type.SUCCESS);
         assertThat(markedAs.getState().getCurrent()).isEqualTo(State.Type.RESTARTED);
     }
+
+    @Test
+    @LoadFlows({"flows/valids/pause_no_tasks.yaml"})
+    void killToState() throws Exception {
+        Execution execution = runnerUtils.runOneUntilPaused(MAIN_TENANT, "io.kestra.tests", "pause_no_tasks");
+        Flow flow = flowRepository.findByExecution(execution);
+
+        Execution killed = executionService.kill(execution, flow, Optional.of(State.Type.CANCELLED));
+
+        assertThat(killed.getState().getCurrent()).isEqualTo(State.Type.CANCELLED);
+        assertThat(killed.findTaskRunsByTaskId("pause").getFirst().getState().getCurrent()).isEqualTo(State.Type.KILLED);
+        assertThat(killed.getState().getHistories()).hasSize(5);
+    }
 }
\ No newline at end of file