feat(kafka-runner): introduce async queue and send log async (#806)

kestra-io · Nov 8, 2022 · a2b6578 · a2b6578
1 parent 8a574b9
commit a2b6578
Show file tree

Hide file tree

Showing 8 changed files with 40 additions and 17 deletions.
diff --git a/core/src/main/java/io/kestra/core/queues/QueueInterface.java b/core/src/main/java/io/kestra/core/queues/QueueInterface.java
@@ -6,6 +6,8 @@
 public interface QueueInterface<T> extends Closeable {
     void emit(T message) throws QueueException;
 
+    void emitAsync(T message) throws QueueException;
+
     void delete(T message) throws QueueException;
 
     Runnable receive(Consumer<T> consumer);

diff --git a/core/src/main/java/io/kestra/core/runners/RunContextLogger.java b/core/src/main/java/io/kestra/core/runners/RunContextLogger.java
@@ -156,7 +156,7 @@ public ContextAppender(QueueInterface<LogEntry> logQueue, LogEntry logEntry) {
         @Override
         protected void append(ILoggingEvent e) {
             logEntries(e, logEntry)
-                .forEach(logQueue::emit);
+                .forEach(logQueue::emitAsync);
         }
     }
 

diff --git a/core/src/main/java/io/kestra/core/services/TaskDefaultService.java b/core/src/main/java/io/kestra/core/services/TaskDefaultService.java
@@ -62,7 +62,7 @@ public Flow injectDefaults(Flow flow, Execution execution) {
                     Execution.loggingEventFromException(e),
                     LogEntry.of(execution)
                 )
-                .forEach(logQueue::emit);
+                .forEach(logQueue::emitAsync);
             return flow;
         }
     }

diff --git a/jdbc/src/main/java/io/kestra/jdbc/runner/JdbcExecutor.java b/jdbc/src/main/java/io/kestra/jdbc/runner/JdbcExecutor.java
@@ -486,7 +486,7 @@ private boolean deduplicateFlowTrigger(Execution execution, ExecutorState execut
     private Executor handleFailedExecutionFromExecutor(Executor executor, Exception e) {
         Execution.FailedExecutionWithLog failedExecutionWithLog = executor.getExecution().failedExecutionFromExecutor(e);
         try {
-            failedExecutionWithLog.getLogs().forEach(logQueue::emit);
+            failedExecutionWithLog.getLogs().forEach(logQueue::emitAsync);
 
             return executor.withExecution(failedExecutionWithLog.getExecution(), "exception");
         } catch (Exception ex) {

diff --git a/jdbc/src/main/java/io/kestra/jdbc/runner/JdbcQueue.java b/jdbc/src/main/java/io/kestra/jdbc/runner/JdbcQueue.java
@@ -118,6 +118,11 @@ public void emit(T message) {
         this.produce(queueService.key(message), message, false);
     }
 
+    @Override
+    public void emitAsync(T message) throws QueueException {
+        this.emit(message);
+    }
+
     @Override
     public void delete(T message) throws QueueException {
         dslContextWrapper.transaction(configuration -> DSL

diff --git a/...r-kafka/src/main/java/io/kestra/runner/kafka/KafkaExecutorProductionExceptionHandler.java b/...r-kafka/src/main/java/io/kestra/runner/kafka/KafkaExecutorProductionExceptionHandler.java
@@ -64,7 +64,7 @@ public ProductionExceptionHandlerResponse handle(ProducerRecord<byte[], byte[]>
                 Execution.FailedExecutionWithLog failedExecutionWithLog = execution.failedExecutionFromExecutor(exception);
                 Execution sendExecution = failedExecutionWithLog.getExecution();
 
-                failedExecutionWithLog.getLogs().forEach(logEntry -> logQueue.emit(logEntry));
+                failedExecutionWithLog.getLogs().forEach(logEntry -> logQueue.emitAsync(logEntry));
 
                 if (exception instanceof RecordTooLargeException) {
                     boolean exit = false;

diff --git a/runner-kafka/src/main/java/io/kestra/runner/kafka/KafkaQueue.java b/runner-kafka/src/main/java/io/kestra/runner/kafka/KafkaQueue.java
@@ -19,6 +19,7 @@
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.producer.KafkaProducer;
 import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.errors.WakeupException;
@@ -28,6 +29,7 @@
 import java.util.*;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
@@ -84,31 +86,40 @@ public KafkaQueue(String topicKey, Class<T> cls, ApplicationContext applicationC
         kafkaAdminService.createIfNotExist(topicKey);
     }
 
+    private Future<RecordMetadata> produceAsync(String key, T message) {
+        return kafkaProducer
+            .send(
+                new ProducerRecord<>(
+                    topicsConfig.getName(),
+                    key, message
+                ),
+                (metadata, e) -> {
+                    if (e != null) {
+                        log.error("Failed to produce on '{}' with key '{}', metadata '{}' ", this.cls, key,  metadata, e);
+                    }
+                }
+            );
+    }
+
     private void produce(String key, T message) {
         try {
-            kafkaProducer
-                .send(
-                    new ProducerRecord<>(
-                        topicsConfig.getName(),
-                        key, message
-                    ),
-                    (metadata, e) -> {
-                        if (e != null) {
-                            log.error("Failed to produce on '{}' with key '{}', metadata '{}' ", this.cls, key,  metadata, e);
-                        }
-                    }
-                )
-                .get();
+            this.produceAsync(key, message).get();
         } catch (InterruptedException | ExecutionException e) {
             throw new QueueException("Failed to produce on '" + this.cls + "' with key '" +  key + "': ", e);
         }
     }
 
+
     @Override
     public void emit(T message) throws QueueException {
         this.produce(this.queueService.key(message), message);
     }
 
+    @Override
+    public void emitAsync(T message) throws QueueException {
+        this.produceAsync(this.queueService.key(message), message);
+    }
+
     @Override
     public void delete(T message) throws QueueException {
         this.produce(this.queueService.key(message), null);

diff --git a/runner-memory/src/main/java/io/kestra/runner/memory/MemoryQueue.java b/runner-memory/src/main/java/io/kestra/runner/memory/MemoryQueue.java
@@ -84,6 +84,11 @@ public void emit(T message) {
         this.produce(queueService.key(message), message);
     }
 
+    @Override
+    public void emitAsync(T message) throws QueueException {
+        this.emit(message);
+    }
+
     @Override
     public void delete(T message) throws QueueException {
         this.produce(queueService.key(message), null);