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:
Dave Syer
2016-01-26 12:43:27 +00:00
parent 0fc2c1a809
commit 5e2c0e9bd4
2 changed files with 18 additions and 26 deletions

View File

@@ -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);
}

View File

@@ -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);
}