Fixed missing wrapping of ExecutorService

with this change the missing methods in ExecutorService are wrapped with LocalComponentTraceCallable. We are wrapping only those callables that have not already been wrapped. Kudos to @PranayPS for finding the issue and proposing a solution.

Fixes #333

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

@@ -15,6 +15,7 @@
  */
 package org.springframework.cloud.sleuth.instrument.async;
 
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
 import java.util.concurrent.Callable;
@@ -25,8 +26,8 @@ import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
 import org.springframework.cloud.sleuth.SpanNamer;
-import org.springframework.cloud.sleuth.Tracer;
 import org.springframework.cloud.sleuth.TraceKeys;
+import org.springframework.cloud.sleuth.Tracer;
 
 /**
  * A decorator class for {@link ExecutorService} to support tracing in Executors
@@ -110,24 +111,35 @@ public class TraceableExecutorService implements ExecutorService {
 
 	@Override
 	public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks) throws InterruptedException {
-		return this.delegate.invokeAll(tasks);
+		return this.delegate.invokeAll(wrapCallableCollection(tasks));
 	}
 
 	@Override
 	public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit)
 			throws InterruptedException {
-		return this.delegate.invokeAll(tasks, timeout, unit);
+		return this.delegate.invokeAll(wrapCallableCollection(tasks), timeout, unit);
 	}
 
 	@Override
 	public <T> T invokeAny(Collection<? extends Callable<T>> tasks) throws InterruptedException, ExecutionException {
-		return this.delegate.invokeAny(tasks);
+		return this.delegate.invokeAny(wrapCallableCollection(tasks));
 	}
 
 	@Override
 	public <T> T invokeAny(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit)
 			throws InterruptedException, ExecutionException, TimeoutException {
-		return this.delegate.invokeAny(tasks, timeout, unit);
+		return this.delegate.invokeAny(wrapCallableCollection(tasks), timeout, unit);
+	}
+
+	private <T> Collection<? extends Callable<T>> wrapCallableCollection(Collection<? extends Callable<T>> tasks) {
+		List<Callable<T>> ts = new ArrayList<>();
+		for (Callable<T> task : tasks) {
+			if (!(task instanceof LocalComponentTraceCallable)) {
+				ts.add(new LocalComponentTraceCallable<>(this.tracer, this.traceKeys,
+						this.spanNamer, this.spanName, task));
+			}
+		}
+		return ts;
 	}
 
 }

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorServiceTests.java

@@ -1,19 +1,27 @@
 package org.springframework.cloud.sleuth.instrument.async;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 import java.util.Queue;
 import java.util.Random;
+import java.util.concurrent.Callable;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 
+import org.hamcrest.Description;
+import org.hamcrest.Matcher;
+import org.hamcrest.TypeSafeMatcher;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.BDDMockito;
 import org.mockito.Mock;
+import org.mockito.Mockito;
 import org.mockito.runners.MockitoJUnitRunner;
 import org.springframework.cloud.sleuth.DefaultSpanNamer;
 import org.springframework.cloud.sleuth.NoOpSpanReporter;
@@ -21,6 +29,7 @@ import org.springframework.cloud.sleuth.Span;
 import org.springframework.cloud.sleuth.SpanNamer;
 import org.springframework.cloud.sleuth.TraceKeys;
 import org.springframework.cloud.sleuth.Tracer;
+import org.springframework.cloud.sleuth.assertions.SleuthAssertions;
 import org.springframework.cloud.sleuth.log.NoOpSpanLogger;
 import org.springframework.cloud.sleuth.sampler.AlwaysSampler;
 import org.springframework.cloud.sleuth.trace.DefaultTracer;
@@ -67,6 +76,55 @@ public class TraceableExecutorServiceTests {
 		then(this.spanVerifyingRunnable.spanIds.stream().distinct().collect(toList())).hasSize(TOTAL_THREADS);
 	}
 
+	@Test
+	@SuppressWarnings("unchecked")
+	public void should_wrap_methods_in_trace_representation_only_for_non_tracing_callables() throws Exception {
+		ExecutorService executorService = Mockito.mock(ExecutorService.class);
+		TraceableExecutorService traceManagerableExecutorService = new TraceableExecutorService(
+				executorService, this.tracer, new TraceKeys(), this.spanNamer);
+
+		traceManagerableExecutorService.invokeAll(callables());
+		BDDMockito.then(executorService).should().invokeAll(BDDMockito.argThat(withOneLocalComponentTraceCallable()));
+
+		traceManagerableExecutorService.invokeAll(callables(), 1L, TimeUnit.DAYS);
+		BDDMockito.then(executorService).should().invokeAll(BDDMockito.argThat(withOneLocalComponentTraceCallable()),
+				BDDMockito.eq(1L) , BDDMockito.eq(TimeUnit.DAYS));
+
+		traceManagerableExecutorService.invokeAny(callables());
+		BDDMockito.then(executorService).should().invokeAny(BDDMockito.argThat(withOneLocalComponentTraceCallable()));
+
+		traceManagerableExecutorService.invokeAny(callables(), 1L, TimeUnit.DAYS);
+		BDDMockito.then(executorService).should().invokeAny(BDDMockito.argThat(withOneLocalComponentTraceCallable()),
+				BDDMockito.eq(1L) , BDDMockito.eq(TimeUnit.DAYS));
+	}
+
+	private Matcher<Collection<? extends Callable<Object>>> withOneLocalComponentTraceCallable() {
+		return new TypeSafeMatcher<Collection<? extends Callable<Object>>>() {
+			@Override
+			protected boolean matchesSafely(Collection<? extends Callable<Object>> item) {
+				try {
+					SleuthAssertions.then(item)
+							.flatExtracting(Object::getClass)
+							.containsExactly(LocalComponentTraceCallable.class);
+				} catch (AssertionError e) {
+					return false;
+				}
+				return true;
+			}
+
+			@Override public void describeTo(Description description) {
+				description.appendText("should contain a single local component trace callable");
+			}
+		};
+	}
+
+	private List callables() {
+		List list = new ArrayList<>();
+		list.add(new LocalComponentTraceCallable<Object>(this.tracer, new TraceKeys(), this.spanNamer, () -> "foo"));
+		list.add((Callable) () -> "bar");
+		return list;
+	}
+
 	@Test
 	public void should_propagate_trace_info_when_compleable_future_is_used() throws Exception {
 		Tracer tracer = this.tracer;