diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/SpanMessageHeaders.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/SpanMessageHeaders.java index 57b119ec5..a431c8a4d 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/SpanMessageHeaders.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/SpanMessageHeaders.java @@ -34,6 +34,19 @@ import org.springframework.util.StringUtils; */ public class SpanMessageHeaders { + public static final String SPAN_HEADER = "X-Current-Span"; + + public static Span getSpanFromHeader(Message message) { + if (message==null) { + return null; + } + Object object = message.getHeaders().get(SPAN_HEADER); + if (object instanceof Span) { + return (Span) object; + } + return null; + } + public static Message addSpanHeaders(TraceKeys traceKeys, Message message, Span span) { if (span == null) { @@ -60,7 +73,8 @@ public class SpanMessageHeaders { else { addHeader(headers, Span.NOT_SAMPLED_NAME, ""); } - return MessageBuilder.fromMessage(message).copyHeaders(headers).build(); + return MessageBuilder.fromMessage(message).copyHeaders(headers) + .setHeader(SPAN_HEADER, span).build(); } public static void addAnnotations(TraceKeys traceKeys, Message message, @@ -100,12 +114,6 @@ public class SpanMessageHeaders { } } - private static void addHeader(Map headers, String name, Long value) { - if (value != null) { - addHeader(headers, name, Span.toHex(value)); - } - } - private static Long getFirst(List parents) { return parents.isEmpty() ? null : parents.get(0); } diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceChannelInterceptor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceChannelInterceptor.java index eb36c9b43..e33b114a2 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceChannelInterceptor.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceChannelInterceptor.java @@ -22,7 +22,6 @@ import org.springframework.cloud.sleuth.Span; import org.springframework.cloud.sleuth.Tracer; import org.springframework.cloud.sleuth.instrument.TraceKeys; import org.springframework.cloud.sleuth.sampler.IsTracingSampler; -import org.springframework.cloud.sleuth.trace.SpanContextHolder; import org.springframework.messaging.Message; import org.springframework.messaging.MessageChannel; @@ -32,36 +31,21 @@ import org.springframework.messaging.MessageChannel; */ public class TraceChannelInterceptor extends AbstractTraceChannelInterceptor { - private final ThreadLocal traceHolder = new ThreadLocal<>(); - public TraceChannelInterceptor(Tracer tracer, TraceKeys traceKeys, Random random) { super(tracer, traceKeys, random); } @Override public void postSend(Message message, MessageChannel channel, boolean sent) { - Span trace = closeAllParentMessageRelatedSpans(this.traceHolder.get()); - this.traceHolder.remove(); - if (trace != null) { - SpanContextHolder.setCurrentSpan(trace); - } - } - - private Span closeAllParentMessageRelatedSpans(Span trace) { - Span traceToClose = trace; - while (traceToClose != null && traceToClose.getName() != null && - traceToClose.getName().startsWith(MESSAGE_NAME_PREFIX)) { - traceToClose = getTracer().close(traceToClose); - } - return traceToClose; + getTracer().close(SpanMessageHeaders.getSpanFromHeader(message)); } @Override public Message preSend(Message message, MessageChannel channel) { - Span parentSpan = getTracer().isTracing() ? getTracer().getCurrentSpan() : buildSpan(message); + Span parentSpan = getTracer().isTracing() ? getTracer().getCurrentSpan() + : buildSpan(message); String name = getMessageChannelName(channel); Span span = startSpan(parentSpan, name, message); - this.traceHolder.set(span); return SpanMessageHeaders.addSpanHeaders(getTraceKeys(), message, span); }