diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/ExecutorBeanPostProcessor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/ExecutorBeanPostProcessor.java index 449f47bf4..e56686925 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/ExecutorBeanPostProcessor.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/ExecutorBeanPostProcessor.java @@ -34,6 +34,7 @@ import org.springframework.aop.framework.ProxyFactoryBean; import org.springframework.beans.BeansException; import org.springframework.beans.factory.BeanFactory; import org.springframework.beans.factory.config.BeanPostProcessor; +import org.springframework.core.task.AsyncTaskExecutor; import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor; import org.springframework.util.ReflectionUtils; @@ -67,7 +68,8 @@ class ExecutorBeanPostProcessor implements BeanPostProcessor { @Override public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException { - if (bean instanceof ThreadPoolTaskExecutor) { + if (bean instanceof ThreadPoolTaskExecutor + && !(bean instanceof LazyTraceThreadPoolTaskExecutor)) { if (isProxyNeeded(beanName)) { return wrapThreadPoolTaskExecutor(bean); } @@ -75,7 +77,8 @@ class ExecutorBeanPostProcessor implements BeanPostProcessor { log.info("Not instrumenting bean " + beanName); } } - else if (bean instanceof ExecutorService) { + else if (bean instanceof ExecutorService + && !(bean instanceof TraceableExecutorService)) { if (isProxyNeeded(beanName)) { return wrapExecutorService(bean); } @@ -83,7 +86,16 @@ class ExecutorBeanPostProcessor implements BeanPostProcessor { log.info("Not instrumenting bean " + beanName); } } - else if (bean instanceof Executor) { + else if (bean instanceof AsyncTaskExecutor + && !(bean instanceof LazyTraceAsyncTaskExecutor)) { + if (isProxyNeeded(beanName)) { + return wrapAsyncTaskExecutor(bean); + } + else { + log.info("Not instrumenting bean " + beanName); + } + } + else if (bean instanceof Executor && !(bean instanceof LazyTraceExecutor)) { return wrapExecutor(bean); } return bean; @@ -128,6 +140,13 @@ class ExecutorBeanPostProcessor implements BeanPostProcessor { return createExecutorServiceProxy(bean, cglibProxy, executor); } + private Object wrapAsyncTaskExecutor(Object bean) { + boolean classFinal = Modifier.isFinal(bean.getClass().getModifiers()); + boolean cglibProxy = !classFinal; + AsyncTaskExecutor executor = (AsyncTaskExecutor) bean; + return createAsyncTaskExecutorProxy(bean, cglibProxy, executor); + } + boolean isProxyNeeded(String beanName) { SleuthAsyncProperties sleuthAsyncProperties = asyncConfigurationProperties(); return !sleuthAsyncProperties.getIgnoredBeans().contains(beanName); @@ -145,6 +164,12 @@ class ExecutorBeanPostProcessor implements BeanPostProcessor { () -> new TraceableExecutorService(this.beanFactory, executor)); } + Object createAsyncTaskExecutorProxy(Object bean, boolean cglibProxy, + AsyncTaskExecutor executor) { + return getProxiedObject(bean, cglibProxy, executor, + () -> new LazyTraceAsyncTaskExecutor(this.beanFactory, executor)); + } + private Object getProxiedObject(Object bean, boolean cglibProxy, Executor executor, Supplier supplier) { ProxyFactoryBean factory = new ProxyFactoryBean(); diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceAsyncTaskExecutor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceAsyncTaskExecutor.java new file mode 100644 index 000000000..4a61a1846 --- /dev/null +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceAsyncTaskExecutor.java @@ -0,0 +1,120 @@ +/* + * Copyright 2013-2019 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.instrument.async; + +import java.util.concurrent.Callable; +import java.util.concurrent.Future; + +import brave.Tracing; +import org.apache.commons.logging.Log; +import org.apache.commons.logging.LogFactory; + +import org.springframework.beans.factory.BeanFactory; +import org.springframework.beans.factory.NoSuchBeanDefinitionException; +import org.springframework.cloud.sleuth.DefaultSpanNamer; +import org.springframework.cloud.sleuth.SpanNamer; +import org.springframework.core.task.AsyncTaskExecutor; + +/** + * {@link AsyncTaskExecutor} that wraps {@link Runnable} and {@link Callable} in a trace + * representation. + * + * @author Marcin Grzejszczak + * @since 2.1.0 + */ +public class LazyTraceAsyncTaskExecutor implements AsyncTaskExecutor { + + private static final Log log = LogFactory.getLog(LazyTraceAsyncTaskExecutor.class); + + private final BeanFactory beanFactory; + + private final AsyncTaskExecutor delegate; + + private Tracing tracing; + + private SpanNamer spanNamer; + + public LazyTraceAsyncTaskExecutor(BeanFactory beanFactory, + AsyncTaskExecutor delegate) { + this.beanFactory = beanFactory; + this.delegate = delegate; + } + + @Override + public void execute(Runnable task) { + Runnable taskToRun = task; + if (!ContextUtil.isContextInCreation(this.beanFactory)) { + taskToRun = new TraceRunnable(tracing(), spanNamer(), task); + } + this.delegate.execute(taskToRun); + } + + @Override + public void execute(Runnable task, long startTimeout) { + Runnable taskToRun = task; + if (!ContextUtil.isContextInCreation(this.beanFactory)) { + taskToRun = new TraceRunnable(tracing(), spanNamer(), task); + } + this.delegate.execute(taskToRun, startTimeout); + } + + @Override + public Future submit(Runnable task) { + Runnable taskToRun = task; + if (!ContextUtil.isContextInCreation(this.beanFactory)) { + taskToRun = new TraceRunnable(tracing(), spanNamer(), task); + } + return this.delegate.submit(taskToRun); + } + + @Override + public Future submit(Callable task) { + Callable taskToRun = task; + if (!ContextUtil.isContextInCreation(this.beanFactory)) { + taskToRun = new TraceCallable<>(tracing(), spanNamer(), task); + } + return this.delegate.submit(taskToRun); + } + + // due to some race conditions trace keys might not be ready yet + private SpanNamer spanNamer() { + if (this.spanNamer == null) { + try { + this.spanNamer = this.beanFactory.getBean(SpanNamer.class); + } + catch (NoSuchBeanDefinitionException e) { + log.warn( + "SpanNamer bean not found - will provide a manually created instance"); + return new DefaultSpanNamer(); + } + } + return this.spanNamer; + } + + private Tracing tracing() { + if (this.tracing == null) { + try { + this.tracing = this.beanFactory.getBean(Tracing.class); + } + catch (NoSuchBeanDefinitionException e) { + return null; + } + } + return this.tracing; + } + +} diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceAsyncIntegrationTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceAsyncIntegrationTests.java index 4226d749f..aa0689377 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceAsyncIntegrationTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceAsyncIntegrationTests.java @@ -16,6 +16,7 @@ package org.springframework.cloud.sleuth.instrument.web; +import java.util.concurrent.Executor; import java.util.concurrent.atomic.AtomicReference; import brave.Span; @@ -34,6 +35,7 @@ import org.springframework.cloud.sleuth.instrument.DefaultTestAutoConfiguration; import org.springframework.cloud.sleuth.util.ArrayListSpanReporter; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; +import org.springframework.core.task.SimpleAsyncTaskExecutor; import org.springframework.scheduling.annotation.Async; import org.springframework.scheduling.annotation.EnableAsync; import org.springframework.test.context.junit4.SpringRunner; @@ -189,6 +191,16 @@ public class TraceAsyncIntegrationTests { return new ArrayListSpanReporter(); } + @Bean + Executor fooExecutor() { + return new SimpleAsyncTaskExecutor(); + } + + @Bean + Executor barExecutor() { + return new SimpleAsyncTaskExecutor(); + } + } static class ClassPerformingAsyncLogic { @@ -201,7 +213,7 @@ public class TraceAsyncIntegrationTests { this.tracer = tracer; } - @Async + @Async("fooExecutor") public void invokeAsynchronousLogic() { this.span.set(this.tracer.currentSpan()); }