feat(): Zombie task resubmit

kestra-io · Sep 18, 2023 · 848d933 · 848d933
1 parent 201ee9b
commit 848d933
Show file tree

Hide file tree

Showing 22 changed files with 470 additions and 53 deletions.
diff --git a/cli/src/main/resources/application.yml b/cli/src/main/resources/application.yml
@@ -124,6 +124,9 @@ kestra:
       workerheartbeat:
         cls: io.kestra.core.runners.WorkerHeartbeat
         table: "worker_heartbeat"
+      workerjobrunning:
+        cls: io.kestra.core.runners.WorkerJobRunning
+        table: "worker_job_running"
 
     queues:
       min-poll-interval: 25ms
@@ -170,5 +173,5 @@ kestra:
     fixed-delay: 1h
 
   heartbeat:
-    frequency: 10
+    frequency: 5
     heartbeat-missed: 3
diff --git a/core/src/main/java/io/kestra/core/queues/QueueService.java b/core/src/main/java/io/kestra/core/queues/QueueService.java
@@ -1,10 +1,7 @@
 package io.kestra.core.queues;
 
 import io.kestra.core.models.Setting;
-import io.kestra.core.models.executions.Execution;
-import io.kestra.core.models.executions.ExecutionKilled;
-import io.kestra.core.models.executions.LogEntry;
-import io.kestra.core.models.executions.MetricEntry;
+import io.kestra.core.models.executions.*;
 import io.kestra.core.models.flows.Flow;
 import io.kestra.core.models.templates.Template;
 import io.kestra.core.models.topologies.FlowTopology;

diff --git a/core/src/main/java/io/kestra/core/repositories/WorkerJobRunningInterface.java b/core/src/main/java/io/kestra/core/repositories/WorkerJobRunningInterface.java
@@ -0,0 +1,16 @@
+package io.kestra.core.repositories;
+
+import io.kestra.core.runners.WorkerJobRunning;
+
+import java.util.List;
+import java.util.Optional;
+
+public interface WorkerJobRunningInterface {
+    Optional<WorkerJobRunning> findByTaskRunId(String taskRunId);
+
+    void delete(String taskRunId);
+
+    WorkerJobRunning save(WorkerJobRunning workerJobRunning);
+
+    List<WorkerJobRunning> getWorkerJobWithWorkerDead(List<String> workersAlive);
+}
diff --git a/core/src/main/java/io/kestra/core/runners/WorkerHeartbeat.java b/core/src/main/java/io/kestra/core/runners/WorkerHeartbeat.java
@@ -1,9 +1,6 @@
 package io.kestra.core.runners;
 
-import lombok.Builder;
-import lombok.Data;
-import lombok.NoArgsConstructor;
-import lombok.ToString;
+import lombok.*;
 import lombok.experimental.SuperBuilder;
 
 import java.time.Instant;
@@ -19,6 +16,16 @@ public class WorkerHeartbeat extends WorkerInstance {
     @Builder.Default
     private Instant heartbeatDate = Instant.now();
 
+    public WorkerInstance toWorkerInstance() {
+        return WorkerInstance.builder()
+            .workerUuid(this.getWorkerUuid())
+            .hostname(this.getHostname())
+            .port(this.getPort())
+            .managementPort(this.getManagementPort())
+            .workerGroup(this.getWorkerGroup())
+            .build();
+    }
+
     public enum Status {
         UP, DEAD
     }

diff --git a/core/src/main/java/io/kestra/core/runners/WorkerInstance.java b/core/src/main/java/io/kestra/core/runners/WorkerInstance.java
@@ -1,9 +1,6 @@
 package io.kestra.core.runners;
 
-import lombok.Builder;
-import lombok.Data;
-import lombok.NoArgsConstructor;
-import lombok.ToString;
+import lombok.*;
 import lombok.experimental.SuperBuilder;
 
 import javax.validation.constraints.NotNull;
@@ -15,6 +12,7 @@
 @SuperBuilder(toBuilder = true)
 @ToString
 @NoArgsConstructor
+@Getter
 public class WorkerInstance {
     @NotNull
     private UUID workerUuid;

diff --git a/core/src/main/java/io/kestra/core/runners/WorkerJob.java b/core/src/main/java/io/kestra/core/runners/WorkerJob.java
@@ -12,4 +12,6 @@ public abstract class WorkerJob {
     abstract public String getType();
 
     abstract public String uid();
+
+    abstract public String taskRunId();
 }
diff --git a/core/src/main/java/io/kestra/core/runners/WorkerTask.java b/core/src/main/java/io/kestra/core/runners/WorkerTask.java
@@ -41,4 +41,9 @@ public Logger logger() {
     public String uid() {
         return this.taskRun.getTaskId();
     }
+
+    @Override
+    public String taskRunId() {
+        return this.taskRun.getId();
+    }
 }
diff --git a/core/src/main/java/io/kestra/core/runners/WorkerTrigger.java b/core/src/main/java/io/kestra/core/runners/WorkerTrigger.java
@@ -31,4 +31,9 @@ public class WorkerTrigger extends WorkerJob {
     public String uid() {
         return triggerContext.uid();
     }
+
+    @Override
+    public String taskRunId() {
+        return triggerContext.uid();
+    }
 }
diff --git a/core/src/main/java/io/kestra/core/runners/WorkerTriggerRunning.java b/core/src/main/java/io/kestra/core/runners/WorkerTriggerRunning.java
@@ -4,7 +4,6 @@
 import io.kestra.core.models.conditions.ConditionContext;
 import io.kestra.core.models.triggers.AbstractTrigger;
 import io.kestra.core.models.triggers.TriggerContext;
-import lombok.Builder;
 import lombok.Data;
 import lombok.NoArgsConstructor;
 import lombok.experimental.SuperBuilder;

diff --git a/jdbc-h2/src/main/java/io/kestra/repository/h2/H2WorkerJobRunningRepository.java b/jdbc-h2/src/main/java/io/kestra/repository/h2/H2WorkerJobRunningRepository.java
@@ -0,0 +1,16 @@
+package io.kestra.repository.h2;
+
+import io.kestra.core.runners.WorkerJobRunning;
+import io.kestra.jdbc.repository.AbstractJdbcWorkerJobRunningRepository;
+import io.micronaut.context.ApplicationContext;
+import jakarta.inject.Inject;
+import jakarta.inject.Singleton;
+
+@Singleton
+@H2RepositoryEnabled
+public class H2WorkerJobRunningRepository extends AbstractJdbcWorkerJobRunningRepository {
+    @Inject
+    public H2WorkerJobRunningRepository(ApplicationContext applicationContext) {
+        super(new H2Repository<>(WorkerJobRunning.class, applicationContext));
+    }
+}
diff --git a/jdbc-h2/src/main/java/io/kestra/runner/h2/H2WorkerJobQueue.java b/jdbc-h2/src/main/java/io/kestra/runner/h2/H2WorkerJobQueue.java
@@ -4,27 +4,76 @@
 import io.kestra.core.queues.QueueFactoryInterface;
 import io.kestra.core.queues.QueueInterface;
 import io.kestra.core.queues.WorkerJobQueueInterface;
-import io.kestra.core.runners.WorkerJob;
+import io.kestra.core.runners.*;
 import io.kestra.core.utils.Either;
+import io.kestra.jdbc.repository.AbstractJdbcWorkerJobRunningRepository;
+import io.kestra.jdbc.runner.JdbcHeartbeat;
+import io.kestra.jdbc.runner.JdbcQueue;
 import io.micronaut.context.ApplicationContext;
 import io.micronaut.inject.qualifiers.Qualifiers;
+import lombok.extern.slf4j.Slf4j;
 
+import java.net.UnknownHostException;
 import java.util.function.Consumer;
 
+@Slf4j
 public class H2WorkerJobQueue implements WorkerJobQueueInterface {
-    QueueInterface<WorkerJob> workerTaskQueue;
+    JdbcQueue<WorkerJob> workerTaskQueue;
+    JdbcHeartbeat jdbcHeartbeat;
+    AbstractJdbcWorkerJobRunningRepository jdbcWorkerJobRunningRepository;
 
     @SuppressWarnings("unchecked")
     public H2WorkerJobQueue(ApplicationContext applicationContext) {
-        this.workerTaskQueue = (QueueInterface<WorkerJob>) applicationContext.getBean(
+        this.workerTaskQueue = (JdbcQueue<WorkerJob>) applicationContext.getBean(
             QueueInterface.class,
             Qualifiers.byName(QueueFactoryInterface.WORKERJOB_NAMED)
         );
+        this.jdbcHeartbeat = applicationContext.getBean(JdbcHeartbeat.class);
+        this.jdbcWorkerJobRunningRepository = applicationContext.getBean(AbstractJdbcWorkerJobRunningRepository.class);
     }
 
     @Override
     public Runnable receive(String consumerGroup, Class<?> queueType, Consumer<Either<WorkerJob, DeserializationException>> consumer) {
-        return workerTaskQueue.receive(consumerGroup, queueType, consumer);
+        return workerTaskQueue.receiveTransaction(consumerGroup, queueType, (dslContext, eithers) -> {
+
+            eithers.forEach(either -> {
+                if (either.isRight()) {
+                    log.error("Unable to deserialize an execution: {}", either.getRight().getMessage());
+                    return;
+                }
+                WorkerJob workerJob = either.getLeft();
+                WorkerInstance workerInstance = null;
+                try {
+                    workerInstance = jdbcHeartbeat.get().toWorkerInstance();
+                } catch (UnknownHostException e) {
+                    throw new RuntimeException(e);
+                }
+                WorkerJobRunning workerJobRunning;
+
+                if (workerJob instanceof WorkerTask workerTask) {
+                    workerJobRunning = WorkerTaskRunning.of(
+                        workerTask,
+                        workerInstance,
+                        0
+                    );
+                } else if (workerJob instanceof WorkerTrigger workerTrigger) {
+                    workerJobRunning = WorkerTriggerRunning.of(
+                        workerTrigger,
+                        workerInstance,
+                        0
+                    );
+                }
+                else {
+                    throw new IllegalArgumentException("Message is of type " + workerJob.getClass() + " which should never occurs");
+                }
+
+                jdbcWorkerJobRunningRepository.save(workerJobRunning, dslContext);
+
+                log.trace("Sending a workerJobRunning: {}", workerJobRunning);
+            });
+
+            eithers.forEach(consumer);
+        });
     }
 
     @Override

diff --git a/jdbc-h2/src/main/resources/migrations/h2/V1_2__initial.sql b/jdbc-h2/src/main/resources/migrations/h2/V1_2__initial.sql
@@ -9,4 +9,12 @@ CREATE TABLE IF NOT EXISTS worker_heartbeat (
     "worker_group" VARCHAR(150) GENERATED ALWAYS AS (JQ_STRING("value",'.workerGroup')),
     "status" VARCHAR(10) NOT NULL GENERATED ALWAYS AS (JQ_STRING("value",'.status')),
     "heartbeat_date" TIMESTAMP NOT NULL GENERATED ALWAYS AS (PARSEDATETIME(JQ_STRING("value", '.heartbeatDate'), 'yyyy-MM-dd''T''HH:mm:ss.SSSXXX'))
-    );
+);
+
+/* ----------------------- worker_job_running ----------------------- */
+CREATE TABLE IF NOT EXISTS worker_job_running (
+    "key" VARCHAR(250) NOT NULL PRIMARY KEY,
+    "value" TEXT NOT NULL,
+    "worker_uuid" VARCHAR(36) NOT NULL GENERATED ALWAYS AS (JQ_STRING("value",'.workerInstance.workerUuid')),
+    "taskrun_id" VARCHAR(150) NOT NULL GENERATED ALWAYS AS (JQ_STRING("value",'.taskRun.id'))
+);
diff --git a/jdbc-mysql/src/main/java/io/kestra/repository/mysql/MysqlWorkerJobRunningRepository.java b/jdbc-mysql/src/main/java/io/kestra/repository/mysql/MysqlWorkerJobRunningRepository.java
@@ -0,0 +1,16 @@
+package io.kestra.repository.mysql;
+
+import io.kestra.core.runners.WorkerJobRunning;
+import io.kestra.jdbc.repository.AbstractJdbcWorkerJobRunningRepository;
+import io.micronaut.context.ApplicationContext;
+import jakarta.inject.Inject;
+import jakarta.inject.Singleton;
+
+@Singleton
+@MysqlRepositoryEnabled
+public class MysqlWorkerJobRunningRepository extends AbstractJdbcWorkerJobRunningRepository {
+    @Inject
+    public MysqlWorkerJobRunningRepository(ApplicationContext applicationContext) {
+        super(new MysqlRepository<>(WorkerJobRunning.class, applicationContext));
+    }
+}
diff --git a/jdbc-mysql/src/main/java/io/kestra/runner/mysql/MysqlWorkerJobQueue.java b/jdbc-mysql/src/main/java/io/kestra/runner/mysql/MysqlWorkerJobQueue.java
@@ -4,27 +4,76 @@
 import io.kestra.core.queues.QueueFactoryInterface;
 import io.kestra.core.queues.QueueInterface;
 import io.kestra.core.queues.WorkerJobQueueInterface;
-import io.kestra.core.runners.WorkerJob;
+import io.kestra.core.runners.*;
 import io.kestra.core.utils.Either;
+import io.kestra.jdbc.repository.AbstractJdbcWorkerJobRunningRepository;
+import io.kestra.jdbc.runner.JdbcHeartbeat;
+import io.kestra.jdbc.runner.JdbcQueue;
 import io.micronaut.context.ApplicationContext;
 import io.micronaut.inject.qualifiers.Qualifiers;
+import lombok.extern.slf4j.Slf4j;
 
+import java.net.UnknownHostException;
 import java.util.function.Consumer;
 
+@Slf4j
 public class MysqlWorkerJobQueue implements WorkerJobQueueInterface {
-    QueueInterface<WorkerJob> workerTaskQueue;
+    JdbcQueue<WorkerJob> workerTaskQueue;
+    JdbcHeartbeat jdbcHeartbeat;
+    AbstractJdbcWorkerJobRunningRepository jdbcWorkerJobRunningRepository;
 
     @SuppressWarnings("unchecked")
     public MysqlWorkerJobQueue(ApplicationContext applicationContext) {
-        this.workerTaskQueue = (QueueInterface<WorkerJob>) applicationContext.getBean(
+        this.workerTaskQueue = (JdbcQueue<WorkerJob>) applicationContext.getBean(
             QueueInterface.class,
             Qualifiers.byName(QueueFactoryInterface.WORKERJOB_NAMED)
         );
+        this.jdbcHeartbeat = applicationContext.getBean(JdbcHeartbeat.class);
+        this.jdbcWorkerJobRunningRepository = applicationContext.getBean(AbstractJdbcWorkerJobRunningRepository.class);
     }
 
     @Override
     public Runnable receive(String consumerGroup, Class<?> queueType, Consumer<Either<WorkerJob, DeserializationException>> consumer) {
-        return workerTaskQueue.receive(consumerGroup, queueType, consumer);
+        return workerTaskQueue.receiveTransaction(consumerGroup, queueType, (dslContext, eithers) -> {
+
+            eithers.forEach(either -> {
+                if (either.isRight()) {
+                    log.error("Unable to deserialize an execution: {}", either.getRight().getMessage());
+                    return;
+                }
+                WorkerJob workerJob = either.getLeft();
+                WorkerInstance workerInstance = null;
+                try {
+                    workerInstance = jdbcHeartbeat.get().toWorkerInstance();
+                } catch (UnknownHostException e) {
+                    throw new RuntimeException(e);
+                }
+                WorkerJobRunning workerJobRunning;
+
+                if (workerJob instanceof WorkerTask workerTask) {
+                    workerJobRunning = WorkerTaskRunning.of(
+                        workerTask,
+                        workerInstance,
+                        0
+                    );
+                } else if (workerJob instanceof WorkerTrigger workerTrigger) {
+                    workerJobRunning = WorkerTriggerRunning.of(
+                        workerTrigger,
+                        workerInstance,
+                        0
+                    );
+                }
+                else {
+                    throw new IllegalArgumentException("Message is of type " + workerJob.getClass() + " which should never occurs");
+                }
+
+                jdbcWorkerJobRunningRepository.save(workerJobRunning, dslContext);
+
+                log.trace("Sending a workerJobRunning: {}", workerJobRunning);
+            });
+
+            eithers.forEach(consumer);
+        });
     }
 
     @Override

diff --git a/jdbc-mysql/src/main/resources/migrations/mysql/V1_2__initial.sql b/jdbc-mysql/src/main/resources/migrations/mysql/V1_2__initial.sql
@@ -2,12 +2,19 @@
 CREATE TABLE IF NOT EXISTS worker_heartbeat (
     `key` VARCHAR(250) NOT NULL PRIMARY KEY,
     `value` JSON NOT NULL,
-    `worker_uuid` VARCHAR(36)GENERATED ALWAYS AS (value ->> '$.workerUuid') STORED NOT NULL ,
-    `hostname` VARCHAR(150)GENERATED ALWAYS AS (value ->> '$.hostname') STORED NOT NULL ,
+    `worker_uuid` VARCHAR(36) GENERATED ALWAYS AS (value ->> '$.workerUuid') STORED NOT NULL,
+    `hostname` VARCHAR(150) GENERATED ALWAYS AS (value ->> '$.hostname') STORED NOT NULL,
     `port` INT UNSIGNED GENERATED ALWAYS AS (value ->> '$.port') STORED,
     `management_port` INT UNSIGNED GENERATED ALWAYS AS (value ->> '$.managementPort') STORED,
     `worker_group` VARCHAR(150) GENERATED ALWAYS AS (value ->> '$.workerGroup') STORED,
-    `status` VARCHAR(10)GENERATED ALWAYS AS (value ->> '$.status') STORED NOT NULL ,
+    `status` VARCHAR(10)GENERATED ALWAYS AS (value ->> '$.status') STORED NOT NULL,
     `heartbeat_date` DATETIME(6) GENERATED ALWAYS AS (STR_TO_DATE(value ->> '$.heartbeatDate' , '%Y-%m-%dT%H:%i:%s.%fZ')) STORED NOT NULL
     ) ENGINE INNODB CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;
 
+/* ----------------------- worker_job_running ----------------------- */
+CREATE TABLE IF NOT EXISTS worker_job_running (
+    `key` VARCHAR(250) NOT NULL PRIMARY KEY,
+    `value` JSON NOT NULL,
+    `worker_uuid` VARCHAR(36) GENERATED ALWAYS AS (value ->> '$.workerInstance.workerUuid') STORED NOT NULL,
+    `taskrun_id` VARCHAR(150) GENERATED ALWAYS AS (value ->> '$.taskRun.id') STORED NOT NULL
+    ) ENGINE INNODB CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;
diff --git a/...tgres/src/main/java/io/kestra/repository/postgres/PostgresWorkerJobRunningRepository.java b/...tgres/src/main/java/io/kestra/repository/postgres/PostgresWorkerJobRunningRepository.java
@@ -0,0 +1,16 @@
+package io.kestra.repository.postgres;
+
+import io.kestra.core.runners.WorkerJobRunning;
+import io.kestra.jdbc.repository.AbstractJdbcWorkerJobRunningRepository;
+import io.micronaut.context.ApplicationContext;
+import jakarta.inject.Inject;
+import jakarta.inject.Singleton;
+
+@Singleton
+@PostgresRepositoryEnabled
+public class PostgresWorkerJobRunningRepository extends AbstractJdbcWorkerJobRunningRepository {
+    @Inject
+    public PostgresWorkerJobRunningRepository(ApplicationContext applicationContext) {
+        super(new PostgresRepository<>(WorkerJobRunning.class, applicationContext));
+    }
+}