Use Message headers to propagate span not thread local
Otherwise we can end up closing spans that we shouldn't, or not closing ones that we should.
This commit is contained in:
@@ -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<String, String> headers, String name, Long value) {
|
||||
if (value != null) {
|
||||
addHeader(headers, name, Span.toHex(value));
|
||||
}
|
||||
}
|
||||
|
||||
private static Long getFirst(List<Long> parents) {
|
||||
return parents.isEmpty() ? null : parents.get(0);
|
||||
}
|
||||
|
||||
@@ -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<Span> 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);
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user