diff --git a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/async/TraceAsyncDefaultAutoConfiguration.java b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/async/TraceAsyncDefaultAutoConfiguration.java index 18fb2fbdf..d07995c26 100644 --- a/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/async/TraceAsyncDefaultAutoConfiguration.java +++ b/spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/instrument/async/TraceAsyncDefaultAutoConfiguration.java @@ -92,7 +92,7 @@ public class TraceAsyncDefaultAutoConfiguration { @Override public Executor getAsyncExecutor() { Executor delegate = getDefaultExecutor(); - return new LazyTraceExecutor(this.beanFactory, delegate); + return LazyTraceExecutor.wrap(this.beanFactory, delegate); } /** diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/ExecutorInstrumentor.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/ExecutorInstrumentor.java index 4157fb4fe..24d29bf61 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/ExecutorInstrumentor.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/ExecutorInstrumentor.java @@ -20,10 +20,13 @@ import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Method; import java.lang.reflect.Modifier; import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Executor; import java.util.concurrent.ExecutorService; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.function.Function; import java.util.function.Supplier; import org.aopalliance.aop.Advice; @@ -176,43 +179,43 @@ public class ExecutorInstrumentor { Object createThreadPoolTaskExecutorProxy(Object bean, boolean cglibProxy, ThreadPoolTaskExecutor executor, String beanName) { if (!cglibProxy) { - return new LazyTraceThreadPoolTaskExecutor(this.beanFactory, executor, beanName); + return LazyTraceThreadPoolTaskExecutor.wrap(this.beanFactory, executor, beanName); } return getProxiedObject(bean, beanName, true, executor, - () -> new LazyTraceThreadPoolTaskExecutor(this.beanFactory, executor, beanName)); + () -> LazyTraceThreadPoolTaskExecutor.wrap(this.beanFactory, executor, beanName)); } Supplier createThreadPoolTaskSchedulerProxy(ThreadPoolTaskScheduler executor, String beanName) { - return () -> new LazyTraceThreadPoolTaskScheduler(this.beanFactory, executor, beanName); + return () -> LazyTraceThreadPoolTaskScheduler.wrap(this.beanFactory, executor, beanName); } Supplier createScheduledThreadPoolExecutorProxy(ScheduledThreadPoolExecutor executor, String beanName) { - return () -> new LazyTraceScheduledThreadPoolExecutor(executor.getCorePoolSize(), executor.getThreadFactory(), + return () -> LazyTraceScheduledThreadPoolExecutor.wrap(executor.getCorePoolSize(), executor.getThreadFactory(), executor.getRejectedExecutionHandler(), this.beanFactory, executor, beanName); } Object createExecutorServiceProxy(Object bean, boolean cglibProxy, ExecutorService executor, String beanName) { return getProxiedObject(bean, beanName, cglibProxy, executor, () -> { if (executor instanceof ScheduledExecutorService) { - return new TraceableScheduledExecutorService(this.beanFactory, executor, beanName); + return TraceableScheduledExecutorService.wrap(this.beanFactory, executor, beanName); } - return new TraceableExecutorService(this.beanFactory, executor, beanName); + return TraceableExecutorService.wrap(this.beanFactory, executor, beanName); }); } Object createScheduledExecutorServiceProxy(Object bean, boolean cglibProxy, ScheduledExecutorService executor, String beanName) { return getProxiedObject(bean, beanName, cglibProxy, executor, - () -> new TraceableScheduledExecutorService(this.beanFactory, executor, beanName)); + () -> TraceableScheduledExecutorService.wrap(this.beanFactory, executor, beanName)); } Object createAsyncTaskExecutorProxy(Object bean, boolean cglibProxy, AsyncTaskExecutor executor, String beanName) { return getProxiedObject(bean, beanName, cglibProxy, executor, () -> { if (bean instanceof ThreadPoolTaskScheduler) { - return new LazyTraceThreadPoolTaskScheduler(this.beanFactory, (ThreadPoolTaskScheduler) executor, + return LazyTraceThreadPoolTaskScheduler.wrap(this.beanFactory, (ThreadPoolTaskScheduler) executor, beanName); } - return new LazyTraceAsyncTaskExecutor(this.beanFactory, executor, beanName); + return LazyTraceAsyncTaskExecutor.wrap(this.beanFactory, executor, beanName); }); } @@ -259,7 +262,7 @@ public class ExecutorInstrumentor { factory.addAdvice(new ExecutorMethodInterceptor(executor, this.beanFactory, beanName) { @Override Executor executor(BeanFactory beanFactory, Executor executor, String beanName) { - return supplier.get(); + return executorFromCache(beanFactory, executor, beanName, e -> supplier.get()); } }); factory.setTarget(bean); @@ -316,6 +319,8 @@ class ExecutorMethodInterceptor implements MethodInterceptor private final String beanName; + private static final Map CACHE = new ConcurrentHashMap<>(); + ExecutorMethodInterceptor(T delegate, BeanFactory beanFactory, String beanName) { this.delegate = delegate; this.beanFactory = beanFactory; @@ -344,8 +349,15 @@ class ExecutorMethodInterceptor implements MethodInterceptor return ReflectionUtils.findMethod(object.getClass(), method.getName(), method.getParameterTypes()); } + @SuppressWarnings("unchecked") T executor(BeanFactory beanFactory, T executor, String beanName) { - return (T) new LazyTraceExecutor(beanFactory, executor, beanName); + return executorFromCache(beanFactory, executor, beanName, + e -> (T) LazyTraceExecutor.wrap(beanFactory, e, beanName)); + } + + @SuppressWarnings("unchecked") + T executorFromCache(BeanFactory beanFactory, T executor, String beanName, Function function) { + return (T) CACHE.computeIfAbsent(executor, function); } } diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceAsyncCustomizer.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceAsyncCustomizer.java index 883f3c7c8..e7db11f4e 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceAsyncCustomizer.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceAsyncCustomizer.java @@ -47,7 +47,7 @@ public class LazyTraceAsyncCustomizer extends AsyncConfigurerSupport { if (this.delegate.getAsyncExecutor() instanceof LazyTraceExecutor) { return this.delegate.getAsyncExecutor(); } - return new LazyTraceExecutor(this.beanFactory, this.delegate.getAsyncExecutor()); + return LazyTraceExecutor.wrap(this.beanFactory, this.delegate.getAsyncExecutor()); } @Override diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceAsyncTaskExecutor.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceAsyncTaskExecutor.java index 4755ceace..d450c0f41 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceAsyncTaskExecutor.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceAsyncTaskExecutor.java @@ -16,7 +16,9 @@ package org.springframework.cloud.sleuth.instrument.async; +import java.util.Map; import java.util.concurrent.Callable; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Future; import org.apache.commons.logging.Log; @@ -29,6 +31,7 @@ import org.springframework.cloud.sleuth.Tracer; import org.springframework.cloud.sleuth.internal.ContextUtil; import org.springframework.cloud.sleuth.internal.DefaultSpanNamer; import org.springframework.core.task.AsyncTaskExecutor; +import org.springframework.lang.NonNull; /** * {@link AsyncTaskExecutor} that wraps {@link Runnable} and {@link Callable} in a trace @@ -40,6 +43,8 @@ import org.springframework.core.task.AsyncTaskExecutor; // public as most types in this package were documented for use public class LazyTraceAsyncTaskExecutor implements AsyncTaskExecutor { + private static final Map CACHE = new ConcurrentHashMap<>(); + private static final Log log = LogFactory.getLog(LazyTraceAsyncTaskExecutor.class); private final BeanFactory beanFactory; @@ -64,6 +69,28 @@ public class LazyTraceAsyncTaskExecutor implements AsyncTaskExecutor { this.beanName = beanName; } + /** + * Wraps the Executor in a trace instance. + * @param beanFactory bean factory + * @param delegate delegate to wrap + * @param beanName bean name + * @return traced instance + */ + public static LazyTraceAsyncTaskExecutor wrap(BeanFactory beanFactory, @NonNull AsyncTaskExecutor delegate, + String beanName) { + return CACHE.computeIfAbsent(delegate, e -> new LazyTraceAsyncTaskExecutor(beanFactory, delegate, beanName)); + } + + /** + * Wraps the Executor in a trace instance. + * @param beanFactory bean factory + * @param delegate delegate to wrap + * @return traced instance + */ + public static LazyTraceAsyncTaskExecutor wrap(BeanFactory beanFactory, @NonNull AsyncTaskExecutor delegate) { + return CACHE.computeIfAbsent(delegate, e -> new LazyTraceAsyncTaskExecutor(beanFactory, delegate, null)); + } + @Override public void execute(Runnable task) { Runnable taskToRun = task; diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceExecutor.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceExecutor.java index d7e26471d..8394aba5a 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceExecutor.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceExecutor.java @@ -16,6 +16,8 @@ package org.springframework.cloud.sleuth.instrument.async; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Executor; import org.apache.commons.logging.Log; @@ -27,6 +29,7 @@ import org.springframework.cloud.sleuth.SpanNamer; import org.springframework.cloud.sleuth.Tracer; import org.springframework.cloud.sleuth.internal.ContextUtil; import org.springframework.cloud.sleuth.internal.DefaultSpanNamer; +import org.springframework.lang.NonNull; /** * {@link Executor} that wraps {@link Runnable} in a trace representation. @@ -39,6 +42,8 @@ public class LazyTraceExecutor implements Executor { private static final Log log = LogFactory.getLog(LazyTraceExecutor.class); + private static final Map CACHE = new ConcurrentHashMap<>(); + private final BeanFactory beanFactory; private final Executor delegate; @@ -61,6 +66,27 @@ public class LazyTraceExecutor implements Executor { this.beanName = beanName; } + /** + * Wraps the Executor in a trace instance. + * @param beanFactory bean factory + * @param delegate delegate to wrap + * @param beanName bean name + * @return traced instance + */ + public static LazyTraceExecutor wrap(BeanFactory beanFactory, @NonNull Executor delegate, String beanName) { + return CACHE.computeIfAbsent(delegate, e -> new LazyTraceExecutor(beanFactory, delegate, beanName)); + } + + /** + * Wraps the Executor in a trace instance. + * @param beanFactory bean factory + * @param delegate delegate to wrap + * @return traced instance + */ + public static LazyTraceExecutor wrap(BeanFactory beanFactory, @NonNull Executor delegate) { + return CACHE.computeIfAbsent(delegate, e -> new LazyTraceExecutor(beanFactory, delegate, null)); + } + @Override public void execute(Runnable command) { if (ContextUtil.isContextUnusable(this.beanFactory)) { diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceScheduledThreadPoolExecutor.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceScheduledThreadPoolExecutor.java index b8d534f65..605757545 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceScheduledThreadPoolExecutor.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceScheduledThreadPoolExecutor.java @@ -20,8 +20,10 @@ import java.lang.reflect.Method; import java.util.ArrayList; import java.util.Collection; import java.util.List; +import java.util.Map; import java.util.concurrent.BlockingQueue; import java.util.concurrent.Callable; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutionException; import java.util.concurrent.Future; import java.util.concurrent.RejectedExecutionHandler; @@ -42,6 +44,7 @@ import org.springframework.cloud.sleuth.SpanNamer; import org.springframework.cloud.sleuth.Tracer; import org.springframework.cloud.sleuth.internal.ContextUtil; import org.springframework.cloud.sleuth.internal.DefaultSpanNamer; +import org.springframework.lang.NonNull; import org.springframework.util.ReflectionUtils; /** @@ -56,6 +59,8 @@ class LazyTraceScheduledThreadPoolExecutor extends ScheduledThreadPoolExecutor { private static final Log log = LogFactory.getLog(LazyTraceScheduledThreadPoolExecutor.class); + private static final Map CACHE = new ConcurrentHashMap<>(); + private final BeanFactory beanFactory; private final ScheduledThreadPoolExecutor delegate; @@ -175,6 +180,19 @@ class LazyTraceScheduledThreadPoolExecutor extends ScheduledThreadPoolExecutor { this.newTaskForCallable = makeAccessibleIfNotNullAndOverridden(newTaskForCallable); } + static LazyTraceScheduledThreadPoolExecutor wrap(int corePoolSize, BeanFactory beanFactory, + @NonNull ScheduledThreadPoolExecutor delegate, String beanName) { + return CACHE.computeIfAbsent(delegate, + e -> new LazyTraceScheduledThreadPoolExecutor(corePoolSize, beanFactory, delegate, beanName)); + } + + static LazyTraceScheduledThreadPoolExecutor wrap(int corePoolSize, ThreadFactory threadFactory, + RejectedExecutionHandler handler, BeanFactory beanFactory, @NonNull ScheduledThreadPoolExecutor delegate, + String beanName) { + return CACHE.computeIfAbsent(delegate, e -> new LazyTraceScheduledThreadPoolExecutor(corePoolSize, + threadFactory, handler, beanFactory, delegate, beanName)); + } + private Runnable traceRunnableWhenContextReady(Runnable delegate) { if (isContextUnusable()) { return delegate; diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskExecutor.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskExecutor.java index 0a5ad92ee..b678784fc 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskExecutor.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskExecutor.java @@ -16,7 +16,9 @@ package org.springframework.cloud.sleuth.instrument.async; +import java.util.Map; import java.util.concurrent.Callable; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Future; import java.util.concurrent.RejectedExecutionHandler; import java.util.concurrent.ThreadFactory; @@ -32,6 +34,7 @@ import org.springframework.cloud.sleuth.Tracer; import org.springframework.cloud.sleuth.internal.ContextUtil; import org.springframework.cloud.sleuth.internal.DefaultSpanNamer; import org.springframework.core.task.TaskDecorator; +import org.springframework.lang.NonNull; import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor; import org.springframework.util.concurrent.ListenableFuture; @@ -46,6 +49,8 @@ public class LazyTraceThreadPoolTaskExecutor extends ThreadPoolTaskExecutor { private static final Log log = LogFactory.getLog(LazyTraceThreadPoolTaskExecutor.class); + private static final Map CACHE = new ConcurrentHashMap<>(); + private final BeanFactory beanFactory; private final ThreadPoolTaskExecutor delegate; @@ -68,6 +73,30 @@ public class LazyTraceThreadPoolTaskExecutor extends ThreadPoolTaskExecutor { this.beanName = beanName; } + /** + * Wraps the Executor in a trace instance. + * @param beanFactory bean factory + * @param delegate delegate to wrap + * @param beanName bean name + * @return traced instance + */ + public static LazyTraceThreadPoolTaskExecutor wrap(BeanFactory beanFactory, + @NonNull ThreadPoolTaskExecutor delegate, String beanName) { + return CACHE.computeIfAbsent(delegate, + e -> new LazyTraceThreadPoolTaskExecutor(beanFactory, delegate, beanName)); + } + + /** + * Wraps the Executor in a trace instance. + * @param beanFactory bean factory + * @param delegate delegate to wrap + * @return traced instance + */ + public static LazyTraceThreadPoolTaskExecutor wrap(BeanFactory beanFactory, + @NonNull ThreadPoolTaskExecutor delegate) { + return CACHE.computeIfAbsent(delegate, e -> new LazyTraceThreadPoolTaskExecutor(beanFactory, delegate, null)); + } + @Override public void execute(Runnable task) { this.delegate.execute(wrap(task)); diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskScheduler.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskScheduler.java index 136ef830a..ff1dd6fcb 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskScheduler.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskScheduler.java @@ -20,7 +20,9 @@ import java.lang.reflect.Method; import java.time.Duration; import java.time.Instant; import java.util.Date; +import java.util.Map; import java.util.concurrent.Callable; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; import java.util.concurrent.RejectedExecutionHandler; @@ -39,6 +41,7 @@ import org.springframework.cloud.sleuth.SpanNamer; import org.springframework.cloud.sleuth.Tracer; import org.springframework.cloud.sleuth.internal.ContextUtil; import org.springframework.cloud.sleuth.internal.DefaultSpanNamer; +import org.springframework.lang.NonNull; import org.springframework.lang.Nullable; import org.springframework.scheduling.Trigger; import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler; @@ -59,6 +62,8 @@ class LazyTraceThreadPoolTaskScheduler extends ThreadPoolTaskScheduler { private static final Log log = LogFactory.getLog(LazyTraceThreadPoolTaskScheduler.class); + private static final Map CACHE = new ConcurrentHashMap<>(); + private final BeanFactory beanFactory; private final ThreadPoolTaskScheduler delegate; @@ -97,6 +102,19 @@ class LazyTraceThreadPoolTaskScheduler extends ThreadPoolTaskScheduler { makeAccessibleIfNotNull(this.getDefaultThreadNamePrefix); } + /** + * Wraps the Executor in a trace instance. + * @param beanFactory bean factory + * @param delegate delegate to wrap + * @param beanName bean name + * @return traced instance + */ + static LazyTraceThreadPoolTaskScheduler wrap(BeanFactory beanFactory, @NonNull ThreadPoolTaskScheduler delegate, + String beanName) { + return CACHE.computeIfAbsent(delegate, + e -> new LazyTraceThreadPoolTaskScheduler(beanFactory, delegate, beanName)); + } + private void makeAccessibleIfNotNull(Method method) { if (method != null) { ReflectionUtils.makeAccessible(method); @@ -147,7 +165,7 @@ class LazyTraceThreadPoolTaskScheduler extends ThreadPoolTaskScheduler { if (executorService instanceof TraceableScheduledExecutorService) { return executorService; } - return new TraceableExecutorService(this.beanFactory, executorService, this.beanName); + return TraceableExecutorService.wrap(this.beanFactory, executorService, this.beanName); } private RejectedExecutionHandler traceRejectedExecutionHandler(RejectedExecutionHandler rejectedExecutionHandler) { @@ -172,14 +190,14 @@ class LazyTraceThreadPoolTaskScheduler extends ThreadPoolTaskScheduler { if (executorService instanceof TraceableScheduledExecutorService) { return executorService; } - return new TraceableScheduledExecutorService(this.beanFactory, executorService, this.beanName); + return TraceableScheduledExecutorService.wrap(this.beanFactory, executorService, this.beanName); } @Override public ScheduledExecutorService getScheduledExecutor() throws IllegalStateException { ScheduledExecutorService executor = this.delegate.getScheduledExecutor(); return executor instanceof TraceableScheduledExecutorService ? executor - : new TraceableScheduledExecutorService(this.beanFactory, executor, this.beanName); + : TraceableScheduledExecutorService.wrap(this.beanFactory, executor, this.beanName); } @Override @@ -188,7 +206,7 @@ class LazyTraceThreadPoolTaskScheduler extends ThreadPoolTaskScheduler { if (executor instanceof LazyTraceScheduledThreadPoolExecutor) { return executor; } - return new LazyTraceScheduledThreadPoolExecutor(executor.getCorePoolSize(), executor.getThreadFactory(), + return LazyTraceScheduledThreadPoolExecutor.wrap(executor.getCorePoolSize(), executor.getThreadFactory(), executor.getRejectedExecutionHandler(), this.beanFactory, executor, this.beanName); } diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorService.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorService.java index 0b36394d2..cc3ad4aad 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorService.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorService.java @@ -19,7 +19,9 @@ package org.springframework.cloud.sleuth.instrument.async; import java.util.ArrayList; import java.util.Collection; import java.util.List; +import java.util.Map; import java.util.concurrent.Callable; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; @@ -40,6 +42,8 @@ import org.springframework.cloud.sleuth.internal.ContextUtil; // public as most types in this package were documented for use public class TraceableExecutorService implements ExecutorService { + private static final Map CACHE = new ConcurrentHashMap<>(); + final ExecutorService delegate; final String spanName; @@ -60,6 +64,27 @@ public class TraceableExecutorService implements ExecutorService { this.spanName = spanName; } + /** + * Wraps the Executor in a trace instance. + * @param beanFactory bean factory + * @param delegate delegate to wrap + * @param beanName bean name + * @return traced instance + */ + public static TraceableExecutorService wrap(BeanFactory beanFactory, ExecutorService delegate, String beanName) { + return CACHE.computeIfAbsent(delegate, e -> new TraceableExecutorService(beanFactory, delegate, beanName)); + } + + /** + * Wraps the Executor in a trace instance. + * @param beanFactory bean factory + * @param delegate delegate to wrap + * @return traced instance + */ + public static TraceableExecutorService wrap(BeanFactory beanFactory, ExecutorService delegate) { + return CACHE.computeIfAbsent(delegate, e -> new TraceableExecutorService(beanFactory, delegate, null)); + } + @Override public void execute(Runnable command) { this.delegate.execute(ContextUtil.isContextUnusable(this.beanFactory) ? command diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableScheduledExecutorService.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableScheduledExecutorService.java index abccba742..eb98ca685 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableScheduledExecutorService.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableScheduledExecutorService.java @@ -16,7 +16,9 @@ package org.springframework.cloud.sleuth.instrument.async; +import java.util.Map; import java.util.concurrent.Callable; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutorService; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; @@ -34,6 +36,8 @@ import org.springframework.cloud.sleuth.internal.ContextUtil; // public as most types in this package were documented for use public class TraceableScheduledExecutorService extends TraceableExecutorService implements ScheduledExecutorService { + private static final Map CACHE = new ConcurrentHashMap<>(); + public TraceableScheduledExecutorService(BeanFactory beanFactory, final ExecutorService delegate) { super(beanFactory, delegate); } @@ -42,6 +46,29 @@ public class TraceableScheduledExecutorService extends TraceableExecutorService super(beanFactory, delegate, beanName); } + /** + * Wraps the Executor in a trace instance. + * @param beanFactory bean factory + * @param delegate delegate to wrap + * @param beanName bean name + * @return traced instance + */ + public static TraceableScheduledExecutorService wrap(BeanFactory beanFactory, ExecutorService delegate, + String beanName) { + return CACHE.computeIfAbsent(delegate, + e -> new TraceableScheduledExecutorService(beanFactory, delegate, beanName)); + } + + /** + * Wraps the Executor in a trace instance. + * @param beanFactory bean factory + * @param delegate delegate to wrap + * @return traced instance + */ + public static TraceableScheduledExecutorService wrap(BeanFactory beanFactory, ExecutorService delegate) { + return CACHE.computeIfAbsent(delegate, e -> new TraceableScheduledExecutorService(beanFactory, delegate, null)); + } + private ScheduledExecutorService getScheduledExecutorService() { return (ScheduledExecutorService) this.delegate; } diff --git a/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceAsyncCustomizerTest.java b/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceAsyncCustomizerTest.java index 0b5992738..48a3d0fad 100644 --- a/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceAsyncCustomizerTest.java +++ b/spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceAsyncCustomizerTest.java @@ -17,10 +17,12 @@ package org.springframework.cloud.sleuth.instrument.async; import java.util.concurrent.Executor; +import java.util.concurrent.Executors; import org.assertj.core.api.BDDAssertions; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.BDDMockito; import org.mockito.InjectMocks; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; @@ -45,6 +47,8 @@ public class LazyTraceAsyncCustomizerTest { @Test public void should_wrap_async_executor_in_trace_version() throws Exception { + BDDMockito.given(this.asyncConfigurer.getAsyncExecutor()).willReturn(Executors.newSingleThreadExecutor()); + Executor executor = this.lazyTraceAsyncCustomizer.getAsyncExecutor(); BDDAssertions.then(executor).isExactlyInstanceOf(LazyTraceExecutor.class); diff --git a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskSchedulerTests.java b/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskSchedulerTests.java index 3aa01c258..aae098478 100644 --- a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskSchedulerTests.java +++ b/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskSchedulerTests.java @@ -20,6 +20,7 @@ import java.time.Duration; import java.time.Instant; import java.util.Date; import java.util.concurrent.Callable; +import java.util.concurrent.Executors; import java.util.concurrent.RejectedExecutionHandler; import java.util.concurrent.ThreadFactory; @@ -94,6 +95,8 @@ public abstract class LazyTraceThreadPoolTaskSchedulerTests implements TestTraci @Test public void getScheduledExecutor() { + BDDMockito.given(this.delegate.getScheduledExecutor()).willReturn(Executors.newScheduledThreadPool(1)); + this.executor.getScheduledExecutor(); BDDMockito.then(this.delegate).should().getScheduledExecutor();