diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceCallable.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceCallable.java index f7431e8c5..8d37a4055 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceCallable.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceCallable.java @@ -41,6 +41,7 @@ public class TraceCallable extends TraceDelegate> implements Call @Override public V call() throws Exception { + ensureThatThreadIsNotPollutedByPreviousTraces(); Trace trace = startSpan(); try { return this.getDelegate().call(); diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceDelegate.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceDelegate.java index ec42dca57..c50e0d2e6 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceDelegate.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceDelegate.java @@ -21,6 +21,7 @@ import org.springframework.cloud.sleuth.Trace; import org.springframework.cloud.sleuth.TraceManager; import lombok.Getter; +import org.springframework.cloud.sleuth.trace.TraceContextHolder; /** * @author Spencer Gibb @@ -62,4 +63,8 @@ public abstract class TraceDelegate { protected String getSpanName() { return this.name == null ? Thread.currentThread().getName() : this.name; } + + protected void ensureThatThreadIsNotPollutedByPreviousTraces() { + TraceContextHolder.removeCurrentTrace(); + } } diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceRunnable.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceRunnable.java index d9eb0c49d..a4695adfd 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceRunnable.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/TraceRunnable.java @@ -39,6 +39,7 @@ public class TraceRunnable extends TraceDelegate implements Runnable { @Override public void run() { + ensureThatThreadIsNotPollutedByPreviousTraces(); Trace trace = startSpan(); try { this.getDelegate().run(); diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/hystrix/TraceCommand.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/hystrix/TraceCommand.java index df5cf1254..7a77664da 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/hystrix/TraceCommand.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/hystrix/TraceCommand.java @@ -71,7 +71,7 @@ public abstract class TraceCommand extends HystrixCommand { @Override protected R run() throws Exception { - enforceThatHystrixThreadIsNotPolutedByPreviousTraces(); + enforceThatHystrixThreadIsNotPollutedByPreviousTraces(); Trace trace = this.traceManager.startSpan(getCommandKey().name(), parentSpan); try { return doRun(); @@ -80,8 +80,8 @@ public abstract class TraceCommand extends HystrixCommand { } } - // TODO: Do more analysis why this is nor removed properly - private void enforceThatHystrixThreadIsNotPolutedByPreviousTraces() { + // TODO: Do more analysis why this is not removed properly + private void enforceThatHystrixThreadIsNotPollutedByPreviousTraces() { TraceContextHolder.removeCurrentTrace(); } diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/executor/TraceableExecutorServiceTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/executor/TraceableExecutorServiceTests.java index 996d0bd99..afb25cc32 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/executor/TraceableExecutorServiceTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/executor/TraceableExecutorServiceTests.java @@ -1,156 +1,87 @@ package org.springframework.cloud.sleuth.instrument.executor; +import lombok.SneakyThrows; import org.junit.After; import org.junit.Before; import org.junit.Test; -import org.mockito.ArgumentCaptor; -import org.mockito.Mockito; +import org.junit.runner.RunWith; +import org.mockito.Mock; +import org.mockito.runners.MockitoJUnitRunner; import org.springframework.cloud.sleuth.Span; import org.springframework.cloud.sleuth.Trace; import org.springframework.cloud.sleuth.TraceManager; -import org.springframework.cloud.sleuth.event.SpanAcquiredEvent; -import org.springframework.cloud.sleuth.event.SpanReleasedEvent; -import org.springframework.cloud.sleuth.instrument.TraceRunnable; import org.springframework.cloud.sleuth.sampler.AlwaysSampler; import org.springframework.cloud.sleuth.trace.DefaultTraceManager; import org.springframework.cloud.sleuth.trace.TraceContextHolder; -import org.springframework.context.ApplicationEvent; import org.springframework.context.ApplicationEventPublisher; import org.springframework.util.JdkIdGenerator; import java.util.ArrayList; import java.util.List; -import java.util.concurrent.CountDownLatch; +import java.util.Queue; +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 java.util.concurrent.atomic.AtomicInteger; -import static org.hamcrest.Matchers.is; -import static org.junit.Assert.assertThat; -import static org.mockito.Matchers.isA; -import static org.mockito.Mockito.*; +import static java.util.stream.Collectors.toList; +import static org.assertj.core.api.BDDAssertions.then; +@RunWith(MockitoJUnitRunner.class) public class TraceableExecutorServiceTests { - private ApplicationEventPublisher publisher; - private ExecutorService traceManagerableExecutorService; - private TraceManager traceManager; - private ExecutorService executorService; + private static int TOTAL_THREADS = 10; - private int NUM_SPANS = 11; - private int TOTAL_THREADS = 10; + @Mock ApplicationEventPublisher publisher; + TraceManager traceManager; + ExecutorService executorService = Executors.newFixedThreadPool(3); + ExecutorService traceManagerableExecutorService; + SpanVerifyingRunnable spanVerifyingRunnable = new SpanVerifyingRunnable(); @Before - public void setUp() throws Exception { - this.publisher = Mockito.mock(ApplicationEventPublisher.class); - this.traceManager = new DefaultTraceManager(new AlwaysSampler(), new JdkIdGenerator(), this.publisher); - ExecutorService es = Executors.newFixedThreadPool(3); - this.traceManagerableExecutorService = new TraceableExecutorService(es, this.traceManager); - this.executorService = Executors.newFixedThreadPool(3); + public void setup() { + //traceManager = new DefaultTraceManager(new AlwaysSampler(), new RandomLongSpanIdGenerator(), publisher); + traceManager = new DefaultTraceManager(new AlwaysSampler(), new JdkIdGenerator(), publisher); + traceManagerableExecutorService = new TraceableExecutorService(executorService, traceManager); + TraceContextHolder.removeCurrentTrace(); } @After public void tearDown() throws Exception { - this.traceManager = null; - this.traceManagerableExecutorService.shutdown(); - this.executorService.shutdown(); + traceManager = null; + traceManagerableExecutorService.shutdown(); + executorService.shutdown(); TraceContextHolder.removeCurrentTrace(); } @Test - public void test_whenTraceContextOfWorkerThreadIsNotClosed_thenException() { - //THis test case ideally should fail but it is not failing because of the - // https://github.com/spring-cloud/spring-cloud-sleuth/issues/60 comment two - final AtomicInteger counter = new AtomicInteger(0); - final CountDownLatch latch = new CountDownLatch(this.TOTAL_THREADS); - Trace trace = this.traceManager.startSpan("PARENT"); - for (int i = 0; i < this.TOTAL_THREADS; i++) { - this.traceManagerableExecutorService.execute(new MyRunnable(counter, latch)); - } + @SneakyThrows + public void should_propagate_trace_id_and_set_new_span_when_traceable_executor_service_is_executed() { + Trace trace = traceManager.startSpan("PARENT"); + CompletableFuture.allOf(runnablesExecutedViaTraceManagerableExecutorService()).get(); + traceManager.close(trace); - try { - latch.await(); - } catch (InterruptedException e) { - e.printStackTrace(); - } - - this.traceManager.close(trace); - - verify(this.publisher, times(this.NUM_SPANS)).publishEvent(isA(SpanAcquiredEvent.class)); - verify(this.publisher, times(this.NUM_SPANS)).publishEvent(isA(SpanReleasedEvent.class)); - - ArgumentCaptor captor = ArgumentCaptor - .forClass(ApplicationEvent.class); - verify(this.publisher, atLeast(this.NUM_SPANS)).publishEvent(captor.capture()); - - List spans = new ArrayList<>(); - for (ApplicationEvent event : captor.getAllValues()) { - if (event instanceof SpanReleasedEvent) { - spans.add(((SpanReleasedEvent) event).getSpan()); - } - } - - assertThat("spans was wrong size", spans.size(), is(this.NUM_SPANS)); + then(spanVerifyingRunnable.traceIds.stream().distinct().collect(toList())).containsOnly(trace.getSpan().getTraceId()); + then(spanVerifyingRunnable.spanIds.stream().distinct().collect(toList())).hasSize(TOTAL_THREADS); } - @Test - public void test_whenTraceContextOfWorkerThreadIsClosed_thenNoException() { - final AtomicInteger counter = new AtomicInteger(0); - final CountDownLatch latch = new CountDownLatch(this.TOTAL_THREADS); - Trace trace = this.traceManager.startSpan("PARENT"); - for (int i = 0; i < this.TOTAL_THREADS; i++) { - final Runnable command = new TraceRunnable(this.traceManager, new MyRunnable(counter, latch)); - this.executorService.execute(command); + private CompletableFuture[] runnablesExecutedViaTraceManagerableExecutorService() { + List futures = new ArrayList<>(); + for (int i = 0; i < TOTAL_THREADS; i++) { + futures.add(CompletableFuture.runAsync(spanVerifyingRunnable, traceManagerableExecutorService)); } - - try { - latch.await(); - } catch (InterruptedException e) { - e.printStackTrace(); - } - - this.traceManager.close(trace); - - verify(this.publisher, times(this.NUM_SPANS)).publishEvent(isA(SpanAcquiredEvent.class)); - verify(this.publisher, times(this.NUM_SPANS)).publishEvent(isA(SpanReleasedEvent.class)); - - ArgumentCaptor captor = ArgumentCaptor - .forClass(ApplicationEvent.class); - verify(this.publisher, atLeast(this.NUM_SPANS)).publishEvent(captor.capture()); - - List spans = new ArrayList<>(); - for (ApplicationEvent event : captor.getAllValues()) { - if (event instanceof SpanReleasedEvent) { - spans.add(((SpanReleasedEvent) event).getSpan()); - } - } - - assertThat("spans was wrong size", spans.size(), is(this.NUM_SPANS)); + return futures.toArray(new CompletableFuture[futures.size()]); } - class MyRunnable implements Runnable { - private final AtomicInteger counter; - private final CountDownLatch latch; + class SpanVerifyingRunnable implements Runnable { - MyRunnable(final AtomicInteger counter, final CountDownLatch latch) { - this.counter = counter; - this.latch = latch; - } + Queue traceIds = new ConcurrentLinkedQueue<>(); + Queue spanIds = new ConcurrentLinkedQueue<>(); @Override public void run() { - try { - try { - TimeUnit.MILLISECONDS.sleep(100l); - } catch (InterruptedException e) { - e.printStackTrace(); - } - } - finally { - this.counter.incrementAndGet(); - this.latch.countDown(); - - } + Span span = TraceContextHolder.getCurrentSpan(); + traceIds.add(span.getTraceId()); + spanIds.add(span.getSpanId()); } }