Add explicit autoconfiguration support for STOMP and websockets

Adds channel interceptors to initiate a span on an incoming message
and also inject it into the thread context when it is handled
asynchronously.

spring-cloud-sleuth-core/pom.xml

@@ -33,6 +33,11 @@
 			<artifactId>spring-boot-starter-web</artifactId>
 			<optional>true</optional>
 		</dependency>
+		<dependency>
+			<groupId>org.springframework.boot</groupId>
+			<artifactId>spring-boot-starter-websocket</artifactId>
+			<optional>true</optional>
+		</dependency>
 		<dependency>
 			<groupId>org.springframework.boot</groupId>
 			<artifactId>spring-boot-configuration-processor</artifactId>

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();

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,

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

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;
 

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()));
+	}
+}

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,\

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;
 

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

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;

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;