feat(core,jdbc): small trigger / scheduler improvements

Author: loicmathieu

Diff for: core/src/main/java/io/kestra/core/schedulers/AbstractScheduler.java

@@ -337,11 +337,6 @@ private void handle() {
                     .conditionContext(flowWithTriggers.getConditionContext())
                     .triggerContext(flowWithTriggers.TriggerContext.toBuilder().date(now()).stopAfter(flowWithTriggers.getAbstractTrigger().getStopAfter()).build())
                     .build())
-                .peek(f -> {
-                    if (f.getTriggerContext().getEvaluateRunningDate() != null || isExecutionNotRunning(f)) {
-                        this.triggerState.unlock(f.getTriggerContext());
-                    }
-                })
                 .filter(f -> f.getTriggerContext().getEvaluateRunningDate() == null)
                 .filter(this::isExecutionNotRunning)
                 .map(FlowWithPollingTriggerNextDate::of)

Diff for: core/src/main/java/io/kestra/core/schedulers/DefaultScheduleContext.java

@@ -1,4 +1,11 @@
 package io.kestra.core.schedulers;
 
+import java.util.function.Consumer;
+
 // For tests purpose
-public class DefaultScheduleContext implements ScheduleContextInterface {}
+public class DefaultScheduleContext implements ScheduleContextInterface {
+    @Override
+    public void doInTransaction(Consumer<ScheduleContextInterface> consumer) {
+        consumer.accept(this);
+    }
+}

Diff for: core/src/main/java/io/kestra/core/schedulers/ScheduleContextInterface.java

@@ -1,4 +1,14 @@
 package io.kestra.core.schedulers;
 
+import java.util.function.Consumer;
+
+/**
+ * This context is used by the Scheduler to allow evaluating and updating triggers in a transaction from the main evaluation loop.
+ * See AbstractScheduler.handle().
+ */
 public interface ScheduleContextInterface {
+    /**
+     * Do trigger retrieval and updating in a single transaction.
+     */
+    void doInTransaction(Consumer<ScheduleContextInterface> consumer);
 }

Diff for: core/src/main/java/io/kestra/core/schedulers/SchedulerTriggerStateInterface.java

@@ -24,15 +24,14 @@ public interface SchedulerTriggerStateInterface {
 
     Trigger update(Flow flow, AbstractTrigger abstractTrigger, ConditionContext conditionContext) throws Exception;
 
-    List<Trigger> findByNextExecutionDateReadyForAllTenants(ZonedDateTime now, ScheduleContextInterface scheduleContext);
 
     /**
-     * Required for Kafka
+     * Used by the JDBC implementation: find triggers in all tenants.
      */
-    List<Trigger> findByNextExecutionDateReadyForGivenFlows(List<Flow> flows, ZonedDateTime now, ScheduleContextInterface scheduleContext);
+    List<Trigger> findByNextExecutionDateReadyForAllTenants(ZonedDateTime now, ScheduleContextInterface scheduleContext);
 
     /**
-     * Required for Kafka
+     * Used by the Kafka implementation: find triggers in the scheduler assigned flow (as in Kafka partition assignment).
      */
-    void unlock(Trigger trigger);
+    List<Trigger> findByNextExecutionDateReadyForGivenFlows(List<Flow> flows, ZonedDateTime now, ScheduleContextInterface scheduleContext);
 }

Diff for: jdbc/src/main/java/io/kestra/jdbc/runner/JdbcScheduler.java

@@ -100,7 +100,7 @@ public void run() {
     public void handleNext(List<Flow> flows, ZonedDateTime now, BiConsumer<List<Trigger>, ScheduleContextInterface> consumer) {
         JdbcSchedulerContext schedulerContext = new JdbcSchedulerContext(this.dslContextWrapper);
 
-        schedulerContext.startTransaction(scheduleContextInterface -> {
+        schedulerContext.doInTransaction(scheduleContextInterface -> {
             List<Trigger> triggers = this.triggerState.findByNextExecutionDateReadyForAllTenants(now, scheduleContextInterface);
 
             consumer.accept(triggers, scheduleContextInterface);

Diff for: jdbc/src/main/java/io/kestra/jdbc/runner/JdbcSchedulerContext.java

@@ -18,17 +18,14 @@ public JdbcSchedulerContext(JooqDSLContextWrapper dslContextWrapper) {
         this.dslContextWrapper = dslContextWrapper;
     }
 
-    public void startTransaction(Consumer<ScheduleContextInterface> consumer) {
+    @Override
+    public void doInTransaction(Consumer<ScheduleContextInterface> consumer) {
         this.dslContextWrapper.transaction(configuration -> {
             this.context = DSL.using(configuration);
 
             consumer.accept(this);
 
-            this.commit();
+            this.context.commit();
         });
     }
-
-    public void commit() {
-        this.context.commit();
-    }
 }

Diff for: jdbc/src/main/java/io/kestra/jdbc/runner/JdbcSchedulerTriggerState.java

@@ -84,7 +84,4 @@ public List<Trigger> findByNextExecutionDateReadyForAllTenants(ZonedDateTime now
     public List<Trigger> findByNextExecutionDateReadyForGivenFlows(List<Flow> flows, ZonedDateTime now, ScheduleContextInterface scheduleContext) {
         throw new NotImplementedException();
     }
-
-    @Override
-    public void unlock(Trigger trigger) {}
 }

Diff for: runner-memory/src/main/java/io/kestra/runner/memory/MemorySchedulerTriggerState.java

@@ -79,7 +79,4 @@ public List<Trigger> findByNextExecutionDateReadyForAllTenants(ZonedDateTime now
     public List<Trigger> findByNextExecutionDateReadyForGivenFlows(List<Flow> flows, ZonedDateTime now, ScheduleContextInterface scheduleContext) {
         throw new NotImplementedException();
     }
-
-    @Override
-    public void unlock(Trigger trigger) {}
 }