Address the missing packages of TaskExecutorCustomizer and TaskSchedulerCustomizer

Author: Binbing Hou

genie-agent/src/main/java/com/netflix/genie/agent/spring/autoconfigure/AgentAutoConfiguration.java

@@ -23,10 +23,10 @@
 import com.netflix.genie.agent.utils.locks.impl.FileLockFactory;
 import com.netflix.genie.common.internal.util.GenieHostInfo;
 import com.netflix.genie.common.internal.util.HostnameUtil;
+import org.springframework.beans.factory.config.BeanPostProcessor;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
+import org.springframework.boot.autoconfigure.task.TaskExecutionAutoConfiguration;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
-import org.springframework.boot.task.TaskExecutorCustomizer;
-import org.springframework.boot.task.TaskSchedulerCustomizer;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.Lazy;
@@ -51,6 +51,9 @@
 )
 public class AgentAutoConfiguration {
 
+    // Common bean names used by Spring Boot
+    private static final String TASK_SCHEDULER_BEAN_NAME = "taskScheduler";
+
     /**
      * Provide a bean of type {@link GenieHostInfo} if none already exists.
      *
@@ -151,34 +154,51 @@ public TaskScheduler heartBeatServiceTaskScheduler(final AgentProperties agentPr
     }
 
     /**
-     * Customizer for Spring's task executor.
+     * Bean post processor to customize the task executor created by Spring Boot.
+     * This provides the functionality of TaskExecutorCustomizer which is no longer available in AWS Cloud v3.
      *
      * @param agentProperties the agent properties
-     * @return a customizer for the task executor
+     * @return A BeanPostProcessor that customizes task executors
      */
-    @Bean
-    TaskExecutorCustomizer taskExecutorCustomizer(final AgentProperties agentProperties) {
-        return taskExecutor -> {
-            taskExecutor.setWaitForTasksToCompleteOnShutdown(true);
-            taskExecutor.setAwaitTerminationSeconds(
-                (int) agentProperties.getShutdown().getSystemExecutorLeeway().getSeconds()
-            );
+    @Bean(name = "taskExecutorCustomizer")
+    public BeanPostProcessor taskExecutorCustomizer(final AgentProperties agentProperties) {
+        return new BeanPostProcessor() {
+            @Override
+            public Object postProcessAfterInitialization(final Object bean, final String beanName) {
+                // Customize the application task executor
+                if (bean instanceof ThreadPoolTaskExecutor executor && TaskExecutionAutoConfiguration.APPLICATION_TASK_EXECUTOR_BEAN_NAME.equals(beanName)) {
+                    executor.setWaitForTasksToCompleteOnShutdown(true);
+                    executor.setAwaitTerminationSeconds(
+                        (int) agentProperties.getShutdown().getSystemExecutorLeeway().getSeconds()
+                    );
+                }
+                return bean;
+            }
         };
     }
 
     /**
-     * Customizer for Spring's task scheduler.
+     * Bean post processor to customize the task scheduler created by Spring Boot.
+     * This provides the functionality of TaskSchedulerCustomizer which is no longer available in AWS Cloud v3.
      *
      * @param agentProperties the agent properties
-     * @return a customizer for the task scheduler
+     * @return A BeanPostProcessor that customizes task schedulers
      */
-    @Bean
-    TaskSchedulerCustomizer taskSchedulerCustomizer(final AgentProperties agentProperties) {
-        return taskScheduler -> {
-            taskScheduler.setWaitForTasksToCompleteOnShutdown(true);
-            taskScheduler.setAwaitTerminationSeconds(
-                (int) agentProperties.getShutdown().getSystemSchedulerLeeway().getSeconds()
-            );
+    @Bean(name = "taskSchedulerCustomizer")
+    public BeanPostProcessor taskSchedulerCustomizer(final AgentProperties agentProperties) {
+        return new BeanPostProcessor() {
+            @Override
+            public Object postProcessAfterInitialization(final Object bean, final String beanName) {
+                // Customize the task scheduler
+                if (bean instanceof ThreadPoolTaskScheduler scheduler
+                    && TASK_SCHEDULER_BEAN_NAME.equals(beanName)) {
+                    scheduler.setWaitForTasksToCompleteOnShutdown(true);
+                    scheduler.setAwaitTerminationSeconds(
+                        (int) agentProperties.getShutdown().getSystemSchedulerLeeway().getSeconds()
+                    );
+                }
+                return bean;
+            }
         };
     }
 }