diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorService.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorService.java index a3fb9b7a1..4bd0d578e 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorService.java +++ b/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 List> invokeAll(Collection> tasks) throws InterruptedException { - return this.delegate.invokeAll(tasks); + return this.delegate.invokeAll(wrapCallableCollection(tasks)); } @Override public List> invokeAll(Collection> 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 invokeAny(Collection> tasks) throws InterruptedException, ExecutionException { - return this.delegate.invokeAny(tasks); + return this.delegate.invokeAny(wrapCallableCollection(tasks)); } @Override public T invokeAny(Collection> 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 Collection> wrapCallableCollection(Collection> tasks) { + List> ts = new ArrayList<>(); + for (Callable task : tasks) { + if (!(task instanceof LocalComponentTraceCallable)) { + ts.add(new LocalComponentTraceCallable<>(this.tracer, this.traceKeys, + this.spanNamer, this.spanName, task)); + } + } + return ts; } } diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorServiceTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorServiceTests.java index e44077ccf..cc51a77f5 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorServiceTests.java +++ b/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>> withOneLocalComponentTraceCallable() { + return new TypeSafeMatcher>>() { + @Override + protected boolean matchesSafely(Collection> 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(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;