feat(core): improve the Worker usage of workerThreadReferences

kestra-io · May 29, 2023 · 1de9135 · 1de9135
1 parent 0cac32c
commit 1de9135
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 10 deletions.
diff --git a/core/src/main/java/io/kestra/core/endpoints/WorkerEndpoint.java b/core/src/main/java/io/kestra/core/endpoints/WorkerEndpoint.java
@@ -30,11 +30,11 @@ public WorkerEndpointResult running() {
                 .sum()
             )
             .runnings(
-                worker.getWorkerThreadReferences()
+                worker.getWorkerThreadTasks()
                     .stream()
-                    .map(workerThread -> new WorkerEndpointWorkerTask(
-                        workerThread.getWorkerTask().getTaskRun(),
-                        workerThread.getWorkerTask().getTask())
+                    .map(workerTask -> new WorkerEndpointWorkerTask(
+                        workerTask.getTaskRun(),
+                        workerTask.getTask())
                     )
                     .collect(Collectors.toList())
             )

diff --git a/core/src/main/java/io/kestra/core/runners/Worker.java b/core/src/main/java/io/kestra/core/runners/Worker.java
@@ -40,6 +40,7 @@
 import java.time.temporal.ChronoUnit;
 import java.util.*;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -62,7 +63,6 @@ public class Worker implements Runnable, Closeable {
     @Getter
     private final Map<Long, AtomicInteger> metricRunningCount = new ConcurrentHashMap<>();
 
-    @Getter
     private final List<WorkerThread> workerThreadReferences = new ArrayList<>();
 
     @SuppressWarnings("unchecked")
@@ -266,7 +266,7 @@ private WorkerTaskResult run(WorkerTask workerTask, Boolean cleanUp) throws Queu
         finalWorkerTask = finalWorkerTask.withTaskRun(finalWorkerTask.getTaskRun().withState(state));
 
         // if resulting object can't be emitted (mostly size of message), we just can't emit it like that.
-        // So we just tryed to failed the status of the worker task, in this case, no log can't be happend, just
+        // So we just tried to fail the status of the worker task, in this case, no log can't be happend, just
         // changing status must work in order to finish current task (except if we are near the upper bound size).
         try {
             WorkerTaskResult workerTaskResult = new WorkerTaskResult(finalWorkerTask, dynamicWorkerResults);
@@ -325,7 +325,6 @@ private WorkerTask runAttempt(WorkerTask workerTask) {
         metricRunningCount.incrementAndGet();
 
         WorkerThread workerThread = new WorkerThread(logger, workerTask, task, runContext, metricRegistry);
-        workerThread.start();
 
         // emit attempts
         this.workerTaskResultQueue.emit(new WorkerTaskResult(workerTask
@@ -341,6 +340,7 @@ private WorkerTask runAttempt(WorkerTask workerTask) {
             synchronized (this) {
                 workerThreadReferences.add(workerThread);
             }
+            workerThread.start();
             workerThread.join();
             state = workerThread.getTaskState();
         } catch (InterruptedException e) {
@@ -430,7 +430,7 @@ public void close() throws IOException {
 
         Await.until(
             () -> {
-                if (this.executors.isTerminated() && this.getWorkerThreadReferences().size() == 0) {
+                if (this.executors.isTerminated() && this.workerThreadReferences.size() == 0) {
                     log.info("No more workers busy, shutting down!");
 
                     // we ensure that last produce message are send
@@ -445,7 +445,7 @@ public void close() throws IOException {
 
                 log.warn(
                     "Waiting worker with still {} thread(s) running, waiting!",
-                    this.getWorkerThreadReferences().size()
+                    this.workerThreadReferences.size()
                 );
 
                 return false;
@@ -459,6 +459,10 @@ public void close() throws IOException {
         metricEntryQueue.close();
     }
 
+    public List<WorkerTask> getWorkerThreadTasks() {
+        return this.workerThreadReferences.stream().map(thread -> thread.workerTask).toList();
+    }
+
     @Getter
     public static class WorkerThread extends Thread {
         Logger logger;
@@ -528,12 +532,12 @@ public void kill() {
             this.interrupt();
         }
 
-        @Synchronized
         private void exceptionHandler(Thread t, Throwable e) {
             if (!this.killed) {
                 logger.error(e.getMessage(), e);
                 taskState = io.kestra.core.models.flows.State.Type.FAILED;
             }
         }
     }
+
 }