feat(tests): intercept created executions through queue proxy & kill them if running after test

Author: brian-mulier-p

core/src/main/java/io/kestra/core/queues/QueueInterface.java

@@ -3,6 +3,7 @@
 import io.kestra.core.exceptions.DeserializationException;
 import io.kestra.core.models.Pauseable;
 import io.kestra.core.utils.Either;
+import org.apache.commons.lang3.NotImplementedException;
 
 import java.io.Closeable;
 import java.util.List;
@@ -54,4 +55,20 @@ default Runnable receive(String consumerGroup, Class<?> queueType, Consumer<Eith
     }
 
     Runnable receive(String consumerGroup, Class<?> queueType, Consumer<Either<T, DeserializationException>> consumer, boolean forUpdate);
+
+    default void deleteByKey(String key) throws QueueException {
+        throw new NotImplementedException();
+    }
+
+    default void deleteByKeys(List<String> keys) throws QueueException {
+        throw new NotImplementedException();
+    }
+
+    default void emitOnly(String consumerGroup, T message) throws QueueException {
+        throw new NotImplementedException();
+    }
+
+    default Runnable receiveBatch(Class<?> queueType, Consumer<List<Either<T, DeserializationException>>> consumer) {
+        throw new NotImplementedException();
+    }
 }

jdbc/src/main/java/io/kestra/jdbc/runner/JdbcExecutor.java

@@ -293,7 +293,7 @@ public void run() {
 
         Await.until(() -> this.allFlows != null, Duration.ofMillis(100), Duration.ofMinutes(5));
 
-        this.receiveCancellations.addFirst(((JdbcQueue<Execution>) this.executionQueue).receiveBatch(
+        this.receiveCancellations.addFirst(this.executionQueue.receiveBatch(
             Executor.class,
             executions -> {
                 List<CompletableFuture<Void>> futures = executions.stream()
@@ -302,7 +302,7 @@ public void run() {
                 CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join();
             }
         ));
-        this.receiveCancellations.addFirst(((JdbcQueue<WorkerTaskResult>) this.workerTaskResultQueue).receiveBatch(
+        this.receiveCancellations.addFirst(this.workerTaskResultQueue.receiveBatch(
             Executor.class,
             workerTaskResults -> {
                 List<CompletableFuture<Void>> futures = workerTaskResults.stream()
@@ -1121,14 +1121,14 @@ private void toExecution(Executor executor, boolean ignoreFailure) {
             // IMPORTANT: this must be done before emitting the last execution message so that all consumers are notified that the execution ends.
             // NOTE: we may also purge ExecutionKilled events, but as there may not be a lot of them, it may not be worth it.
             if (cleanExecutionQueue && isTerminated) {
-                ((JdbcQueue<Execution>) executionQueue).deleteByKey(executor.getExecution().getId());
+                executionQueue.deleteByKey(executor.getExecution().getId());
             }
 
             // emit for other consumers than the executor if no failure
             if (hasFailure) {
                 this.executionQueue.emit(executor.getExecution());
             } else {
-                ((JdbcQueue<Execution>) this.executionQueue).emitOnly(null, executor.getExecution());
+                this.executionQueue.emitOnly(null, executor.getExecution());
             }
 
             Execution execution = executor.getExecution();
@@ -1206,8 +1206,8 @@ private void toExecution(Executor executor, boolean ignoreFailure) {
                     List<String> taskRunKeys = executor.getExecution().getTaskRunList().stream()
                         .map(taskRun -> taskRun.getId())
                         .toList();
-                    ((JdbcQueue<WorkerTaskResult>) workerTaskResultQueue).deleteByKeys(taskRunKeys);
-                    ((JdbcQueue<WorkerJob>) workerJobQueue).deleteByKeys(taskRunKeys);
+                    workerTaskResultQueue.deleteByKeys(taskRunKeys);
+                    workerJobQueue.deleteByKeys(taskRunKeys);
                 }
             }
         } catch (QueueException e) {

jdbc/src/main/java/io/kestra/jdbc/runner/JdbcIndexer.java

@@ -39,10 +39,10 @@
 @JdbcRunnerEnabled
 public class JdbcIndexer implements Indexer {
     private final LogRepositoryInterface logRepository;
-    private final JdbcQueue<LogEntry> logQueue;
+    private final QueueInterface<LogEntry> logQueue;
 
     private final MetricRepositoryInterface metricRepository;
-    private final JdbcQueue<MetricEntry> metricQueue;
+    private final QueueInterface<MetricEntry> metricQueue;
     private final MetricRegistry metricRegistry;
     private final List<Runnable> receiveCancellations = new ArrayList<>();
 
@@ -67,9 +67,9 @@ public JdbcIndexer(
         QueueService queueService
     ) {
         this.logRepository = logRepository;
-        this.logQueue = (JdbcQueue<LogEntry>) logQueue;
+        this.logQueue = logQueue;
         this.metricRepository = metricRepositor;
-        this.metricQueue = (JdbcQueue<MetricEntry>) metricQueue;
+        this.metricQueue = metricQueue;
         this.metricRegistry = metricRegistry;
         this.eventPublisher = eventPublisher;
         this.skipExecutionService = skipExecutionService;
@@ -91,7 +91,7 @@ protected void startQueues() {
         this.sendBatch(metricQueue, metricRepository);
     }
 
-    protected <T> void sendBatch(JdbcQueue<T> queueInterface, SaveRepositoryInterface<T> saveRepositoryInterface) {
+    protected <T> void sendBatch(QueueInterface<T> queueInterface, SaveRepositoryInterface<T> saveRepositoryInterface) {
         this.receiveCancellations.addFirst(queueInterface.receiveBatch(Indexer.class, eithers -> {
             // first, log all deserialization issues
             eithers.stream().filter(either -> either.isRight()).forEach(either -> log.error("unable to deserialize an item: {}", either.getRight().getMessage()));

tests/src/main/java/io/kestra/core/junit/extensions/KestraTestExtension.java

@@ -1,13 +1,24 @@
 package io.kestra.core.junit.extensions;
 
 import io.kestra.core.junit.annotations.KestraTest;
+import io.kestra.core.models.executions.ExecutionKilled;
+import io.kestra.core.models.executions.ExecutionKilledExecution;
+import io.kestra.core.queues.QueueException;
+import io.kestra.core.queues.QueueFactoryInterface;
+import io.kestra.core.queues.QueueInterface;
+import io.kestra.core.repositories.ExecutionRepositoryInterface;
 import io.kestra.core.runners.TestRunner;
 import io.kestra.core.utils.TestsUtils;
+import io.micronaut.inject.qualifiers.Qualifiers;
 import io.micronaut.test.annotation.MicronautTestValue;
 import io.micronaut.test.extensions.junit5.MicronautJunit5Extension;
+import lombok.extern.slf4j.Slf4j;
 import org.junit.jupiter.api.extension.ExtensionContext;
 import org.junit.platform.commons.support.AnnotationSupport;
 
+import java.util.ConcurrentModificationException;
+
+@Slf4j
 public class KestraTestExtension extends MicronautJunit5Extension {
     private static final ExtensionContext.Namespace NAMESPACE = ExtensionContext.Namespace.create(KestraTestExtension.class);
 
@@ -47,9 +58,9 @@ public void beforeAll(ExtensionContext extensionContext) throws Exception {
         KestraTest kestraTest = extensionContext.getTestClass()
             .orElseThrow()
             .getAnnotation(KestraTest.class);
-        if (kestraTest.startRunner()){
+        if (kestraTest.startRunner()) {
             TestRunner runner = applicationContext.getBean(TestRunner.class);
-            if (!runner.isRunning()){
+            if (!runner.isRunning()) {
                 runner.setSchedulerEnabled(kestraTest.startScheduler());
                 runner.setWorkerEnabled(kestraTest.startWorker());
                 runner.run();
@@ -62,5 +73,49 @@ public void afterTestExecution(ExtensionContext context) throws Exception {
         super.afterTestExecution(context);
 
         TestsUtils.queueConsumersCleanup();
+
+        KestraTest kestraTest = context.getTestClass()
+            .orElseThrow()
+            .getAnnotation(KestraTest.class);
+        if (kestraTest.startRunner()
+            && applicationContext.containsBean(ExecutionRepositoryInterface.class)
+            && applicationContext.containsBean(QueueInterface.class, Qualifiers.byName(QueueFactoryInterface.KILL_NAMED))) {
+            ExecutionRepositoryInterface executionRepository = applicationContext.getBean(ExecutionRepositoryInterface.class);
+            QueueInterface<ExecutionKilled> killQueue = applicationContext.getBean(QueueInterface.class, Qualifiers.byName(QueueFactoryInterface.KILL_NAMED));
+
+            retryingExecutionKill(executionRepository, killQueue, 10);
+
+            TestRunner.testExecutions.get().clear();
+        }
+    }
+
+    private void retryingExecutionKill(ExecutionRepositoryInterface executionRepository, QueueInterface<ExecutionKilled> killQueue, int retriesLeft) throws InterruptedException {
+        try {
+            TestRunner.testExecutions.get().stream()
+                .flatMap(launchedExecution -> executionRepository.findById(launchedExecution.getTenantId(), launchedExecution.getId()).stream())
+                .filter(inRepository -> inRepository.getState().isRunning())
+                .forEach(inRepository -> {
+                    log.warn("Execution {} is still running after test execution, killing it", inRepository.getId());
+                    try {
+                        killQueue.emit(ExecutionKilledExecution.builder()
+                            .tenantId(inRepository.getTenantId())
+                            .executionId(inRepository.getId())
+                            .state(ExecutionKilled.State.REQUESTED)
+                            .isOnKillCascade(true)
+                            .build()
+                        );
+                    } catch (QueueException e) {
+                        log.warn("Couldn't kill execution {} after test execution", inRepository.getId(), e);
+                    }
+                });
+        } catch (ConcurrentModificationException e) {
+            // We intentionally don't use a CopyOnWriteArrayList to retry on concurrent modification exceptions to make sure to get rid of flakiness due to overflowing executions
+            if (retriesLeft <= 0) {
+                log.warn("Couldn't kill executions after test execution, due to concurrent modifications, this could impact further tests", e);
+                return;
+            }
+            Thread.sleep(100);
+            retryingExecutionKill(executionRepository, killQueue, retriesLeft - 1);
+        }
     }
 }

tests/src/main/java/io/kestra/core/queues/TestQueueFactory.java

@@ -0,0 +1,44 @@
+package io.kestra.core.queues;
+
+import io.kestra.core.models.executions.Execution;
+import io.micronaut.context.annotation.Bean;
+import io.micronaut.context.annotation.Factory;
+import io.micronaut.context.annotation.Replaces;
+import jakarta.inject.Named;
+import jakarta.inject.Singleton;
+
+import java.lang.reflect.Proxy;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Optional;
+
+@Factory
+public class TestQueueFactory {
+    private QueueInterface<Execution> delegate;
+    private List<Execution> testExecutions;
+
+    public TestQueueFactory(QueueFactoryInterface queueFactoryInterface) {
+        this.delegate = queueFactoryInterface.execution();
+    }
+
+    public void setTestExecutionsList(List<Execution> testExecutions) {
+        this.testExecutions = testExecutions;
+    }
+
+    @Singleton
+    @Replaces(named = QueueFactoryInterface.EXECUTION_NAMED)
+    @Named(QueueFactoryInterface.EXECUTION_NAMED)
+    @Bean(preDestroy = "close")
+    public QueueInterface<Execution> execution() {
+        return (QueueInterface<Execution>) Proxy.newProxyInstance(this.getClass().getClassLoader(), new Class[]{QueueInterface.class}, (proxy, method, args) -> {
+            try {
+                if (method.getName().contains("emit")) {
+                    Arrays.stream(args).filter(arg -> arg instanceof Execution).forEach(arg -> testExecutions.add((Execution) arg));
+                }
+                return method.invoke(this.delegate, args);
+            } catch (Exception e) {
+                throw Optional.ofNullable(e.getCause()).orElse(e);
+            }
+        });
+    }
+}

tests/src/main/java/io/kestra/core/runners/TestRunner.java

@@ -1,11 +1,15 @@
 package io.kestra.core.runners;
 
+import io.kestra.core.models.executions.Execution;
+import io.kestra.core.queues.TestQueueFactory;
+import io.kestra.core.queues.QueueFactoryInterface;
 import io.kestra.core.server.Service;
 import io.kestra.core.utils.Await;
 import io.kestra.core.utils.ExecutorsUtils;
 import io.kestra.worker.DefaultWorker;
 import io.micronaut.context.ApplicationContext;
 import io.micronaut.context.annotation.Value;
+import io.micronaut.inject.qualifiers.Qualifiers;
 import jakarta.annotation.PreDestroy;
 import jakarta.inject.Inject;
 import jakarta.inject.Singleton;
@@ -25,6 +29,8 @@
 @Slf4j
 @Singleton
 public class TestRunner implements Runnable, AutoCloseable {
+    public static final ThreadLocal<List<Execution>> testExecutions = ThreadLocal.withInitial(ArrayList::new);
+
     @Setter private int workerThread = Math.max(3, Runtime.getRuntime().availableProcessors()) * 16;
     @Setter private boolean schedulerEnabled = true;
     @Setter private boolean workerEnabled = true;
@@ -46,6 +52,9 @@ public class TestRunner implements Runnable, AutoCloseable {
 
     @Override
     public void run() {
+        TestQueueFactory executionQueue = applicationContext.getBean(TestQueueFactory.class);
+        executionQueue.setTestExecutionsList(testExecutions.get());
+
         running.set(true);
 
         poolExecutor = executorsUtils.cachedThreadPool("standalone-runner");