Check termination status for TaskSchedulerRouter instances correctly (#…

…43)

CHANGELOG.md

@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres
 to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [2.14.4] - 2024-04-10
+
+### Fixed
+
+* Check termination status for `TaskSchedulerRouter` instances correctly.
+
 ## [2.14.3] - 2024-02-22
 
 ### Changed

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

@@ -14,7 +14,6 @@
 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
@@ -94,6 +93,8 @@ protected Mono<Boolean> getResourceGracefulTerminationStatus(TaskScheduler resou
     return Mono.fromCallable(() -> {
       if (resource instanceof Executor) {
         return ExecutorShutdownUtils.isTerminated((Executor) resource);
+      } else if (taskSchedulerRouter != null && taskSchedulerRouter.isInstance(resource)) {
+        return getTaskSchedulerRouterTerminationStatus(resource);
       } else {
         log.warn("Unknown TaskScheduler to check termination: {}. Return true.", resource.getClass());
         return true;
@@ -134,4 +135,16 @@ private void shutdownTaskSchedulerRouterForced(TaskScheduler scheduler) {
     }
   }
 
+  private boolean getTaskSchedulerRouterTerminationStatus(TaskScheduler scheduler) {
+    try {
+      Executor executor = (Executor) taskSchedulerRouterLocalExecutorField.get(scheduler);
+      if (executor != null) {
+        return ExecutorShutdownUtils.isTerminated(executor);
+      }
+    } catch (IllegalAccessException e) {
+      log.warn("Couldn't get TaskSchedulerRouter termination status during graceful shutdown", e);
+    }
+    return true;
+  }
+
 }

gradle.properties

@@ -1 +1 @@
-version=2.14.3
+version=2.14.4