Added LazyTraceAsyncTaskExecutor and added some more tests

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<Executor> supplier) {
 		ProxyFactoryBean factory = new ProxyFactoryBean();

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 <T> Future<T> submit(Callable<T> task) {
+		Callable<T> 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;
+	}
+
+}

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());
 		}