From ab385a77bb00fdec11cc78c716901d5cc0d6b9bc Mon Sep 17 00:00:00 2001
From: Peeter Karolin <43540664+tw-peeterkarolin@users.noreply.github.com>
Date: Wed, 10 Apr 2024 12:52:49 +0300
Subject: [PATCH] Check termination status for TaskSchedulerRouter instances
 correctly (#43)

---
 CHANGELOG.md                                      |  6 ++++++
 .../TaskSchedulersGracefulShutdownStrategy.java   | 15 ++++++++++++++-
 gradle.properties                                 |  2 +-
 3 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0755cdd..bfb778b 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/core/src/main/java/com/transferwise/common/gracefulshutdown/strategies/TaskSchedulersGracefulShutdownStrategy.java b/core/src/main/java/com/transferwise/common/gracefulshutdown/strategies/TaskSchedulersGracefulShutdownStrategy.java
index 6f3b0e4..cac089a 100644
--- a/core/src/main/java/com/transferwise/common/gracefulshutdown/strategies/TaskSchedulersGracefulShutdownStrategy.java
+++ b/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;
+  }
+
 }
diff --git a/gradle.properties b/gradle.properties
index f57b13b..1c546c8 100644
--- a/gradle.properties
+++ b/gradle.properties
@@ -1 +1 @@
-version=2.14.3
+version=2.14.4