From a5d3668a56118265e300d489823974ef7dac4caf Mon Sep 17 00:00:00 2001
From: liyang <249007302@qq.com>
Date: Wed, 3 Apr 2024 12:22:21 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E8=A7=A3=E5=86=B3Spring=20=E4=B8=AD?=
 =?UTF-8?q?=E7=9A=84=20ThreadPoolTaskExecutor=20=E8=A2=AB=E6=A1=86?=
 =?UTF-8?q?=E6=9E=B6=E7=AE=A1=E7=90=86=EF=BC=8C@Bean=20=E5=AE=9A=E4=B9=89?=
 =?UTF-8?q?=E6=97=B6=E5=8A=A0=20@DynamicTp=20=E6=B3=A8=E8=A7=A3=E5=90=8E?=
 =?UTF-8?q?=EF=BC=8CThreadPoolTaskExecutor=E4=B8=AD=E5=AE=9A=E4=B9=89?=
 =?UTF-8?q?=E7=9A=84=E7=BA=BF=E7=A8=8B=E6=B1=A0=E8=A3=85=E9=A5=B0=E5=99=A8?=
 =?UTF-8?q?=E5=A4=B1=E6=95=88=E9=97=AE=E9=A2=98=20=EF=BC=88https://gitee.c?=
 =?UTF-8?q?om/dromara/dynamic-tp/issues/I9D31H=EF=BC=89?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../core/spring/DtpPostProcessor.java         | 21 ++++++++++++++++++-
 1 file changed, 20 insertions(+), 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
index 6d0470a4..a99f41ad 100644
--- 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.DynamicTp;
 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.config.BeanPostProcessor;
 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 @@ public class DtpPostProcessor implements BeanPostProcessor, BeanFactoryAware, Pr
 
     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);
-- 
Gitee