diff --git a/spring-cloud-sleuth-core/pom.xml b/spring-cloud-sleuth-core/pom.xml index 6200dd847..0a3f8e749 100644 --- a/spring-cloud-sleuth-core/pom.xml +++ b/spring-cloud-sleuth-core/pom.xml @@ -33,6 +33,11 @@ spring-boot-starter-web true + + org.springframework.boot + spring-boot-starter-websocket + true + org.springframework.boot spring-boot-configuration-processor diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/AbstractTraceChannelInterceptor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/AbstractTraceChannelInterceptor.java similarity index 73% rename from spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/AbstractTraceChannelInterceptor.java rename to spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/AbstractTraceChannelInterceptor.java index e9102cb84..cb825cadb 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/AbstractTraceChannelInterceptor.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/AbstractTraceChannelInterceptor.java @@ -1,4 +1,4 @@ -package org.springframework.cloud.sleuth.instrument.integration; +package org.springframework.cloud.sleuth.instrument.messaging; import java.util.Random; @@ -10,13 +10,15 @@ import org.springframework.integration.context.IntegrationObjectSupport; import org.springframework.messaging.Message; import org.springframework.messaging.MessageChannel; import org.springframework.messaging.support.ChannelInterceptorAdapter; +import org.springframework.messaging.support.ExecutorChannelInterceptor; +import org.springframework.util.ClassUtils; /** * Abstraction over classes related to channel intercepting * * @author Marcin Grzejszczak */ -abstract class AbstractTraceChannelInterceptor extends ChannelInterceptorAdapter { +abstract class AbstractTraceChannelInterceptor extends ChannelInterceptorAdapter implements ExecutorChannelInterceptor { protected static final String MESSAGE_NAME_PREFIX = "message/"; @@ -26,7 +28,8 @@ abstract class AbstractTraceChannelInterceptor extends ChannelInterceptorAdapter private final TraceKeys traceKeys; - protected AbstractTraceChannelInterceptor(Tracer tracer, TraceKeys traceKeys, Random random) { + protected AbstractTraceChannelInterceptor(Tracer tracer, TraceKeys traceKeys, + Random random) { this.tracer = tracer; this.traceKeys = traceKeys; this.random = random; @@ -44,12 +47,14 @@ abstract class AbstractTraceChannelInterceptor extends ChannelInterceptorAdapter * Returns a span given the message and a channel. Returns null when there was no * trace id passed initially. */ - Span buildSpan(Message message) { - if (!hasHeader(message, Span.TRACE_ID_NAME) || !hasHeader(message, Span.SPAN_ID_NAME)) { + protected Span buildSpan(Message message) { + if (!hasHeader(message, Span.TRACE_ID_NAME) + || !hasHeader(message, Span.SPAN_ID_NAME)) { return null; // cannot build a span without ids } - long spanId = hasHeader(message, Span.SPAN_ID_NAME) ? - Span.fromHex(getHeader(message, Span.SPAN_ID_NAME)) : this.random.nextLong(); + long spanId = hasHeader(message, Span.SPAN_ID_NAME) + ? Span.fromHex(getHeader(message, Span.SPAN_ID_NAME)) + : this.random.nextLong(); long traceId = Span.fromHex(getHeader(message, Span.TRACE_ID_NAME)); Span.SpanBuilder span = Span.builder().traceId(traceId).spanId(spanId); if (message.getHeaders().containsKey(Span.NOT_SAMPLED_NAME)) { @@ -85,11 +90,15 @@ abstract class AbstractTraceChannelInterceptor extends ChannelInterceptorAdapter String getChannelName(MessageChannel channel) { String name = null; - if (channel instanceof IntegrationObjectSupport) { - name = ((IntegrationObjectSupport) channel).getComponentName(); - } - if (name == null && channel instanceof AbstractMessageChannel) { - name = ((AbstractMessageChannel) channel).getFullChannelName(); + if (ClassUtils.isPresent( + "org.springframework.integration.context.IntegrationObjectSupport", + null)) { + if (channel instanceof IntegrationObjectSupport) { + name = ((IntegrationObjectSupport) channel).getComponentName(); + } + if (name == null && channel instanceof AbstractMessageChannel) { + name = ((AbstractMessageChannel) channel).getFullChannelName(); + } } if (name == null) { name = channel.toString(); 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/messaging/SpanMessageHeaders.java similarity index 78% rename from spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/SpanMessageHeaders.java rename to spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/SpanMessageHeaders.java index 4bd471632..5924a4f5b 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/messaging/SpanMessageHeaders.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.cloud.sleuth.instrument.integration; +package org.springframework.cloud.sleuth.instrument.messaging; import java.util.HashMap; import java.util.List; @@ -22,8 +22,10 @@ import java.util.Map; import org.springframework.cloud.sleuth.Span; import org.springframework.cloud.sleuth.instrument.TraceKeys; -import org.springframework.integration.support.MessageBuilder; import org.springframework.messaging.Message; +import org.springframework.messaging.support.MessageBuilder; +import org.springframework.messaging.support.MessageHeaderAccessor; +import org.springframework.messaging.support.NativeMessageHeaderAccessor; import org.springframework.util.StringUtils; /** @@ -37,7 +39,7 @@ public class SpanMessageHeaders { public static final String SPAN_HEADER = "X-Current-Span"; public static Span getSpanFromHeader(Message message) { - if (message==null) { + if (message == null) { return null; } Object object = message.getHeaders().get(SPAN_HEADER); @@ -49,10 +51,14 @@ public class SpanMessageHeaders { public static Message addSpanHeaders(TraceKeys traceKeys, Message message, Span span) { + + MessageHeaderAccessor accessor = MessageHeaderAccessor + .getMutableAccessor(message); if (span == null) { if (!message.getHeaders().containsKey(Span.NOT_SAMPLED_NAME)) { - return MessageBuilder.fromMessage(message) - .setHeader(Span.NOT_SAMPLED_NAME, "true").build(); + accessor.setHeader(Span.NOT_SAMPLED_NAME, "true"); + return MessageBuilder.createMessage(message.getPayload(), + accessor.getMessageHeaders()); } return message; } @@ -73,8 +79,16 @@ public class SpanMessageHeaders { else { addHeader(headers, Span.NOT_SAMPLED_NAME, "true"); } - return MessageBuilder.fromMessage(message).copyHeaders(headers) - .setHeader(SPAN_HEADER, span).build(); + accessor.setHeader(SPAN_HEADER, span); + accessor.copyHeaders(headers); + if (accessor instanceof NativeMessageHeaderAccessor) { + NativeMessageHeaderAccessor nativeAccessor = (NativeMessageHeaderAccessor) accessor; + for (String name : headers.keySet()) { + nativeAccessor.setNativeHeader(name, headers.get(name)); + } + } + return MessageBuilder.createMessage(message.getPayload(), + accessor.getMessageHeaders()); } public static void addAnnotations(TraceKeys traceKeys, Message message, 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/messaging/TraceChannelInterceptor.java similarity index 79% rename from spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceChannelInterceptor.java rename to spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceChannelInterceptor.java index c294e6e4f..146e04d7c 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/messaging/TraceChannelInterceptor.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.cloud.sleuth.instrument.integration; +package org.springframework.cloud.sleuth.instrument.messaging; import java.util.Random; @@ -24,6 +24,7 @@ import org.springframework.cloud.sleuth.instrument.TraceKeys; import org.springframework.cloud.sleuth.sampler.NeverSampler; import org.springframework.messaging.Message; import org.springframework.messaging.MessageChannel; +import org.springframework.messaging.MessageHandler; /** * @author Dave Syer @@ -59,4 +60,17 @@ public class TraceChannelInterceptor extends AbstractTraceChannelInterceptor { return getTracer().startTrace(name); } + @Override + public Message beforeHandle(Message message, MessageChannel channel, + MessageHandler handler) { + getTracer().continueSpan(SpanMessageHeaders.getSpanFromHeader(message)); + return message; + } + + @Override + public void afterMessageHandled(Message message, MessageChannel channel, + MessageHandler handler, Exception ex) { + getTracer().detach(SpanMessageHeaders.getSpanFromHeader(message)); + } + } diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceSpringIntegrationAutoConfiguration.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceSpringIntegrationAutoConfiguration.java similarity index 96% rename from spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceSpringIntegrationAutoConfiguration.java rename to spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceSpringIntegrationAutoConfiguration.java index 01ab7285b..de100b5fc 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceSpringIntegrationAutoConfiguration.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceSpringIntegrationAutoConfiguration.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.cloud.sleuth.instrument.integration; +package org.springframework.cloud.sleuth.instrument.messaging; import java.util.Random; diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/websocket/TraceWebSocketAutoConfiguration.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/websocket/TraceWebSocketAutoConfiguration.java new file mode 100644 index 000000000..77477f5e0 --- /dev/null +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/websocket/TraceWebSocketAutoConfiguration.java @@ -0,0 +1,45 @@ +package org.springframework.cloud.sleuth.instrument.websocket; + +import java.util.Random; + +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.autoconfigure.condition.ConditionalOnBean; +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; +import org.springframework.cloud.sleuth.Tracer; +import org.springframework.cloud.sleuth.instrument.TraceKeys; +import org.springframework.cloud.sleuth.instrument.messaging.TraceChannelInterceptor; +import org.springframework.messaging.simp.config.ChannelRegistration; +import org.springframework.stereotype.Component; +import org.springframework.web.socket.config.annotation.AbstractWebSocketMessageBrokerConfigurer; +import org.springframework.web.socket.config.annotation.DelegatingWebSocketMessageBrokerConfiguration; +import org.springframework.web.socket.config.annotation.StompEndpointRegistry; + +@Component +@ConditionalOnClass(DelegatingWebSocketMessageBrokerConfiguration.class) +@ConditionalOnBean(AbstractWebSocketMessageBrokerConfigurer.class) +public class TraceWebSocketAutoConfiguration + extends AbstractWebSocketMessageBrokerConfigurer { + + @Autowired + private Tracer tracer; + + @Autowired + private TraceKeys traceKeys; + + @Override + public void registerStompEndpoints(StompEndpointRegistry registry) { + // The user must register their own endpoints + } + + @Override + public void configureClientOutboundChannel(ChannelRegistration registration) { + registration.setInterceptors( + new TraceChannelInterceptor(this.tracer, this.traceKeys, new Random())); + } + + @Override + public void configureClientInboundChannel(ChannelRegistration registration) { + registration.setInterceptors( + new TraceChannelInterceptor(this.tracer, this.traceKeys, new Random())); + } +} \ No newline at end of file diff --git a/spring-cloud-sleuth-core/src/main/resources/META-INF/spring.factories b/spring-cloud-sleuth-core/src/main/resources/META-INF/spring.factories index 75a88c8f2..10a524a60 100644 --- a/spring-cloud-sleuth-core/src/main/resources/META-INF/spring.factories +++ b/spring-cloud-sleuth-core/src/main/resources/META-INF/spring.factories @@ -2,7 +2,8 @@ org.springframework.boot.autoconfigure.EnableAutoConfiguration=\ org.springframework.cloud.sleuth.autoconfig.TraceAutoConfiguration,\ org.springframework.cloud.sleuth.log.SleuthLogAutoConfiguration,\ -org.springframework.cloud.sleuth.instrument.integration.TraceSpringIntegrationAutoConfiguration,\ +org.springframework.cloud.sleuth.instrument.messaging.TraceSpringIntegrationAutoConfiguration,\ +org.springframework.cloud.sleuth.instrument.websocket.TraceWebSocketAutoConfiguration,\ org.springframework.cloud.sleuth.instrument.async.AsyncCustomAutoConfiguration,\ org.springframework.cloud.sleuth.instrument.async.AsyncDefaultAutoConfiguration,\ org.springframework.cloud.sleuth.instrument.hystrix.SleuthHystrixAutoConfiguration,\ diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/DefaultTestAutoConfiguration.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/DefaultTestAutoConfiguration.java index 945e20e91..1475f2594 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/DefaultTestAutoConfiguration.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/DefaultTestAutoConfiguration.java @@ -8,7 +8,7 @@ import java.lang.annotation.Target; import org.springframework.boot.autoconfigure.EnableAutoConfiguration; import org.springframework.boot.autoconfigure.jmx.JmxAutoConfiguration; import org.springframework.cloud.client.loadbalancer.LoadBalancerAutoConfiguration; -import org.springframework.cloud.sleuth.instrument.integration.TraceSpringIntegrationAutoConfiguration; +import org.springframework.cloud.sleuth.instrument.messaging.TraceSpringIntegrationAutoConfiguration; import org.springframework.context.annotation.Configuration; import org.springframework.context.annotation.EnableAspectJAutoProxy; diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/SpanMessageHeadersTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/SpanMessageHeadersTests.java new file mode 100644 index 000000000..6e76a6233 --- /dev/null +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/SpanMessageHeadersTests.java @@ -0,0 +1,60 @@ +/* + * Copyright 2015 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.instrument.messaging; + +import static org.assertj.core.api.Assertions.assertThat; + +import org.junit.Test; +import org.springframework.cloud.sleuth.Span; +import org.springframework.cloud.sleuth.instrument.TraceKeys; +import org.springframework.messaging.Message; +import org.springframework.messaging.simp.SimpMessageHeaderAccessor; +import org.springframework.messaging.support.GenericMessage; +import org.springframework.messaging.support.MessageBuilder; +import org.springframework.messaging.support.MessageHeaderAccessor; +import org.springframework.messaging.support.NativeMessageHeaderAccessor; + +/** + * @author Dave Syer + * + */ +public class SpanMessageHeadersTests { + + private TraceKeys traceKeys = new TraceKeys(); + + @Test + public void spanHeadersAdded() { + Span span = Span.builder().name("foo").spanId(1L).traceId(2L).build(); + Message message = new GenericMessage<>("Hello World"); + message = SpanMessageHeaders.addSpanHeaders(this.traceKeys, message, span); + assertThat(message.getHeaders()).containsKey(Span.SPAN_ID_NAME); + } + + @Test + public void nativeSpanHeadersAdded() { + Span span = Span.builder().name("foo").spanId(1L).traceId(2L).build(); + MessageHeaderAccessor accessor = SimpMessageHeaderAccessor.create(); + Message message = MessageBuilder.createMessage("Hello World", accessor.getMessageHeaders()); + message = SpanMessageHeaders.addSpanHeaders(this.traceKeys, message, span); + assertThat(message.getHeaders()) + .containsKey(NativeMessageHeaderAccessor.NATIVE_HEADERS); + MessageHeaderAccessor natives = NativeMessageHeaderAccessor + .getMutableAccessor(message); + assertThat(natives.getMessageHeaders()).containsKey(Span.SPAN_ID_NAME); + } + +} diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceChannelInterceptorTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TraceChannelInterceptorTests.java similarity index 97% rename from spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceChannelInterceptorTests.java rename to spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TraceChannelInterceptorTests.java index a302a41d1..0985d17a3 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceChannelInterceptorTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TraceChannelInterceptorTests.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.cloud.sleuth.instrument.integration; +package org.springframework.cloud.sleuth.instrument.messaging; import static org.assertj.core.api.BDDAssertions.then; import static org.junit.Assert.assertEquals; @@ -37,7 +37,7 @@ import org.springframework.boot.test.SpringApplicationConfiguration; import org.springframework.cloud.sleuth.Span; import org.springframework.cloud.sleuth.Tracer; import org.springframework.cloud.sleuth.event.SpanReleasedEvent; -import org.springframework.cloud.sleuth.instrument.integration.TraceChannelInterceptorTests.App; +import org.springframework.cloud.sleuth.instrument.messaging.TraceChannelInterceptorTests.App; import org.springframework.cloud.sleuth.sampler.AlwaysSampler; import org.springframework.cloud.sleuth.trace.TestSpanContextHolder; import org.springframework.context.annotation.Bean; diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceContextPropagationChannelInterceptorTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TraceContextPropagationChannelInterceptorTests.java similarity index 94% rename from spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceContextPropagationChannelInterceptorTests.java rename to spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TraceContextPropagationChannelInterceptorTests.java index 877fc0c5b..4e27d0a10 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceContextPropagationChannelInterceptorTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TraceContextPropagationChannelInterceptorTests.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.cloud.sleuth.instrument.integration; +package org.springframework.cloud.sleuth.instrument.messaging; import org.junit.After; import org.junit.Test; @@ -26,7 +26,7 @@ import org.springframework.boot.test.IntegrationTest; import org.springframework.boot.test.SpringApplicationConfiguration; import org.springframework.cloud.sleuth.Span; import org.springframework.cloud.sleuth.Tracer; -import org.springframework.cloud.sleuth.instrument.integration.TraceContextPropagationChannelInterceptorTests.App; +import org.springframework.cloud.sleuth.instrument.messaging.TraceContextPropagationChannelInterceptorTests.App; import org.springframework.cloud.sleuth.sampler.AlwaysSampler; import org.springframework.cloud.sleuth.trace.TestSpanContextHolder; import org.springframework.context.annotation.Bean;