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 6d0470a41a42277170413e0d65a32ab4a93febb2..a99f41ad5948fb6e1371ac9c4d3f57b34d6ab64d 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);