Adding ExecutionInterceptor in addition to SchedulerListener

Author: kagkarlsson

db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/ExecutePicked.java

@@ -15,6 +15,8 @@
 
 import static com.github.kagkarlsson.scheduler.ExceptionUtils.describe;
 
+import com.github.kagkarlsson.scheduler.event.ExecutionChain;
+import com.github.kagkarlsson.scheduler.event.ExecutionInterceptor;
 import com.github.kagkarlsson.scheduler.event.SchedulerListener.CandidateEventType;
 import com.github.kagkarlsson.scheduler.event.SchedulerListener.SchedulerEventType;
 import com.github.kagkarlsson.scheduler.event.SchedulerListeners;
@@ -23,21 +25,25 @@
 import com.github.kagkarlsson.scheduler.task.Execution;
 import com.github.kagkarlsson.scheduler.task.ExecutionComplete;
 import com.github.kagkarlsson.scheduler.task.ExecutionContext;
+import com.github.kagkarlsson.scheduler.task.ExecutionHandler;
 import com.github.kagkarlsson.scheduler.task.ExecutionOperations;
 import com.github.kagkarlsson.scheduler.task.Task;
 import java.time.Instant;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Optional;
 import java.util.UUID;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-@SuppressWarnings("rawtypes")
+@SuppressWarnings({"rawtypes", "unchecked"})
 class ExecutePicked implements Runnable {
   private static final Logger LOG = LoggerFactory.getLogger(ExecutePicked.class);
   private final Executor executor;
   private final TaskRepository taskRepository;
   private final SchedulerClient schedulerClient;
   private final SchedulerListeners schedulerListeners;
+  private final List<ExecutionInterceptor> executionInterceptors;
   private final TaskResolver taskResolver;
   private final SchedulerState schedulerState;
   private final ConfigurableLogger failureLogger;
@@ -50,6 +56,7 @@ public ExecutePicked(
       TaskRepository taskRepository,
       SchedulerClient schedulerClient,
       SchedulerListeners schedulerListeners,
+      List<ExecutionInterceptor> executionInterceptors,
       TaskResolver taskResolver,
       SchedulerState schedulerState,
       ConfigurableLogger failureLogger,
@@ -60,6 +67,7 @@ public ExecutePicked(
     this.taskRepository = taskRepository;
     this.schedulerClient = schedulerClient;
     this.schedulerListeners = schedulerListeners;
+    this.executionInterceptors = executionInterceptors;
     this.taskResolver = taskResolver;
     this.schedulerState = schedulerState;
     this.failureLogger = failureLogger;
@@ -97,12 +105,12 @@ private void executePickedExecution(Execution execution, CurrentlyExecuting curr
     Instant executionStarted = clock.now();
     try {
       LOG.debug("Executing: " + execution);
-      CompletionHandler completion =
-          task.get()
-              .execute(
-                  execution.taskInstance,
-                  new ExecutionContext(
-                      schedulerState, execution, schedulerClient, currentlyExecuting));
+      ExecutionHandler handler = task.get();
+      ExecutionContext executionContext =
+          new ExecutionContext(schedulerState, execution, schedulerClient, currentlyExecuting);
+      ExecutionChain chain = new ExecutionChain(new ArrayList<>(executionInterceptors), handler);
+
+      CompletionHandler completion = chain.proceed(execution.taskInstance, executionContext);
       LOG.debug("Execution done: " + execution);
 
       complete(completion, execution, executionStarted);

db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/FetchCandidates.java

@@ -13,6 +13,7 @@
  */
 package com.github.kagkarlsson.scheduler;
 
+import com.github.kagkarlsson.scheduler.event.ExecutionInterceptor;
 import com.github.kagkarlsson.scheduler.event.SchedulerListener.CandidateEventType;
 import com.github.kagkarlsson.scheduler.event.SchedulerListener.SchedulerEventType;
 import com.github.kagkarlsson.scheduler.event.SchedulerListeners;
@@ -32,6 +33,7 @@ public class FetchCandidates implements PollStrategy {
   private final TaskRepository taskRepository;
   private final SchedulerClient schedulerClient;
   private final SchedulerListeners schedulerListeners;
+  private final List<ExecutionInterceptor> executionInterceptors;
   private final SchedulerState schedulerState;
   private final ConfigurableLogger failureLogger;
   private final TaskResolver taskResolver;
@@ -49,6 +51,7 @@ public FetchCandidates(
       SchedulerClient schedulerClient,
       int threadpoolSize,
       SchedulerListeners schedulerListeners,
+      List<ExecutionInterceptor> executionInterceptors,
       SchedulerState schedulerState,
       ConfigurableLogger failureLogger,
       TaskResolver taskResolver,
@@ -60,6 +63,7 @@ public FetchCandidates(
     this.taskRepository = taskRepository;
     this.schedulerClient = schedulerClient;
     this.schedulerListeners = schedulerListeners;
+    this.executionInterceptors = executionInterceptors;
     this.schedulerState = schedulerState;
     this.failureLogger = failureLogger;
     this.taskResolver = taskResolver;
@@ -103,6 +107,7 @@ public void run() {
                             taskRepository,
                             schedulerClient,
                             schedulerListeners,
+                            executionInterceptors,
                             taskResolver,
                             schedulerState,
                             failureLogger,

db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/LockAndFetchCandidates.java

@@ -13,6 +13,7 @@
  */
 package com.github.kagkarlsson.scheduler;
 
+import com.github.kagkarlsson.scheduler.event.ExecutionInterceptor;
 import com.github.kagkarlsson.scheduler.event.SchedulerListener.SchedulerEventType;
 import com.github.kagkarlsson.scheduler.event.SchedulerListeners;
 import com.github.kagkarlsson.scheduler.logging.ConfigurableLogger;
@@ -29,6 +30,7 @@ public class LockAndFetchCandidates implements PollStrategy {
   private final TaskRepository taskRepository;
   private final SchedulerClient schedulerClient;
   private final SchedulerListeners schedulerListeners;
+  private final List<ExecutionInterceptor> executionInterceptors;
   private final TaskResolver taskResolver;
   private final SchedulerState schedulerState;
   private final ConfigurableLogger failureLogger;
@@ -46,6 +48,7 @@ public LockAndFetchCandidates(
       SchedulerClient schedulerClient,
       int threadpoolSize,
       SchedulerListeners schedulerListeners,
+      List<ExecutionInterceptor> executionInterceptors,
       SchedulerState schedulerState,
       ConfigurableLogger failureLogger,
       TaskResolver taskResolver,
@@ -57,6 +60,7 @@ public LockAndFetchCandidates(
     this.taskRepository = taskRepository;
     this.schedulerClient = schedulerClient;
     this.schedulerListeners = schedulerListeners;
+    this.executionInterceptors = executionInterceptors;
     this.taskResolver = taskResolver;
     this.schedulerState = schedulerState;
     this.failureLogger = failureLogger;
@@ -102,6 +106,7 @@ public void run() {
               taskRepository,
               schedulerClient,
               schedulerListeners,
+              executionInterceptors,
               taskResolver,
               schedulerState,
               failureLogger,

db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/Scheduler.java

@@ -17,6 +17,7 @@
 import static java.util.stream.Collectors.toList;
 
 import com.github.kagkarlsson.scheduler.SchedulerState.SettableSchedulerState;
+import com.github.kagkarlsson.scheduler.event.ExecutionInterceptor;
 import com.github.kagkarlsson.scheduler.event.SchedulerListener;
 import com.github.kagkarlsson.scheduler.event.SchedulerListener.SchedulerEventType;
 import com.github.kagkarlsson.scheduler.event.SchedulerListeners;
@@ -81,6 +82,7 @@ protected Scheduler(
       Duration heartbeatInterval,
       int numberOfMissedHeartbeatsBeforeDead,
       List<SchedulerListener> schedulerListeners,
+      List<ExecutionInterceptor> executionInterceptors,
       PollingStrategyConfig pollingStrategyConfig,
       Duration deleteUnresolvedAfter,
       Duration shutdownMaxWait,
@@ -120,6 +122,7 @@ protected Scheduler(
               this,
               threadpoolSize,
               this.schedulerListeners,
+              executionInterceptors,
               schedulerState,
               failureLogger,
               taskResolver,
@@ -135,6 +138,7 @@ protected Scheduler(
               this,
               threadpoolSize,
               this.schedulerListeners,
+              executionInterceptors,
               schedulerState,
               failureLogger,
               taskResolver,

db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/SchedulerBuilder.java

@@ -17,6 +17,7 @@
 import static com.github.kagkarlsson.scheduler.Scheduler.THREAD_PREFIX;
 import static java.util.Optional.ofNullable;
 
+import com.github.kagkarlsson.scheduler.event.ExecutionInterceptor;
 import com.github.kagkarlsson.scheduler.event.SchedulerListener;
 import com.github.kagkarlsson.scheduler.jdbc.AutodetectJdbcCustomization;
 import com.github.kagkarlsson.scheduler.jdbc.JdbcCustomization;
@@ -77,6 +78,7 @@ public class SchedulerBuilder {
   private int numberOfMissedHeartbeatsBeforeDead = DEFAULT_MISSED_HEARTBEATS_LIMIT;
   private boolean alwaysPersistTimestampInUTC = false;
   private List<SchedulerListener> schedulerListeners = new ArrayList<>();
+  private List<ExecutionInterceptor> executionInterceptors = new ArrayList<>();
 
   public SchedulerBuilder(DataSource dataSource, List<Task<?>> knownTasks) {
     this.dataSource = dataSource;
@@ -144,6 +146,11 @@ public SchedulerBuilder addSchedulerListener(SchedulerListener schedulerListener
     return this;
   }
 
+  public SchedulerBuilder addExecutionInterceptor(ExecutionInterceptor interceptor) {
+    this.executionInterceptors.add(interceptor);
+    return this;
+  }
+
   public SchedulerBuilder schedulerName(SchedulerName schedulerName) {
     this.schedulerName = schedulerName;
     return this;
@@ -301,6 +308,7 @@ public Scheduler build() {
             heartbeatInterval,
             numberOfMissedHeartbeatsBeforeDead,
             schedulerListeners,
+            executionInterceptors,
             pollingStrategyConfig,
             deleteUnresolvedAfter,
             shutdownMaxWait,

db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/event/ExecutionChain.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright (C) Gustav Karlsson
+ *
+ * <p>Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
+ * except in compliance with the License. You may obtain a copy of the License at
+ *
+ * <p>http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * <p>Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.github.kagkarlsson.scheduler.event;
+
+import com.github.kagkarlsson.scheduler.task.CompletionHandler;
+import com.github.kagkarlsson.scheduler.task.ExecutionContext;
+import com.github.kagkarlsson.scheduler.task.ExecutionHandler;
+import com.github.kagkarlsson.scheduler.task.TaskInstance;
+import java.util.List;
+
+@SuppressWarnings({"rawtypes", "unchecked"})
+public class ExecutionChain {
+
+  private final List<ExecutionInterceptor> interceptors;
+  private final ExecutionHandler<?> executionHandler;
+
+  public ExecutionChain(
+      List<ExecutionInterceptor> interceptors, ExecutionHandler<?> executionHandler) {
+    this.interceptors = interceptors;
+    this.executionHandler = executionHandler;
+  }
+
+  public CompletionHandler<?> proceed(
+      TaskInstance taskInstance, ExecutionContext executionContext) {
+    if (interceptors.isEmpty()) {
+      return executionHandler.execute(taskInstance, executionContext);
+    } else {
+      ExecutionInterceptor nextInterceptor = interceptors.get(0);
+      List<ExecutionInterceptor> rest = interceptors.subList(1, interceptors.size());
+      return nextInterceptor.execute(
+          taskInstance, executionContext, new ExecutionChain(rest, executionHandler));
+    }
+  }
+}

db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/event/ExecutionInterceptor.java

@@ -0,0 +1,23 @@
+/*
+ * Copyright (C) Gustav Karlsson
+ *
+ * <p>Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
+ * except in compliance with the License. You may obtain a copy of the License at
+ *
+ * <p>http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * <p>Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.github.kagkarlsson.scheduler.event;
+
+import com.github.kagkarlsson.scheduler.task.CompletionHandler;
+import com.github.kagkarlsson.scheduler.task.ExecutionContext;
+import com.github.kagkarlsson.scheduler.task.TaskInstance;
+
+public interface ExecutionInterceptor {
+  CompletionHandler<?> execute(
+      TaskInstance<?> taskInstance, ExecutionContext executionContext, ExecutionChain chain);
+}

db-scheduler/src/main/java/com/github/kagkarlsson/scheduler/testhelper/ManualScheduler.java

@@ -19,6 +19,7 @@
 import com.github.kagkarlsson.scheduler.task.OnStartup;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ScheduledExecutorService;
@@ -60,6 +61,7 @@ public class ManualScheduler extends Scheduler {
         heartbeatInterval,
         SchedulerBuilder.DEFAULT_MISSED_HEARTBEATS_LIMIT,
         schedulerListeners,
+        new ArrayList<>(),
         pollingStrategyConfig,
         deleteUnresolvedAfter,
         Duration.ZERO,

db-scheduler/src/test/java/com/github/kagkarlsson/scheduler/event/ExecutionChainTest.java

@@ -0,0 +1,68 @@
+package com.github.kagkarlsson.scheduler.event;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import com.github.kagkarlsson.scheduler.task.CompletionHandler;
+import com.github.kagkarlsson.scheduler.task.ExecutionContext;
+import com.github.kagkarlsson.scheduler.task.ExecutionHandler;
+import com.github.kagkarlsson.scheduler.task.TaskInstance;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.junit.jupiter.api.Test;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+@SuppressWarnings("rawtypes")
+class ExecutionChainTest {
+
+  @Test
+  public void happy() {
+    AtomicInteger counter = new AtomicInteger(0);
+    CountingExecutionHandler executionHandler = new CountingExecutionHandler(counter);
+    ExecutionChain chain =
+        new ExecutionChain(
+            List.of(new CountingInterceptor(counter), new CountingInterceptor(counter)),
+            executionHandler);
+
+    chain.proceed(new TaskInstance("task1", "id1"), dummyContext());
+
+    assertEquals(3, counter.get());
+  }
+
+  private static ExecutionContext dummyContext() {
+    return new ExecutionContext(null, null, null, null);
+  }
+
+  private static class CountingInterceptor implements ExecutionInterceptor {
+    private static final Logger LOG = LoggerFactory.getLogger(CountingInterceptor.class);
+    private final AtomicInteger counter;
+
+    public CountingInterceptor(AtomicInteger counter) {
+      this.counter = counter;
+    }
+
+    @Override
+    public CompletionHandler<?> execute(
+        TaskInstance<?> taskInstance, ExecutionContext executionContext, ExecutionChain chain) {
+
+      counter.incrementAndGet();
+      return chain.proceed(taskInstance, executionContext);
+    }
+  }
+
+  private static class CountingExecutionHandler implements ExecutionHandler<Void> {
+
+    private final AtomicInteger counter;
+
+    public CountingExecutionHandler(AtomicInteger counter) {
+      this.counter = counter;
+    }
+
+    @Override
+    public CompletionHandler<Void> execute(
+        TaskInstance<Void> taskInstance, ExecutionContext executionContext) {
+      counter.incrementAndGet();
+      return null;
+    }
+  }
+}