Add proper support for @scheduled annotation processing by Spring Boo…

…t 3.2

.github/workflows/build.yml

@@ -18,9 +18,8 @@ jobs:
       max-parallel: 100
       matrix:
         spring_boot_version:
-          - 3.2.2
+          - 3.2.3
           - 3.1.2
-          - 3.0.7
           - 2.7.13
     env:
       SPRING_BOOT_VERSION: ${{ matrix.spring_boot_version }}

CHANGELOG.md

@@ -10,6 +10,8 @@ to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 ### Changed
 
 * Add support for Spring Boot 3.2 and bump some dependencies.
+  * In Spring Boot 3.2 logic for assigning a scheduler for executing `@Scheduled` annotated methods changed, so needed to refactored auto configuration and shutdown logic of `TaskSchedulersGracefulShutdownStrategy` to always ask the `ScheduledTaskRegistrar` for the actual `TaskScheduler` before shutting it down.
+  * Add support for `TaskSchedulerRouter` in `TaskSchedulersGracefulShutdownStrategy`
 
 ## [2.14.2] - 2023-07-28
 

build.libraries.gradle

@@ -1,5 +1,5 @@
 ext {
-    springBootVersion = "${System.getenv("SPRING_BOOT_VERSION") ?: '2.6.15'}"
+    springBootVersion = "${System.getenv("SPRING_BOOT_VERSION") ?: '2.7.13'}"
 
     libraries = [
             // version defined

core/src/main/java/com/transferwise/common/gracefulshutdown/GracefulShutdownAutoConfiguration.java

@@ -116,7 +116,7 @@ public KagkarlssonDbScheduledTaskShutdownStrategy kagkarlssonDbScheduledTaskShut
   }
 
   @Configuration
-  @ConditionalOnBean(type = "org.springframework.scheduling.TaskScheduler")
+  @ConditionalOnBean(type = {"org.springframework.scheduling.annotation.ScheduledAnnotationBeanPostProcessor"})
   @ConditionalOnProperty(value = "tw-graceful-shutdown.spring-task-scheduler.enabled", matchIfMissing = true)
   protected static class SpringTaskSchedulerConfiguration {
 
@@ -126,27 +126,12 @@ public TaskSchedulersGracefulShutdownStrategy taskSchedulersGracefulShutdownStra
         @Autowired GracefulShutdownProperties gracefulShutdownProperties) {
       return new TaskSchedulersGracefulShutdownStrategy(applicationContext, gracefulShutdownProperties);
     }
-  }
-
-  @Configuration
-  @ConditionalOnMissingBean(type = "org.springframework.scheduling.TaskScheduler")
-  @ConditionalOnBean(type = {"org.springframework.scheduling.annotation.ScheduledAnnotationBeanPostProcessor",
-      "org.springframework.scheduling.annotation.SchedulingConfigurer"})
-  @ConditionalOnProperty(value = "tw-graceful-shutdown.spring-task-scheduler.enabled", matchIfMissing = true)
-  protected static class SpringTaskSchedulerAlternativeConfiguration {
-
-    @Bean
-    public TaskSchedulersGracefulShutdownStrategy taskSchedulersGracefulShutdownStrategy(
-        @Autowired ApplicationContext applicationContext,
-        @Autowired GracefulShutdownProperties gracefulShutdownProperties) {
-      return new TaskSchedulersGracefulShutdownStrategy(applicationContext, gracefulShutdownProperties);
-    }
 
     @Bean
-    @Order
     public SchedulingConfigurer twGsSchedulingConfigurer(TaskSchedulersGracefulShutdownStrategy taskSchedulersGracefulShutdownStrategy) {
-      return taskRegistrar -> taskSchedulersGracefulShutdownStrategy.addResource(taskRegistrar.getScheduler());
+      return taskSchedulersGracefulShutdownStrategy::setTaskRegistrar;
     }
+
   }
 
   @Configuration

core/src/main/java/com/transferwise/common/gracefulshutdown/strategies/TaskSchedulersGracefulShutdownStrategy.java

@@ -2,22 +2,55 @@
 
 import com.transferwise.common.gracefulshutdown.config.GracefulShutdownProperties;
 import com.transferwise.common.gracefulshutdown.utils.ExecutorShutdownUtils;
+import java.lang.reflect.Field;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
 import java.time.Duration;
+import java.util.Set;
 import java.util.concurrent.Executor;
 import lombok.NonNull;
+import lombok.Setter;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.context.ApplicationContext;
 import org.springframework.scheduling.TaskScheduler;
+import org.springframework.scheduling.config.ScheduledTaskRegistrar;
+import org.springframework.scheduling.config.Task;
 import reactor.core.publisher.Mono;
 
 @Slf4j
 public class TaskSchedulersGracefulShutdownStrategy extends BaseReactiveResourceShutdownStrategy<TaskScheduler> {
 
+  private static final Class taskSchedulerRouter;
+  private static final Field taskSchedulerRouterLocalExecutorField;
+  private static final Method taskSchedulerRouterDestoryMethod;
+
+  static {
+    Class clazz = null;
+    Field field = null;
+    Method method = null;
+    try {
+      clazz = Class.forName("org.springframework.scheduling.config.TaskSchedulerRouter", true, ExecutorShutdownUtils.class.getClassLoader());
+      field = clazz.getDeclaredField("localExecutor");
+      field.setAccessible(true);
+      method = clazz.getDeclaredMethod("destroy");
+    } catch (ClassNotFoundException e) {
+      // ignore as it's normal in pre Spring 6.1 environment.
+    } catch (NoSuchFieldException | NoSuchMethodException e) {
+      throw new RuntimeException("org.springframework.scheduling.config.TaskSchedulerRouter class is missing expected field or method. Contact the SRE team.", e);
+    }
+    taskSchedulerRouter = clazz;
+    taskSchedulerRouterLocalExecutorField = field;
+    taskSchedulerRouterDestoryMethod = method;
+  }
+
+  @Setter
+  private ScheduledTaskRegistrar taskRegistrar;
+
   @Override
   protected Duration getStrategyShutdownDelay() {
     // Handling of @Scheduled annotation should stop immediately.
     // Those are not related to clients' requests.
-    
+
     return Duration.ofSeconds(0);
   }
 
@@ -30,6 +63,8 @@ protected Mono<Void> shutdownResourceGraceful(@NonNull TaskScheduler resource) {
     return Mono.fromRunnable(() -> {
       if (resource instanceof Executor) {
         ExecutorShutdownUtils.shutdownExecutor((Executor) resource, false);
+      } else if (taskSchedulerRouter != null && taskSchedulerRouter.isInstance(resource)) {
+        shutdownTaskSchedulerRouter(resource);
       } else {
         log.info("Shutting down unknown task scheduler '{}' using it's 'shutdown()' method.", resource);
         ExecutorShutdownUtils.shutdownExecutorWithReflection(resource, true);
@@ -42,6 +77,8 @@ protected Mono<Void> shutdownResourceForced(@NonNull TaskScheduler resource) {
     return Mono.fromRunnable(() -> {
       if (resource instanceof Executor) {
         ExecutorShutdownUtils.shutdownExecutorForced((Executor) resource);
+      } else if (taskSchedulerRouter != null && taskSchedulerRouter.isInstance(resource)) {
+        shutdownTaskSchedulerRouterForced(resource);
       } else {
         log.warn("Unknown TaskScheduler to force shutdown: {}. Skipping.", resource.getClass());
       }
@@ -65,11 +102,32 @@ protected Mono<Boolean> getResourceForcedTerminationStatus(TaskScheduler resourc
     return getResourceGracefulTerminationStatus(resource);
   }
 
-  /**
-   * Will shut down gracefully added resources during app shutdown.
-   * @param taskScheduler TaskScheduler to shut down gracefully.
-   */
-  public void addTaskScheduler(TaskScheduler taskScheduler) {
-    addResource(taskScheduler);
+  @Override
+  public Set<TaskScheduler> getResourcesForShutdown() {
+    Set<TaskScheduler> resourcesForShutdown = super.getResourcesForShutdown();
+    if (taskRegistrar != null) {
+      resourcesForShutdown.add(taskRegistrar.getScheduler());
+    }
+    return resourcesForShutdown;
+  }
+
+  private static void shutdownTaskSchedulerRouter(TaskScheduler scheduler) {
+    try {
+      Executor executor = (Executor) taskSchedulerRouterLocalExecutorField.get(scheduler);
+      if (executor != null) {
+        ExecutorShutdownUtils.shutdownExecutor(executor, true);
+      }
+    } catch (IllegalAccessException e) {
+      log.warn("Couldn't shutdown TaskSchedulerRouter during graceful shutdown", e);
+    }
+  }
+
+  private void shutdownTaskSchedulerRouterForced(TaskScheduler scheduler) {
+    try {
+      taskSchedulerRouterDestoryMethod.invoke(scheduler);
+    } catch (IllegalAccessException | InvocationTargetException e) {
+      log.warn("Couldn't force shutdown TaskSchedulerRouter during graceful shutdown", e);
+    }
   }
+
 }

core/src/main/java/com/transferwise/common/gracefulshutdown/utils/ExecutorShutdownUtils.java

@@ -1,5 +1,6 @@
 package com.transferwise.common.gracefulshutdown.utils;
 
+import com.transferwise.common.baseutils.concurrency.ScheduledTaskExecutor;
 import java.lang.reflect.Method;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
@@ -19,6 +20,7 @@ public abstract class ExecutorShutdownUtils {
    *                                     contact SRE
    */
   public static void shutdownExecutor(Executor executor, boolean askToReportOnUnknownExecutor) {
+
     if (executor instanceof ThreadPoolTaskScheduler) {
       log.info("Shutting down thread pool task scheduler '{}'.", executor);
       shutdownThreadPoolTaskScheduler((ThreadPoolTaskScheduler) executor);
@@ -139,4 +141,5 @@ private static void shutdownScheduledThreadPoolExecutor(ScheduledThreadPoolExecu
     scheduledThreadPoolExecutor.getQueue().clear();
     scheduledThreadPoolExecutor.shutdown();
   }
+
 }

core/src/test/java/com/transferwise/common/gracefulshutdown/AlternativeSchedulingShutdownerIntTest.java → core/src/test/java/com/transferwise/common/gracefulshutdown/CustomSchedulerShutdownerIntTest.java

@@ -6,15 +6,15 @@
 import com.transferwise.common.gracefulshutdown.strategies.TaskSchedulersGracefulShutdownStrategy;
 import com.transferwise.common.gracefulshutdown.test.BaseTestEnvironment;
 import com.transferwise.common.gracefulshutdown.test.TestApplication;
-import com.transferwise.common.gracefulshutdown.test.TestBApplication;
+import com.transferwise.common.gracefulshutdown.testcustomscheduler.CustomSchedulerConfiguration;
 import lombok.SneakyThrows;
 import org.junit.jupiter.api.Test;
 import org.springframework.beans.factory.annotation.Autowired;
-import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.test.context.ContextConfiguration;
 
 @BaseTestEnvironment
-@SpringBootTest(classes = {TestBApplication.class})
-class AlternativeSchedulingShutdownerIntTest {
+@ContextConfiguration(classes = {CustomSchedulerConfiguration.class})
+class CustomSchedulerShutdownerIntTest {
 
   @Autowired
   private GracefulShutdowner gracefulShutdowner;

core/src/test/java/com/transferwise/common/gracefulshutdown/GracefulShutdownerIntTest.java

@@ -14,6 +14,7 @@
 import org.junit.jupiter.api.Test;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.actuate.health.Status;
+import org.springframework.context.ApplicationContext;
 import org.springframework.test.context.ActiveProfiles;
 
 @BaseTestEnvironment
@@ -37,6 +38,8 @@ class GracefulShutdownerIntTest {
   private TaskSchedulersGracefulShutdownStrategy taskSchedulersGracefulShutdownStrategy;
   @Autowired
   private TestApplication testApplication;
+  @Autowired
+  private ApplicationContext context;
 
   @Test
   @SneakyThrows

core/src/test/java/com/transferwise/common/gracefulshutdown/strategies/TaskSchedulersGracefulShutdownStrategyTest.java

@@ -53,7 +53,7 @@ public void shutdown_invoked_on_external_added_classes() {
     );
     ThreadPoolTaskScheduler threadPoolTaskScheduler = new ThreadPoolTaskScheduler();
     threadPoolTaskScheduler.initialize();
-    strategy.addTaskScheduler(threadPoolTaskScheduler);
+    strategy.addResource(threadPoolTaskScheduler);
 
     // WHEN
     strategy.prepareForShutdown();
@@ -74,7 +74,7 @@ public void shutdown_timeout_is_applied_and_called_shutdownNow() {
     );
     ThreadPoolTaskScheduler threadPoolTaskScheduler = new ThreadPoolTaskScheduler();
     threadPoolTaskScheduler.initialize();
-    strategy.addTaskScheduler(threadPoolTaskScheduler);
+    strategy.addResource(threadPoolTaskScheduler);
     threadPoolTaskScheduler.execute(() -> {
       try {
         Thread.sleep(checkMaxWaitTime.toMillis());

core/src/test/java/com/transferwise/common/gracefulshutdown/test/TestBApplication.java → core/src/test/java/com/transferwise/common/gracefulshutdown/testcustomscheduler/CustomSchedulerConfiguration.java

@@ -1,15 +1,13 @@
-package com.transferwise.common.gracefulshutdown.test;
+package com.transferwise.common.gracefulshutdown.testcustomscheduler;
 
 import java.util.concurrent.Executors;
-import org.springframework.boot.autoconfigure.SpringBootApplication;
 import org.springframework.context.annotation.Bean;
-import org.springframework.scheduling.annotation.EnableScheduling;
+import org.springframework.context.annotation.Configuration;
 import org.springframework.scheduling.annotation.SchedulingConfigurer;
 
 // More customized scheduling is set up.
-@SpringBootApplication
-@EnableScheduling
-public class TestBApplication {
+@Configuration
+public class CustomSchedulerConfiguration {
 
   @Bean
   public SchedulingConfigurer mySchedulingConfigurer() {