Repository: kestra-io/kestra (java)

Issue:
[JDBC Executor] The WorkerGroup key is not taken into account when resubmitting a task/trigger on worker failure
https://github.com/kestra-io/kestra/blob/467861652055b514ded47e6dea6ab756a95b3fbc/jdbc/src/main/java/io/kestra/jdbc/runner/JdbcExecutor.java#L370-L404

Pull request:
fix(jdbc): resubmit worker job to the good worker group
Fixes #7730

Patch:
diff --git a/core/src/main/java/io/kestra/core/runners/Worker.java b/core/src/main/java/io/kestra/core/runners/Worker.java
index fe6cfae1c..d249a627d 100644
--- a/core/src/main/java/io/kestra/core/runners/Worker.java
+++ b/core/src/main/java/io/kestra/core/runners/Worker.java
@@ -1019,6 +1019,17 @@ public class Worker implements Service, Runnable, AutoCloseable {
 
     @VisibleForTesting
     public void shutdown() {
+        // initiate shutdown
+        shutdown.compareAndSet(false, true);
+
+        try {
+            // close the WorkerJob queue to stop receiving new JobTask execution.
+            workerJobQueue.close();
+        } catch (IOException e) {
+            log.error("Failed to close the WorkerJobQueue");
+        }
+
+        // close all queues and shutdown now
         this.receiveCancellations.forEach(Runnable::run);
         this.executorService.shutdownNow();
     }
diff --git a/jdbc-mysql/build.gradle b/jdbc-mysql/build.gradle
index dc5089558..c6c27a7fe 100644
--- a/jdbc-mysql/build.gradle
+++ b/jdbc-mysql/build.gradle
@@ -16,4 +16,5 @@ dependencies {
     testImplementation project(':jdbc').sourceSets.test.output
     testImplementation project(':storage-local')
     testImplementation project(':tests')
+    testImplementation("io.micronaut.validation:micronaut-validation") // MysqlServiceLivenessCoordinatorTest fail to init without that
 }
diff --git a/jdbc-postgres/build.gradle b/jdbc-postgres/build.gradle
index 6d0f10168..ab63b1010 100644
--- a/jdbc-postgres/build.gradle
+++ b/jdbc-postgres/build.gradle
@@ -16,4 +16,5 @@ dependencies {
     testImplementation project(':jdbc').sourceSets.test.output
     testImplementation project(':storage-local')
     testImplementation project(':tests')
+    testImplementation("io.micronaut.validation:micronaut-validation") // PostgresServiceLivenessCoordinatorTest fail to init without that
 }
diff --git a/jdbc/src/main/java/io/kestra/jdbc/repository/AbstractJdbcWorkerJobRunningRepository.java b/jdbc/src/main/java/io/kestra/jdbc/repository/AbstractJdbcWorkerJobRunningRepository.java
index f35f6c697..d4ad26e3e 100644
--- a/jdbc/src/main/java/io/kestra/jdbc/repository/AbstractJdbcWorkerJobRunningRepository.java
+++ b/jdbc/src/main/java/io/kestra/jdbc/repository/AbstractJdbcWorkerJobRunningRepository.java
@@ -1,5 +1,6 @@
 package io.kestra.jdbc.repository;
 
+import com.google.common.annotations.VisibleForTesting;
 import io.kestra.core.repositories.WorkerJobRunningRepositoryInterface;
 import io.kestra.core.runners.WorkerJobRunning;
 import lombok.extern.slf4j.Slf4j;
@@ -47,6 +48,20 @@ public abstract class AbstractJdbcWorkerJobRunningRepository extends AbstractJdb
             });
     }
 
+    @VisibleForTesting
+    public List<WorkerJobRunning> findAll() {
+        return this.jdbcRepository
+            .getDslContextWrapper()
+            .transactionResult(configuration -> {
+                var select = DSL
+                    .using(configuration)
+                    .select((field("value")))
+                    .from(this.jdbcRepository.getTable());
+
+                return this.jdbcRepository.fetch(select);
+            });
+    }
+
     public List<WorkerJobRunning> getWorkerJobWithWorkerDead(DSLContext context, List<String> workersToDelete) {
         return context
                 .select(field("value"))
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 1aa44bae3..fc2162b32 100644
--- a/jdbc/src/main/java/io/kestra/jdbc/runner/JdbcExecutor.java
+++ b/jdbc/src/main/java/io/kestra/jdbc/runner/JdbcExecutor.java
@@ -379,7 +379,7 @@ public class JdbcExecutor implements ExecutorInterface, Service {
                         workerJobRunningRepository.deleteByKey(workerTaskRunning.uid());
                     } else {
                         try {
-                            workerJobQueue.emit(WorkerTask.builder()
+                            workerJobQueue.emit(workerTaskRunning.getWorkerInstance().workerGroup(), WorkerTask.builder()
                                 .taskRun(workerTaskRunning.getTaskRun().onRunningResend())
                                 .task(workerTaskRunning.getTask())
                                 .runContext(workerTaskRunning.getRunContext())
@@ -406,7 +406,7 @@ public class JdbcExecutor implements ExecutorInterface, Service {
                 // WorkerTriggerRunning
                 if (workerJobRunning instanceof WorkerTriggerRunning workerTriggerRunning) {
                     try {
-                        workerJobQueue.emit(WorkerTrigger.builder()
+                        workerJobQueue.emit(workerTriggerRunning.getWorkerInstance().workerGroup(), WorkerTrigger.builder()
                             .trigger(workerTriggerRunning.getTrigger())
                             .conditionContext(workerTriggerRunning.getConditionContext())
                             .triggerContext(workerTriggerRunning.getTriggerContext())


Test patch:
diff --git a/jdbc/src/test/java/io/kestra/jdbc/runner/JdbcServiceLivenessCoordinatorTest.java b/jdbc/src/test/java/io/kestra/jdbc/runner/JdbcServiceLivenessCoordinatorTest.java
index fa7f6a167..6ac03b177 100644
--- a/jdbc/src/test/java/io/kestra/jdbc/runner/JdbcServiceLivenessCoordinatorTest.java
+++ b/jdbc/src/test/java/io/kestra/jdbc/runner/JdbcServiceLivenessCoordinatorTest.java
@@ -8,29 +8,36 @@ import io.kestra.core.models.executions.TaskRun;
 import io.kestra.core.models.flows.Flow;
 import io.kestra.core.models.flows.State.Type;
 import io.kestra.core.models.tasks.ResolvedTask;
+import io.kestra.core.models.tasks.WorkerGroup;
 import io.kestra.core.models.triggers.Trigger;
 import io.kestra.core.queues.QueueFactoryInterface;
 import io.kestra.core.queues.QueueInterface;
-import io.kestra.core.repositories.LocalFlowRepositoryLoader;
 import io.kestra.core.runners.*;
 import io.kestra.core.services.SkipExecutionService;
+import io.kestra.core.services.WorkerGroupService;
 import io.kestra.core.tasks.test.SleepTrigger;
-import io.kestra.core.utils.Await;
 import io.kestra.core.utils.IdUtils;
 import io.kestra.core.utils.TestsUtils;
+import io.kestra.core.validations.WorkerGroupValidation;
+import io.kestra.core.validations.validator.WorkerGroupValidator;
 import io.kestra.jdbc.JdbcTestUtils;
+import io.kestra.jdbc.repository.AbstractJdbcWorkerJobRunningRepository;
 import io.kestra.plugin.core.flow.Sleep;
 import io.micronaut.context.ApplicationContext;
 import io.micronaut.context.annotation.Property;
+import io.micronaut.core.annotation.AnnotationValue;
+import io.micronaut.core.annotation.NonNull;
+import io.micronaut.core.annotation.Nullable;
+import io.micronaut.test.annotation.MockBean;
+import io.micronaut.validation.validator.constraints.ConstraintValidatorContext;
 import jakarta.inject.Inject;
 import jakarta.inject.Named;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestInstance;
 import reactor.core.publisher.Flux;
 
-import java.io.IOException;
-import java.net.URISyntaxException;
 import java.time.Duration;
 import java.util.Collections;
 import java.util.List;
@@ -40,6 +47,7 @@ import java.util.concurrent.TimeUnit;
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.*;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
 @KestraTest(environments =  {"test", "liveness"})
@@ -49,42 +57,58 @@ public abstract class JdbcServiceLivenessCoordinatorTest {
     @Inject
     private StandAloneRunner runner;
 
-    @Inject
-    private LocalFlowRepositoryLoader repositoryLoader;
-
     @Inject
     private ApplicationContext applicationContext;
 
     @Inject
-    JdbcTestUtils jdbcTestUtils;
+    private JdbcTestUtils jdbcTestUtils;
 
     @Inject
-    RunContextFactory runContextFactory;
+    private RunContextFactory runContextFactory;
 
     @Inject
     @Named(QueueFactoryInterface.WORKERJOB_NAMED)
-    QueueInterface<WorkerJob> workerJobQueue;
+    private QueueInterface<WorkerJob> workerJobQueue;
 
     @Inject
     @Named(QueueFactoryInterface.WORKERTASKRESULT_NAMED)
-    QueueInterface<WorkerTaskResult> workerTaskResultQueue;
+    private QueueInterface<WorkerTaskResult> workerTaskResultQueue;
 
     @Inject
     @Named(QueueFactoryInterface.WORKERTRIGGERRESULT_NAMED)
-    QueueInterface<WorkerTriggerResult> workerTriggerResultQueue;
+    private QueueInterface<WorkerTriggerResult> workerTriggerResultQueue;
+
+    @Inject
+    @Named(QueueFactoryInterface.TRIGGER_NAMED)
+    private QueueInterface<Trigger> triggerQueue;
+
+    @Inject
+    private JdbcServiceLivenessCoordinator jdbcServiceLivenessHandler;
 
     @Inject
-    JdbcServiceLivenessCoordinator jdbcServiceLivenessHandler;
+    private SkipExecutionService skipExecutionService;
 
     @Inject
-    SkipExecutionService skipExecutionService;
+    private AbstractJdbcWorkerJobRunningRepository workerJobRunningRepository;
 
     @BeforeAll
-    void init() throws IOException, URISyntaxException {
+    void init() {
         jdbcTestUtils.drop();
         jdbcTestUtils.migrate();
+
         // Simulate that executor and workers are not running on the same JVM.
         jdbcServiceLivenessHandler.setServerInstance(IdUtils.create());
+
+        // start the runner
+        runner.setSchedulerEnabled(false);
+        runner.setWorkerEnabled(false);
+        runner.run();
+    }
+
+    @AfterEach
+    void tearDown() {
+        List<WorkerJobRunning> workerJobRunnings = workerJobRunningRepository.findAll();
+        workerJobRunnings.forEach(workerJobRunning -> workerJobRunningRepository.deleteByKey(workerJobRunning.uid()));
     }
 
     @Test
@@ -96,9 +120,6 @@ public abstract class JdbcServiceLivenessCoordinatorTest {
         Worker worker = applicationContext.createBean(Worker.class, IdUtils.create(), 1, null);
         worker.run();
 
-        runner.setSchedulerEnabled(false);
-        runner.setWorkerEnabled(false);
-        runner.run();
         Flux<WorkerTaskResult> receive = TestsUtils.receive(workerTaskResultQueue, either -> {
             if (either.getLeft().getTaskRun().getState().getCurrent() == Type.SUCCESS) {
                 resubmitLatch.countDown();
@@ -109,7 +130,7 @@ public abstract class JdbcServiceLivenessCoordinatorTest {
             }
         });
 
-        workerJobQueue.emit(workerTask(Duration.ofSeconds(10)));
+        workerJobQueue.emit(workerTask(Duration.ofSeconds(5)));
         boolean runningLatchAwait = runningLatch.await(5, TimeUnit.SECONDS);
         assertThat(runningLatchAwait, is(true));
         worker.shutdown(); // stop processing task
@@ -117,7 +138,43 @@ public abstract class JdbcServiceLivenessCoordinatorTest {
         // create second worker (this will revoke previously one).
         Worker newWorker = applicationContext.createBean(Worker.class, IdUtils.create(), 1, null);
         newWorker.run();
-        boolean resubmitLatchAwait = resubmitLatch.await(30, TimeUnit.SECONDS);
+        boolean resubmitLatchAwait = resubmitLatch.await(10, TimeUnit.SECONDS);
+        assertThat(resubmitLatchAwait, is(true));
+        WorkerTaskResult workerTaskResult = receive.blockLast();
+        assertThat(workerTaskResult, notNullValue());
+        assertThat(workerTaskResult.getTaskRun().getState().getCurrent(), is(Type.SUCCESS));
+        assertThat(workerTaskResult.getTaskRun().getAttempts(), hasSize(2));
+        newWorker.shutdown();
+    }
+
+    @Test
+    void shouldReEmitTasksToTheSameWorkerGroup() throws Exception {
+        CountDownLatch runningLatch = new CountDownLatch(1);
+        CountDownLatch resubmitLatch = new CountDownLatch(1);
+
+        // create first worker
+        Worker worker = applicationContext.createBean(Worker.class, IdUtils.create(), 1, "workerGroupKey");
+        worker.run();
+
+        Flux<WorkerTaskResult> receive = TestsUtils.receive(workerTaskResultQueue, either -> {
+            if (either.getLeft().getTaskRun().getState().getCurrent() == Type.SUCCESS) {
+                resubmitLatch.countDown();
+            }
+
+            if (either.getLeft().getTaskRun().getState().getCurrent() == Type.RUNNING) {
+                runningLatch.countDown();
+            }
+        });
+
+        workerJobQueue.emit("workerGroupKey", workerTask(Duration.ofSeconds(5), "workerGroupKey"));
+        boolean runningLatchAwait = runningLatch.await(5, TimeUnit.SECONDS);
+        assertThat(runningLatchAwait, is(true));
+        worker.shutdown(); // stop processing task
+
+        // create second worker (this will revoke previously one).
+        Worker newWorker = applicationContext.createBean(Worker.class, IdUtils.create(), 1, "workerGroupKey");
+        newWorker.run();
+        boolean resubmitLatchAwait = resubmitLatch.await(10, TimeUnit.SECONDS);
         assertThat(resubmitLatchAwait, is(true));
         WorkerTaskResult workerTaskResult = receive.blockLast();
         assertThat(workerTaskResult, notNullValue());
@@ -132,10 +189,8 @@ public abstract class JdbcServiceLivenessCoordinatorTest {
 
         Worker worker = applicationContext.createBean(Worker.class, IdUtils.create(), 8, null);
         worker.run();
-        runner.setSchedulerEnabled(false);
-        runner.setWorkerEnabled(false);
-        runner.run();
-        WorkerTask workerTask = workerTask(Duration.ofSeconds(10));
+
+        WorkerTask workerTask = workerTask(Duration.ofSeconds(5));
         skipExecutionService.setSkipExecutions(List.of(workerTask.getTaskRun().getExecutionId()));
 
         Flux<WorkerTaskResult> receive = TestsUtils.receive(workerTaskResultQueue, either -> {
@@ -168,40 +223,103 @@ public abstract class JdbcServiceLivenessCoordinatorTest {
     void shouldReEmitTriggerWhenWorkerIsDetectedAsNonResponding() throws Exception {
         Worker worker = applicationContext.createBean(Worker.class, IdUtils.create(), 1, null);
         worker.run();
-        runner.setSchedulerEnabled(false);
-        runner.setWorkerEnabled(false);
-        runner.run();
 
         WorkerTrigger workerTrigger = workerTrigger(Duration.ofSeconds(5));
 
         // 2 trigger should happen because of the resubmit
         CountDownLatch countDownLatch = new CountDownLatch(2);
-        Flux<WorkerJob> receive = TestsUtils.receive(workerJobQueue, workerJob -> countDownLatch.countDown());
+        Flux<WorkerTriggerResult> receive = TestsUtils.receive(workerTriggerResultQueue, workerTriggerResult -> countDownLatch.countDown());
 
+        // we wait that the worker receive the trigger
+        CountDownLatch triggerCountDownLatch = new CountDownLatch(1);
+        Flux<Trigger> receiveTrigger = TestsUtils.receive(triggerQueue, either -> {
+            if (either.getLeft().getWorkerId().equals(worker.getId())) {
+                triggerCountDownLatch.countDown();
+            }
+        });
         workerJobQueue.emit(workerTrigger);
-        Await.until(() -> worker.getEvaluateTriggerRunningCount()
-                .get(workerTrigger.getTriggerContext().uid()) != null,
-            Duration.ofMillis(100),
-            Duration.ofSeconds(5)
-        );
+        assertTrue(triggerCountDownLatch.await(10, TimeUnit.SECONDS));
+        receiveTrigger.blockLast();
         worker.shutdown();
 
         Worker newWorker = applicationContext.createBean(Worker.class, IdUtils.create(), 1, null);
         applicationContext.registerSingleton(newWorker);
         newWorker.run();
 
-        boolean lastAwait = countDownLatch.await(15, TimeUnit.SECONDS);
+        boolean lastAwait = countDownLatch.await(10, TimeUnit.SECONDS);
 
         newWorker.shutdown();
         receive.blockLast();
         assertThat(lastAwait, is(true));
     }
 
+    @Test
+    void shouldReEmitTriggerToTheSameWorkerGroup() throws Exception {
+        Worker worker = applicationContext.createBean(Worker.class, IdUtils.create(), 1, "workerGroupKey");
+        worker.run();
+
+        WorkerTrigger workerTrigger = workerTrigger(Duration.ofSeconds(5), "workerGroupKey");
+
+        // 2 trigger should happen because of the resubmit
+        CountDownLatch countDownLatch = new CountDownLatch(2);
+        Flux<WorkerTriggerResult> receive = TestsUtils.receive(workerTriggerResultQueue, workerTriggerResult -> countDownLatch.countDown());
+
+        // we wait that the worker receive the trigger
+        CountDownLatch triggerCountDownLatch = new CountDownLatch(1);
+        Flux<Trigger> receiveTrigger = TestsUtils.receive(triggerQueue, either -> {
+            if (either.getLeft().getWorkerId().equals(worker.getId())) {
+                triggerCountDownLatch.countDown();
+            }
+        });
+        workerJobQueue.emit("workerGroupKey", workerTrigger);
+        assertTrue(triggerCountDownLatch.await(10, TimeUnit.SECONDS));
+        receiveTrigger.blockLast();
+        worker.shutdown();
+
+        Worker newWorker = applicationContext.createBean(Worker.class, IdUtils.create(), 1, "workerGroupKey");
+        applicationContext.registerSingleton(newWorker);
+        newWorker.run();
+
+        boolean lastAwait = countDownLatch.await(10, TimeUnit.SECONDS);
+
+        newWorker.shutdown();
+        receive.blockLast();
+        assertThat(lastAwait, is(true));
+    }
+
+    @MockBean(WorkerGroupService.class)
+    WorkerGroupService workerGroupService() {
+        return new WorkerGroupService() {
+            @Override
+            public String resolveGroupFromKey(String workerGroupKey) {
+                return workerGroupKey;
+            }
+        };
+    }
+
+    @MockBean(WorkerGroupValidator.class)
+    WorkerGroupValidator workerGroupValidator() {
+        return new WorkerGroupValidator() {
+            @Override
+            public boolean isValid(
+                @Nullable WorkerGroup value,
+                @NonNull AnnotationValue<WorkerGroupValidation> annotationMetadata,
+                @NonNull ConstraintValidatorContext context) {
+                return true;
+            }
+        };
+    }
+
     private WorkerTask workerTask(Duration sleep) {
+        return workerTask(sleep, null);
+    }
+
+    private WorkerTask workerTask(Duration sleep, String workerGroupKey) {
         Sleep bash = Sleep.builder()
             .type(Sleep.class.getName())
             .id("unit-test")
             .duration(sleep)
+            .workerGroup(workerGroupKey != null ? new WorkerGroup(workerGroupKey, null) : null)
             .build();
 
         Execution execution = TestsUtils.mockExecution(flowBuilder(sleep), ImmutableMap.of());
@@ -216,10 +334,15 @@ public abstract class JdbcServiceLivenessCoordinatorTest {
     }
 
     private WorkerTrigger workerTrigger(Duration sleep) {
+        return workerTrigger(sleep, null);
+    }
+
+    private WorkerTrigger workerTrigger(Duration sleep, String workerGroupKey) {
         SleepTrigger trigger = SleepTrigger.builder()
             .type(SleepTrigger.class.getName())
             .id("unit-test")
             .duration(sleep.toMillis())
+            .workerGroup(workerGroupKey != null ? new WorkerGroup(workerGroupKey, null) : null)
             .build();
 
         Map.Entry<ConditionContext, Trigger> mockedTrigger = TestsUtils.mockTrigger(runContextFactory, trigger);