diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/template/TraceTemplate.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/template/TraceTemplate.java index 1fd6bacd6..ccd93d6e3 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/template/TraceTemplate.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/template/TraceTemplate.java @@ -26,22 +26,21 @@ import org.springframework.cloud.sleuth.trace.TraceContextHolder; */ public class TraceTemplate implements TraceOperations { - private final TraceManager trace; + private final TraceManager traceManager; - public TraceTemplate(TraceManager trace) { - this.trace = trace; + public TraceTemplate(TraceManager traceManager) { + this.traceManager = traceManager; } @Override public T trace(final TraceCallback callback) { - if (TraceContextHolder.isTracing()) { - DelegateCallback delegate = new DelegateCallback<>(this.trace); + DelegateCallback delegate = new DelegateCallback<>(this.traceManager); Trace traceScope = delegate.startSpan(); try { return callback.doInTrace(traceScope); } finally { - this.trace.close(traceScope); + this.traceManager.close(traceScope); } } else { return callback.doInTrace(null); diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/template/TraceTemplateTest.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/template/TraceTemplateTest.java new file mode 100644 index 000000000..02c3832fe --- /dev/null +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/template/TraceTemplateTest.java @@ -0,0 +1,39 @@ +package org.springframework.cloud.sleuth.template; + +import static org.assertj.core.api.BDDAssertions.then; + +import org.junit.Test; +import org.mockito.Mockito; +import org.springframework.cloud.sleuth.Trace; +import org.springframework.cloud.sleuth.TraceManager; +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.ApplicationEventPublisher; +import org.springframework.util.JdkIdGenerator; + +public class TraceTemplateTest { + + TraceManager traceManager = new DefaultTraceManager(new AlwaysSampler(), new JdkIdGenerator(), Mockito.mock(ApplicationEventPublisher.class)); + + @Test + public void should_pass_trace_to_the_callback_if_tracing_is_active() { + Trace initialTrace = traceManager.startSpan("test"); + TraceTemplate traceTemplate = new TraceTemplate(traceManager); + + Trace traceFromCallback = whenTraceCallbackReturningCurrentTraceIsExecuted(traceTemplate); + + then(traceFromCallback).isNotNull(); + then(traceFromCallback.getSpan().getTraceId()).isEqualTo(initialTrace.getSpan().getTraceId()); + } + + private Trace whenTraceCallbackReturningCurrentTraceIsExecuted(TraceTemplate traceTemplate) { + return traceTemplate.trace(new TraceCallback() { + @Override + public Trace doInTrace(Trace traceScope) { + return TraceContextHolder.getCurrentTrace(); + } + }); + } + +} \ No newline at end of file