Merge branch 'master' of gitee.com:dromara/dynamic-tp

dromara · Apr 3, 2024 · 82dddae · 82dddae
2 parents 56cfa02 + 5c2776c
commit 82dddae
Showing 1 changed file with 20 additions and 1 deletion.
diff --git a/core/src/main/java/org/dromara/dynamictp/core/spring/DtpPostProcessor.java b/core/src/main/java/org/dromara/dynamictp/core/spring/DtpPostProcessor.java
@@ -17,6 +17,7 @@
 
 package org.dromara.dynamictp.core.spring;
 
+import com.google.common.collect.Lists;
 import lombok.extern.slf4j.Slf4j;
 import lombok.val;
 import org.apache.commons.lang3.StringUtils;
@@ -31,6 +32,7 @@
 import org.dromara.dynamictp.core.support.ExecutorWrapper;
 import org.dromara.dynamictp.core.support.ScheduledThreadPoolExecutorProxy;
 import org.dromara.dynamictp.core.support.ThreadPoolExecutorProxy;
+import org.dromara.dynamictp.core.support.task.wrapper.TaskWrapper;
 import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.BeanFactory;
 import org.springframework.beans.factory.BeanFactoryAware;
@@ -41,6 +43,7 @@
 import org.springframework.beans.factory.support.DefaultListableBeanFactory;
 import org.springframework.core.Ordered;
 import org.springframework.core.PriorityOrdered;
+import org.springframework.core.task.TaskDecorator;
 import org.springframework.core.type.MethodMetadata;
 import org.springframework.lang.NonNull;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
@@ -139,8 +142,24 @@ private Object registerAndReturnCommon(Object bean, String beanName) {
 
     private Object doRegisterAndReturnCommon(Object bean, String poolName) {
         if (bean instanceof ThreadPoolTaskExecutor) {
-            val proxy = newProxy(poolName, ((ThreadPoolTaskExecutor) bean).getThreadPoolExecutor());
+            ThreadPoolTaskExecutor poolTaskExecutor = (ThreadPoolTaskExecutor) bean;
+            val proxy = newProxy(poolName, poolTaskExecutor.getThreadPoolExecutor());
             try {
+                Object taskDecorator = ReflectionUtil.getFieldValue("taskDecorator", poolTaskExecutor);
+                if (Objects.nonNull(taskDecorator)) {
+                    TaskWrapper taskWrapper = (taskDecorator instanceof TaskWrapper) ? (TaskWrapper)taskDecorator : new TaskWrapper() {
+                        @Override
+                        public String name() {
+                            return taskDecorator.getClass().getName();
+                        }
+
+                        @Override
+                        public Runnable wrap(Runnable runnable) {
+                            return ((TaskDecorator) taskDecorator).decorate(runnable);
+                        }
+                    };
+                    ReflectionUtil.setFieldValue("taskWrappers",proxy, Lists.newArrayList(taskWrapper));
+                }
                 ReflectionUtil.setFieldValue("threadPoolExecutor", bean, proxy);
             } catch (IllegalAccessException ignored) { }
             DtpRegistry.registerExecutor(new ExecutorWrapper(poolName, proxy), REGISTER_SOURCE);