Wrapping all Executors and ThreadPoolTaskExecutors

without this change there was a gap in passing tracing info to executors.

With this change the Executors are wrapped in LazyTraceExecutor and ThreadPoolTaskExecutors are wrapped in their tracing representation too

fixes #410

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/TraceCallable.java

@@ -73,6 +73,14 @@ public class TraceCallable<V> implements Callable<V> {
 		this.tracer.close(span);
 	}
 
+	protected Span continueSpan(Span span) {
+		return this.tracer.continueSpan(span);
+	}
+
+	protected Span detachSpan(Span span) {
+		return this.tracer.detach(span);
+	}
+
 	public Tracer getTracer() {
 		return this.tracer;
 	}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/TraceRunnable.java

@@ -77,6 +77,14 @@ public class TraceRunnable implements Runnable {
 		this.tracer.close(span);
 	}
 
+	protected Span continueSpan(Span span) {
+		return this.tracer.continueSpan(span);
+	}
+
+	protected Span detachSpan(Span span) {
+		return this.tracer.detach(span);
+	}
+
 	public Tracer getTracer() {
 		return this.tracer;
 	}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/AsyncDefaultAutoConfiguration.java

@@ -24,8 +24,8 @@ import org.springframework.boot.autoconfigure.AutoConfigureAfter;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
-import org.springframework.cloud.sleuth.Tracer;
 import org.springframework.cloud.sleuth.TraceKeys;
+import org.springframework.cloud.sleuth.Tracer;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.core.task.SimpleAsyncTaskExecutor;
@@ -69,4 +69,9 @@ public class AsyncDefaultAutoConfiguration {
 		return new TraceAsyncAspect(tracer, traceKeys);
 	}
 
+	@Bean
+	public TraceExecutorBeanPostProcessor traceExecutorBeanPostProcessor(BeanFactory beanFactory) {
+		return new TraceExecutorBeanPostProcessor(beanFactory);
+	}
+
 }

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskExecutor.java

@@ -0,0 +1,118 @@
+/*
+ * Copyright 2013-2016 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.lang.invoke.MethodHandles;
+import java.util.concurrent.Callable;
+import java.util.concurrent.Future;
+
+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.cloud.sleuth.TraceKeys;
+import org.springframework.cloud.sleuth.Tracer;
+import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
+import org.springframework.util.concurrent.ListenableFuture;
+
+/**
+ * {@link ThreadPoolTaskExecutor} that delegates execution to {@link LazyTraceExecutor}
+ *
+ * @author Marcin Grzejszczak
+ * @since 1.0.10
+ */
+public class LazyTraceThreadPoolTaskExecutor extends ThreadPoolTaskExecutor {
+
+	private static final Log log = LogFactory.getLog(MethodHandles.lookup().lookupClass());
+
+	private Tracer tracer;
+	private final BeanFactory beanFactory;
+	private final ThreadPoolTaskExecutor delegate;
+	private TraceKeys traceKeys;
+	private SpanNamer spanNamer;
+
+	public LazyTraceThreadPoolTaskExecutor(BeanFactory beanFactory,
+			ThreadPoolTaskExecutor delegate) {
+		this.beanFactory = beanFactory;
+		this.delegate = delegate;
+	}
+
+	@Override
+	public void execute(Runnable task) {
+		this.delegate.execute(new SpanContinuingTraceRunnable(tracer(), traceKeys(), spanNamer(), task));
+	}
+
+	@Override
+	public void execute(Runnable task, long startTimeout) {
+		this.delegate.execute(new SpanContinuingTraceRunnable(tracer(), traceKeys(), spanNamer(), task), startTimeout);
+	}
+
+	@Override
+	public Future<?> submit(Runnable task) {
+		return this.delegate.submit(new SpanContinuingTraceRunnable(tracer(), traceKeys(), spanNamer(), task));
+	}
+
+	@Override
+	public <T> Future<T> submit(Callable<T> task) {
+		return this.delegate.submit(new SpanContinuingTraceCallable<>(tracer(), traceKeys(), spanNamer(), task));
+	}
+
+	@Override
+	public ListenableFuture<?> submitListenable(Runnable task) {
+		return this.delegate.submitListenable(new SpanContinuingTraceRunnable(tracer(), traceKeys(), spanNamer(), task));
+	}
+
+	@Override
+	public <T> ListenableFuture<T> submitListenable(Callable<T> task) {
+		return this.delegate.submitListenable(new SpanContinuingTraceCallable<>(tracer(), traceKeys(), spanNamer(), task));
+	}
+
+	private Tracer tracer() {
+		if (this.tracer == null) {
+			this.tracer = this.beanFactory.getBean(Tracer.class);
+		}
+		return this.tracer;
+	}
+
+	private TraceKeys traceKeys() {
+		if (this.traceKeys == null) {
+			try {
+				this.traceKeys = this.beanFactory.getBean(TraceKeys.class);
+			}
+			catch (NoSuchBeanDefinitionException e) {
+				log.warn("TraceKeys bean not found - will provide a manually created instance");
+				return new TraceKeys();
+			}
+		}
+		return this.traceKeys;
+	}
+
+	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;
+	}
+}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/SpanContinuingTraceCallable.java

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2013-2016 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 org.springframework.cloud.sleuth.Span;
+import org.springframework.cloud.sleuth.SpanNamer;
+import org.springframework.cloud.sleuth.TraceCallable;
+import org.springframework.cloud.sleuth.TraceKeys;
+import org.springframework.cloud.sleuth.Tracer;
+
+/**
+ * Runnable that continues a span if there is one and creates new that is a
+ * local component span if there was no tracing present.
+ *
+ * @author Marcin Grzejszczak
+ * @since 1.0.10
+ */
+class SpanContinuingTraceCallable<V> extends TraceCallable<V> {
+
+	private final LocalComponentTraceCallable<V> traceCallable;
+
+	public SpanContinuingTraceCallable(Tracer tracer, TraceKeys traceKeys,
+			SpanNamer spanNamer, Callable<V> delegate) {
+		super(tracer, spanNamer, delegate);
+		this.traceCallable = new LocalComponentTraceCallable<>(tracer, traceKeys, spanNamer, delegate);
+	}
+
+	public SpanContinuingTraceCallable(Tracer tracer, TraceKeys traceKeys,
+			SpanNamer spanNamer, Callable<V> delegate, String name) {
+		super(tracer, spanNamer, delegate, name);
+		this.traceCallable = new LocalComponentTraceCallable<>(tracer, traceKeys, spanNamer, name, delegate);
+	}
+
+	@Override
+	public V call() throws Exception {
+		Span span = startSpan();
+		try {
+			return this.getDelegate().call();
+		}
+		finally {
+			close(span);
+		}
+	}
+
+	@Override
+	protected Span startSpan() {
+		Span span = this.getParent();
+		if (span == null) {
+			return this.traceCallable.startSpan();
+		}
+		return continueSpan(span);
+	}
+
+	@Override protected void close(Span span) {
+		if (this.getParent() == null) {
+			super.close(span);
+		} else {
+			super.detachSpan(span);
+		}
+	}
+}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/SpanContinuingTraceRunnable.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright 2013-2016 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 org.springframework.cloud.sleuth.Span;
+import org.springframework.cloud.sleuth.SpanNamer;
+import org.springframework.cloud.sleuth.TraceKeys;
+import org.springframework.cloud.sleuth.TraceRunnable;
+import org.springframework.cloud.sleuth.Tracer;
+
+/**
+ * Runnable that continues a span if there is one and creates new that is a
+ * local component span if there was no tracing present.
+ *
+ * @author Marcin Grzejszczak
+ * @since 1.0.10
+ */
+class SpanContinuingTraceRunnable extends TraceRunnable {
+
+	private final LocalComponentTraceRunnable traceRunnable;
+
+	public SpanContinuingTraceRunnable(Tracer tracer, TraceKeys traceKeys,
+			SpanNamer spanNamer, Runnable delegate) {
+		super(tracer, spanNamer, delegate);
+		this.traceRunnable = new LocalComponentTraceRunnable(tracer, traceKeys, spanNamer, delegate);
+	}
+
+	public SpanContinuingTraceRunnable(Tracer tracer, TraceKeys traceKeys,
+			SpanNamer spanNamer, Runnable delegate, String name) {
+		super(tracer, spanNamer, delegate, name);
+		this.traceRunnable = new LocalComponentTraceRunnable(tracer, traceKeys, spanNamer, delegate, name);
+	}
+
+	@Override
+	public void run() {
+		Span span = startSpan();
+		try {
+			this.getDelegate().run();
+		}
+		finally {
+			close(span);
+		}
+	}
+
+	@Override
+	protected Span startSpan() {
+		Span span = this.getParent();
+		if (span == null) {
+			return this.traceRunnable.startSpan();
+		}
+		return continueSpan(span);
+	}
+
+	@Override protected void close(Span span) {
+		if (this.getParent() == null) {
+			super.close(span);
+		} else {
+			super.detachSpan(span);
+		}
+	}
+}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceExecutorBeanPostProcessor.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2013-2016 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.lang.invoke.MethodHandles;
+import java.util.concurrent.Executor;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.springframework.beans.BeansException;
+import org.springframework.beans.factory.BeanFactory;
+import org.springframework.beans.factory.config.BeanPostProcessor;
+import org.springframework.scheduling.TaskScheduler;
+import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
+
+/**
+ * Bean post processor that wraps {@link Executor} in its Trace representation
+ *
+ * @author Marcin Grzejszczak
+ * @since 1.0.10
+ */
+class TraceExecutorBeanPostProcessor implements BeanPostProcessor {
+
+	private static final Log log = LogFactory.getLog(MethodHandles.lookup().lookupClass());
+
+	private final BeanFactory beanFactory;
+
+	public TraceExecutorBeanPostProcessor(BeanFactory beanFactory) {
+		this.beanFactory = beanFactory;
+	}
+
+	@Override public Object postProcessBeforeInitialization(Object bean, String beanName)
+			throws BeansException {
+		if (bean instanceof ThreadPoolTaskExecutor && !(bean instanceof TaskScheduler) &&
+				!(bean instanceof LazyTraceThreadPoolTaskExecutor)) {
+			return new LazyTraceThreadPoolTaskExecutor(this.beanFactory, (ThreadPoolTaskExecutor) bean);
+		}
+		if (bean instanceof Executor && !(bean instanceof TaskScheduler) && !(bean instanceof LazyTraceExecutor)) {
+			return new LazyTraceExecutor(this.beanFactory, (Executor) bean);
+		}
+		return bean;
+	}
+
+	@Override public Object postProcessAfterInitialization(Object bean, String beanName)
+			throws BeansException {
+		return bean;
+	}
+}

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/issues/issue410/Issue410Tests.java

@@ -0,0 +1,162 @@
+/*
+ * Copyright 2013-2016 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.issues.issue410;
+
+import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+import org.springframework.boot.autoconfigure.data.web.SpringDataWebAutoConfiguration;
+import org.springframework.boot.test.SpringApplicationConfiguration;
+import org.springframework.boot.test.WebIntegrationTest;
+import org.springframework.cloud.sleuth.Sampler;
+import org.springframework.cloud.sleuth.Span;
+import org.springframework.cloud.sleuth.Tracer;
+import org.springframework.cloud.sleuth.sampler.AlwaysSampler;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.core.env.Environment;
+import org.springframework.scheduling.annotation.Async;
+import org.springframework.scheduling.annotation.EnableAsync;
+import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
+import org.springframework.stereotype.Component;
+import org.springframework.test.context.TestPropertySource;
+import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+import org.springframework.web.bind.annotation.RequestMapping;
+import org.springframework.web.bind.annotation.RestController;
+import org.springframework.web.client.RestTemplate;
+
+import static org.assertj.core.api.BDDAssertions.then;
+
+/**
+ * @author Marcin Grzejszczak
+ */
+@RunWith(SpringJUnit4ClassRunner.class)
+@SpringApplicationConfiguration(classes = Application.class)
+@WebIntegrationTest
+@TestPropertySource(properties = {"ribbon.eureka.enabled=false", "feign.hystrix.enabled=false", "server.port=0"})
+public class Issue410Tests {
+
+	@Autowired Environment environment;
+	@Autowired Tracer tracer;
+	@Autowired AsyncTask asyncTask;
+	@Autowired RestTemplate restTemplate;
+
+	@Test
+	public void should_pass_tracing_info_for_tasks_running_without_a_pool() {
+		Span span = this.tracer.createSpan("foo");
+
+		String response = this.restTemplate.getForObject("http://localhost:" + port() + "/without_pool", String.class);
+
+		then(response).isEqualTo(Span.idToHex(span.getTraceId()));
+		then(this.asyncTask.getSpan().get()).isNotNull();
+		then(this.asyncTask.getSpan().get().getTraceId()).isEqualTo(span.getTraceId());
+	}
+
+	@Test
+	public void should_pass_tracing_info_for_tasks_running_with_a_pool() {
+		Span span = this.tracer.createSpan("foo");
+
+		String response = this.restTemplate.getForObject("http://localhost:" + port() + "/with_pool", String.class);
+
+		then(response).isEqualTo(Span.idToHex(span.getTraceId()));
+		then(this.asyncTask.getSpan().get()).isNotNull();
+		then(this.asyncTask.getSpan().get().getTraceId()).isEqualTo(span.getTraceId());
+	}
+
+	private int port() {
+		return this.environment.getProperty("local.server.port", Integer.class);
+	}
+}
+
+
+@Configuration
+@EnableAsync
+class AppConfig {
+
+	@Bean Sampler testSampler() {
+		return new AlwaysSampler();
+	}
+
+	@Bean RestTemplate restTemplate() {
+		return new RestTemplate();
+	}
+
+	@Bean public Executor poolTaskExecutor() {
+		ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
+		executor.initialize();
+		return executor;
+	}
+
+}
+
+@Component
+class AsyncTask {
+
+	private static final Log log = LogFactory.getLog(AsyncTask.class);
+
+	private AtomicReference<Span> span = new AtomicReference<>();
+
+	@Autowired Tracer tracer;
+
+	@Async("poolTaskExecutor")
+	public void runWithPool() {
+		log.info("This task is running with a pool.");
+		this.span.set(this.tracer.getCurrentSpan());
+	}
+
+	@Async
+	public void runWithoutPool() {
+		log.info("This task is running without a pool.");
+		this.span.set(this.tracer.getCurrentSpan());
+	}
+
+	public AtomicReference<Span> getSpan() {
+		return span;
+	}
+}
+
+@SpringBootApplication(exclude = SpringDataWebAutoConfiguration.class)
+@RestController
+class Application {
+
+	private static final Log log = LogFactory.getLog(Application.class);
+
+	@Autowired AsyncTask asyncTask;
+	@Autowired Tracer tracer;
+
+	@RequestMapping("/with_pool")
+	public String withPool() {
+		log.info("Executing with pool.");
+		this.asyncTask.runWithPool();
+		return Span.idToHex(this.tracer.getCurrentSpan().getTraceId());
+
+	}
+
+	@RequestMapping("/without_pool")
+	public String withoutPool() {
+		log.info("Executing without pool.");
+		this.asyncTask.runWithoutPool();
+		return Span.idToHex(this.tracer.getCurrentSpan().getTraceId());
+	}
+
+}