Merged Support for STOMP messages over websocket

- Merge branch 'master' of https://github.com/gauravrmazra/spring-cloud-sleuth into gauravrmazra-master - Test scope removed for dependency. Needed for stomp message interceptor support - Removed unused import from integration test - Integration tests for TraceStompMessageContextChannelInterceptor - Refactored code and added integration test for TraceStompMessageChannelInterceptor - Fix for method which was not available in jdk1.7 - Fix build fail error due to new changes for stomp messages - Support for STOMP messages over websocket
2016-01-04 10:16:12 +01:00
parent c3ac8f570f
commit 502b6b1f85
7 changed files with 647 additions and 7 deletions
--- a/spring-cloud-sleuth-core/pom.xml
+++ b/spring-cloud-sleuth-core/pom.xml
@@ -1,7 +1,6 @@
 <?xml version="1.0" encoding="UTF-8"?>
-<project xmlns="http://maven.apache.org/POM/4.0.0"
-		 xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
-		 xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+	xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
 	<modelVersion>4.0.0</modelVersion>

 	<artifactId>spring-cloud-sleuth-core</artifactId>
@@ -130,6 +129,11 @@
 			<artifactId>JUnitParams</artifactId>
 			<scope>test</scope>
 		</dependency>
+		<dependency>
+			<groupId>org.springframework</groupId>
+			<artifactId>spring-messaging</artifactId>
+			<version>4.2.3.RELEASE</version>
+		</dependency>
 	</dependencies>

 </project>
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/StompMessageBuilder.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/StompMessageBuilder.java
@@ -0,0 +1,123 @@
+/*
+ * 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.integration;
+
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+
+import org.springframework.cloud.sleuth.Span;
+import org.springframework.cloud.sleuth.Trace;
+import org.springframework.cloud.sleuth.trace.TraceContextHolder;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
+import org.springframework.messaging.simp.SimpMessageType;
+
+/**
+ * Builder class to create STOMP message
+ * 
+ * @author Gaurav Rai Mazra
+ *
+ */
+public class StompMessageBuilder {
+	
+	public static StompMessageBuilder fromMessage(Message<?> message) {
+		return new StompMessageBuilder(message);
+	}
+
+	private Map<String, Object> headers = new TreeMap<String, Object>();
+	private Message<?> message;
+
+	public StompMessageBuilder(final Message<?> message) {
+		this.message = message;
+		this.headers.putAll(message.getHeaders());
+	}
+
+	public StompMessageBuilder setHeader(String key, Object value) {
+		this.headers.put(key, value);
+		return this;
+	}
+
+	public StompMessageBuilder setHeaderIfAbsent(String key, Object value) {
+		if (this.headers.get(key) == null)
+			this.headers.put(key, value);
+
+		return this;
+	}
+
+	public StompMessageBuilder setHeadersFromSpan(final Span span) {
+		if (span != null) {
+			setHeaderIfAbsent(Trace.SPAN_ID_NAME, span.getSpanId());
+			setHeaderIfAbsent(Trace.TRACE_ID_NAME, span.getTraceId());
+			setHeaderIfAbsent(Trace.SPAN_NAME_NAME, span.getName());
+			String parentId = getParentId(TraceContextHolder.getCurrentSpan());
+			if (parentId != null)
+				setHeaderIfAbsent(Trace.PARENT_ID_NAME, parentId);
+
+			String processId = span.getProcessId();
+			if (processId != null)
+				setHeaderIfAbsent(Trace.PROCESS_ID_NAME, processId);
+		}
+		return this;
+	}
+
+	public Message<?> build() {
+		SimpMessageHeaderAccessor headerAccessor = SimpMessageHeaderAccessor.create(SimpMessageType.MESSAGE);
+		String key;
+		Object value;
+		for (Map.Entry<String, Object> entry : this.headers.entrySet()) {
+			key = entry.getKey();
+			if (key != null) {
+				value = entry.getValue();
+				pushHeaders(headerAccessor, key, value);
+			}
+		}
+
+		return org.springframework.messaging.support.MessageBuilder.createMessage(this.message.getPayload(),
+				headerAccessor.getMessageHeaders());
+	}
+
+	private void pushHeaders(final SimpMessageHeaderAccessor assessor, final String key, final Object value) {
+		switch (key) {
+		case SimpMessageHeaderAccessor.DESTINATION_HEADER:
+		case SimpMessageHeaderAccessor.MESSAGE_TYPE_HEADER:
+		case SimpMessageHeaderAccessor.SESSION_ID_HEADER:
+		case SimpMessageHeaderAccessor.SESSION_ATTRIBUTES:
+		case SimpMessageHeaderAccessor.SUBSCRIPTION_ID_HEADER:
+		case SimpMessageHeaderAccessor.USER_HEADER:
+		case SimpMessageHeaderAccessor.CONNECT_MESSAGE_HEADER:
+		case SimpMessageHeaderAccessor.HEART_BEAT_HEADER:
+		case SimpMessageHeaderAccessor.ORIGINAL_DESTINATION:
+		case SimpMessageHeaderAccessor.IGNORE_ERROR:
+		case Trace.NOT_SAMPLED_NAME:
+		case Trace.PARENT_ID_NAME:
+		case Trace.PROCESS_ID_NAME:
+		case Trace.SPAN_ID_NAME:
+		case Trace.SPAN_NAME_NAME:
+		case Trace.TRACE_ID_NAME:
+			assessor.setHeader(key, value);
+			break;
+		default:
+			assessor.setNativeHeader(key, value == null ? null : value.toString());
+		}
+	}
+
+	private String getParentId(final Span currentSpan) {
+		List<String> parents = currentSpan.getParents();
+		return parents == null || parents.isEmpty() ? null : parents.get(0);
+	}
+}
--- 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/integration/TraceSpringIntegrationAutoConfiguration.java
@@ -39,14 +39,25 @@ public class TraceSpringIntegrationAutoConfiguration {
 	@Bean
 	@GlobalChannelInterceptor
 	public TraceContextPropagationChannelInterceptor traceContextPropagationChannelInterceptor(
-			TraceManager trace) {
-		return new TraceContextPropagationChannelInterceptor(trace);
+			TraceManager traceManager) {
+		return new TraceContextPropagationChannelInterceptor(traceManager);
 	}

 	@Bean
 	@GlobalChannelInterceptor
-	public TraceChannelInterceptor traceChannelInterceptor(TraceManager trace) {
-		return new TraceChannelInterceptor(trace);
+	public TraceChannelInterceptor traceChannelInterceptor(TraceManager traceManager) {
+		return new TraceChannelInterceptor(traceManager);
+	}
+
+	@Bean
+	public TraceStompMessageChannelInterceptor traceStompMessageChannelInterceptor(TraceManager traceManager) {
+		return new TraceStompMessageChannelInterceptor(traceManager);
+	}
+
+	@Bean
+	public TraceStompMessageContextPropagationChannelInterceptor traceStompMessageContextPropagationChannelInteceptor(
+			TraceManager traceManager) {
+		return new TraceStompMessageContextPropagationChannelInterceptor(traceManager);
 	}

 }
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceStompMessageChannelInterceptor.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceStompMessageChannelInterceptor.java
@@ -0,0 +1,108 @@
+/*
+ * 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.integration;
+
+import static org.springframework.util.StringUtils.hasText;
+
+import org.springframework.cloud.sleuth.MilliSpan;
+import org.springframework.cloud.sleuth.MilliSpan.MilliSpanBuilder;
+import org.springframework.cloud.sleuth.Trace;
+import org.springframework.cloud.sleuth.TraceManager;
+import org.springframework.integration.channel.AbstractMessageChannel;
+import org.springframework.integration.context.IntegrationObjectSupport;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageChannel;
+import org.springframework.messaging.support.ChannelInterceptor;
+import org.springframework.messaging.support.ChannelInterceptorAdapter;
+
+/**
+ * Interceptor for Stomp Messages sent over websocket
+ * 
+ * @author Gaurav Rai Mazra
+ * 
+ */
+public class TraceStompMessageChannelInterceptor extends ChannelInterceptorAdapter implements ChannelInterceptor {
+	private ThreadLocal<Trace> traceScopeHolder = new ThreadLocal<Trace>();
+	private final TraceManager traceManager;
+
+	public TraceStompMessageChannelInterceptor(final TraceManager traceManager) {
+		this.traceManager = traceManager;
+	}
+
+	@Override
+	public Message<?> preSend(Message<?> message, MessageChannel channel) {
+		final TraceManager traceManager = this.traceManager;
+
+		if (traceManager.isTracing() || message.getHeaders().containsKey(Trace.NOT_SAMPLED_NAME)) {
+			return StompMessageBuilder.fromMessage(message).setHeadersFromSpan(traceManager.getCurrentSpan()).build();
+		}
+		
+		String spanId = getHeader(message, Trace.SPAN_ID_NAME);
+		String traceId = getHeader(message, Trace.TRACE_ID_NAME);
+		String name = "message/" + getChannelName(channel);
+		Trace trace;
+		if (hasText(spanId) && hasText(traceId)) {
+
+			MilliSpanBuilder span = MilliSpan.builder().traceId(traceId).spanId(spanId);
+			String parentId = getHeader(message, Trace.PARENT_ID_NAME);
+			String processId = getHeader(message, Trace.PROCESS_ID_NAME);
+			String spanName = getHeader(message, Trace.SPAN_NAME_NAME);
+			if (spanName != null) {
+				span.name(spanName);
+			}
+			if (processId != null) {
+				span.processId(processId);
+			}
+			if (parentId != null) {
+				span.parent(parentId);
+			}
+			span.remote(true);
+
+			// TODO: traceManager description?
+			trace = traceManager.startSpan(name, span.build());
+		} else {
+			trace = traceManager.startSpan(name);
+		}
+		this.traceScopeHolder.set(trace);
+		return StompMessageBuilder.fromMessage(message).setHeadersFromSpan(trace.getSpan()).build();
+	}
+
+	@Override
+	public void postSend(Message<?> message, MessageChannel channel, boolean sent) {
+		final ThreadLocal<Trace> traceScopeHolder = this.traceScopeHolder;
+		Trace traceInScope = traceScopeHolder.get();
+		this.traceManager.close(traceInScope);
+		traceScopeHolder.remove();
+	}
+
+	private 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 (name == null) {
+			name = channel.toString();
+		}
+		return name;
+	}
+
+	private String getHeader(Message<?> message, String name) {
+		return (String) message.getHeaders().get(name);
+	}
+}
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceStompMessageContextPropagationChannelInterceptor.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceStompMessageContextPropagationChannelInterceptor.java
@@ -0,0 +1,149 @@
+/*
+ * 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.integration;
+
+import java.util.Map;
+
+import org.springframework.aop.support.AopUtils;
+import org.springframework.cloud.sleuth.Span;
+import org.springframework.cloud.sleuth.Trace;
+import org.springframework.cloud.sleuth.TraceManager;
+import org.springframework.integration.channel.DirectChannel;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageChannel;
+import org.springframework.messaging.MessageHandler;
+import org.springframework.messaging.MessageHeaders;
+import org.springframework.messaging.support.ChannelInterceptorAdapter;
+import org.springframework.messaging.support.ExecutorChannelInterceptor;
+import org.springframework.util.Assert;
+
+/**
+ * 
+ * @author Gaurav Rai Mazra
+ *
+ */
+public class TraceStompMessageContextPropagationChannelInterceptor extends ChannelInterceptorAdapter
+		implements ExecutorChannelInterceptor {
+
+	private final TraceManager traceManager;
+	private final static ThreadLocal<Trace> ORIGINAL_CONTEXT = new ThreadLocal<>();
+
+	public TraceStompMessageContextPropagationChannelInterceptor(TraceManager traceManager) {
+		this.traceManager = traceManager;
+	}
+
+	@Override
+	public final Message<?> preSend(Message<?> message, MessageChannel channel) {
+		if (DirectChannel.class.isAssignableFrom(AopUtils.getTargetClass(channel))) {
+			return message;
+		}
+
+		Span span = this.traceManager.getCurrentSpan();
+
+		if (span != null) {
+			return new MessageWithSpan(message, span);
+		} else {
+			return message;
+		}
+	}
+
+	@Override
+	public final Message<?> postReceive(Message<?> message, MessageChannel channel) {
+		if (message instanceof MessageWithSpan) {
+			MessageWithSpan messageWithSpan = (MessageWithSpan) message;
+			Message<?> messageToHandle = messageWithSpan.message;
+			populatePropagatedContext(messageWithSpan.span, messageToHandle, channel);
+
+			return message;
+		}
+		return message;
+	}
+
+	@Override
+	public void afterMessageHandled(Message<?> message, MessageChannel channel, MessageHandler handler, Exception ex) {
+		resetPropagatedContext();
+	}
+
+	@Override
+	public final Message<?> beforeHandle(Message<?> message, MessageChannel channel, MessageHandler handler) {
+		return postReceive(message, channel);
+	}
+
+	protected void populatePropagatedContext(Span span, Message<?> message, MessageChannel channel) {
+		if (span != null) {
+			ORIGINAL_CONTEXT.set(this.traceManager.continueSpan(span).getSavedTrace());
+		}
+	}
+
+	protected void resetPropagatedContext() {
+		Trace originalContext = ORIGINAL_CONTEXT.get();
+		this.traceManager.detach(originalContext);
+		ORIGINAL_CONTEXT.remove();
+	}
+
+	private class MessageWithSpan implements Message<Object> {
+
+		private final Message<?> message;
+		private final Span span;
+
+		public MessageWithSpan(Message<?> message, Span span) {
+			Assert.notNull(message, "message can not be null");
+			Assert.notNull(span, "span can not be null");
+			this.span = span;
+			this.message = StompMessageBuilder.fromMessage(message).setHeadersFromSpan(this.span).build();
+			addAnnotationsToSpanFromMessage(this.message, this.span);
+		}
+		
+		private void addAnnotationsToSpanFromMessage(Message<?> message, Span span) {
+			for (Map.Entry<String, Object> entry : message.getHeaders().entrySet()) {
+				if (!Trace.HEADERS.contains(entry.getKey())) {
+					String key = "/messaging/headers/" + entry.getKey().toLowerCase();
+					String value = entry.getValue() == null ? null : entry.getValue().toString();
+					span.addAnnotation(key, value);
+				}
+			}
+
+			Object payload = message.getPayload();
+			if (payload != null) {
+				span.addAnnotation("/messaging/payload/type", payload.getClass().getCanonicalName());
+
+				if (payload instanceof String) {
+					span.addAnnotation("/messaging/payload/size", String.valueOf(((String) payload).length()));
+				} else if (payload instanceof byte[]) {
+					span.addAnnotation("/messaging/payload/size", String.valueOf(((byte[]) payload).length));
+				}
+			}
+		}
+
+		@Override
+		public Object getPayload() {
+			return this.message.getPayload();
+		}
+
+		@Override
+		public MessageHeaders getHeaders() {
+			return this.message.getHeaders();
+		}
+
+		@Override
+		public String toString() {
+			return "MessageWithSpan{" + "message=" + this.message + ", span=" + this.span + "}";
+		}
+
+	}
+
+}
--- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceStompMessageChannelInterceptorTests.java
+++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceStompMessageChannelInterceptorTests.java
@@ -0,0 +1,133 @@
+package org.springframework.cloud.sleuth.instrument.integration;
+
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertEquals;
+
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Qualifier;
+import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
+import org.springframework.boot.test.IntegrationTest;
+import org.springframework.boot.test.SpringApplicationConfiguration;
+import org.springframework.cloud.sleuth.Trace;
+import org.springframework.cloud.sleuth.TraceManager;
+import org.springframework.cloud.sleuth.instrument.integration.TraceStompMessageChannelInterceptorTests.TestApplication;
+import org.springframework.cloud.sleuth.sampler.AlwaysSampler;
+import org.springframework.cloud.sleuth.trace.TraceContextHolder;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageHandler;
+import org.springframework.messaging.MessagingException;
+import org.springframework.messaging.support.ExecutorSubscribableChannel;
+import org.springframework.messaging.support.GenericMessage;
+import org.springframework.test.annotation.DirtiesContext;
+import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+
+/**
+ * 
+ * @author Gaurav Rai Mazra
+ *
+ */
+@RunWith(SpringJUnit4ClassRunner.class)
+@SpringApplicationConfiguration(classes = TestApplication.class)
+@IntegrationTest
+@DirtiesContext
+public class TraceStompMessageChannelInterceptorTests implements MessageHandler {
+	@Autowired
+	@Qualifier("executorSubscribableChannel")
+	private ExecutorSubscribableChannel channel;
+
+	@Autowired
+	private TraceManager traceManager;
+
+	private Message<?> message;
+	
+	@Override
+	public void handleMessage(Message<?> message) throws MessagingException {
+		this.message = message;
+	}
+	
+	@Before
+	public void init() {
+		this.channel.subscribe(this);
+	}
+
+	@After
+	public void close() {
+		TraceContextHolder.removeCurrentTrace();
+		this.channel.unsubscribe(this);
+	}
+
+	@Test
+	public void test_whenHeaderIsStamped_thenNoSpanCreation() {
+		Message<?> message = StompMessageBuilder.fromMessage(new GenericMessage<String>("Message2")).setHeader(Trace.NOT_SAMPLED_NAME, "").build();
+		
+		this.channel.send(message);
+		assertNotNull("message was null", this.message);
+		
+		String spanId = this.message.getHeaders().get(Trace.SPAN_ID_NAME, String.class);
+		assertNull("spanId was not null", spanId);
+		
+		Assert.assertEquals(message.getPayload(), this.message.getPayload());
+	}
+
+	@Test
+	public void test_whenMessageHeaderIsNotStamped_thenSpanCreation() {
+		Message<?> message = StompMessageBuilder.fromMessage(new GenericMessage<String>("Message2")).build();
+		this.channel.send(message);
+		assertNotNull("message was null", this.message);
+
+		String spanId = this.message.getHeaders().get(Trace.SPAN_ID_NAME, String.class);
+		assertNotNull("spanId was null", spanId);
+
+		String traceId = this.message.getHeaders().get(Trace.TRACE_ID_NAME, String.class);
+		assertNotNull("traceId was null", traceId);
+		assertNull(TraceContextHolder.getCurrentTrace());
+	}
+
+	@Test
+	public void test_whenMessageHeaderNotStamped_thenHeaderCreation() {
+		final TraceManager traceManager = this.traceManager;
+		final Trace trace = traceManager.startSpan("testSendMessage", new AlwaysSampler(), null);
+		Message<?> message = StompMessageBuilder.fromMessage(new GenericMessage<String>("Message2")).build();
+		this.channel.send(message);
+		
+		traceManager.close(trace);
+
+		assertNotNull("message was null", this.message);
+
+		String spanId = this.message.getHeaders().get(Trace.SPAN_ID_NAME, String.class);
+		assertNotNull("spanId was null", spanId);
+
+		String traceId = this.message.getHeaders().get(Trace.TRACE_ID_NAME, String.class);
+		assertNotNull("traceId was null", traceId);
+		
+		assertEquals("Trace context is not continued", trace.getSpan().getTraceId(), traceId);
+		assertEquals("Trace context is not continued", trace.getSpan().getSpanId(), spanId);
+		assertNull(TraceContextHolder.getCurrentTrace());
+	}
+
+	@Configuration
+	@EnableAutoConfiguration
+	static class TestApplication {
+		@Autowired
+		TraceStompMessageChannelInterceptor stompChannelInterceptor;
+		@Bean
+		public ExecutorSubscribableChannel executorSubscribableChannel() {
+			ExecutorSubscribableChannel channel = new ExecutorSubscribableChannel();
+			channel.addInterceptor(stompChannelInterceptor);
+			return channel;
+		}
+
+		@Bean
+		public AlwaysSampler alwaysSampler() {
+			return new AlwaysSampler();
+		}
+	}
+}
--- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceStompMessageContextPropagationChannelInterceptorTests.java
+++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceStompMessageContextPropagationChannelInterceptorTests.java
@@ -0,0 +1,112 @@
+package org.springframework.cloud.sleuth.instrument.integration;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+
+import org.junit.After;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Qualifier;
+import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
+import org.springframework.boot.test.IntegrationTest;
+import org.springframework.boot.test.SpringApplicationConfiguration;
+import org.springframework.cloud.sleuth.Trace;
+import org.springframework.cloud.sleuth.TraceManager;
+import org.springframework.cloud.sleuth.instrument.integration.TraceStompMessageContextPropagationChannelInterceptorTests.TestApplication;
+import org.springframework.cloud.sleuth.sampler.AlwaysSampler;
+import org.springframework.cloud.sleuth.trace.TraceContextHolder;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageHandler;
+import org.springframework.messaging.MessagingException;
+import org.springframework.messaging.support.ExecutorSubscribableChannel;
+import org.springframework.messaging.support.GenericMessage;
+import org.springframework.test.annotation.DirtiesContext;
+import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+
+/**
+ * 
+ * @author Gaurav Rai Mazra
+ *
+ */
+@RunWith(SpringJUnit4ClassRunner.class)
+@SpringApplicationConfiguration(classes=TestApplication.class)
+@IntegrationTest
+@DirtiesContext
+public class TraceStompMessageContextPropagationChannelInterceptorTests implements MessageHandler {
+	@Autowired
+	@Qualifier("executorSubscribableChannel")
+	private ExecutorSubscribableChannel channel;
+
+	@Autowired
+	private TraceManager traceManager;
+	
+	@Autowired
+	private AlwaysSampler sampler;
+	
+	private Message<?> message;
+
+	@Override
+	public void handleMessage(Message<?> message) throws MessagingException {
+		this.message = message;
+	}
+	
+	@org.junit.Before
+	public void init() {
+		this.channel.subscribe(this);
+	}
+	
+	@After
+	public void close() {
+		TraceContextHolder.removeCurrentTrace();
+		this.channel.unsubscribe(this);
+	}
+	
+	@Test
+	public void testSpanPropagation() {
+		final TraceManager traceManager = this.traceManager;
+		
+		Trace trace = traceManager.startSpan("testSendMessage", this.sampler, null);
+		Message<?> m = StompMessageBuilder.fromMessage(new GenericMessage<String>("Message2")).build();
+
+		this.channel.send(m);
+		
+		String expectedSpanId = trace.getSpan().getSpanId();
+		traceManager.close(trace);
+		
+		Message<?> message = this.message;
+
+		assertNotNull("message was null", message);
+
+		String spanId = message.getHeaders().get(Trace.SPAN_ID_NAME, String.class);
+		assertEquals("spanId was wrong", expectedSpanId,  spanId);
+
+		String traceId = message.getHeaders().get(Trace.TRACE_ID_NAME, String.class);
+		assertNotNull("traceId was null", traceId);
+	}
+	
+	@Configuration
+	@EnableAutoConfiguration
+	static class TestApplication {
+		@Autowired
+		TraceStompMessageChannelInterceptor stompChannelInterceptor;
+		
+		@Autowired
+		TraceStompMessageContextPropagationChannelInterceptor stompMessageContextChannelInterceptor;
+		
+		@Bean
+		public ExecutorSubscribableChannel executorSubscribableChannel() {
+			ExecutorSubscribableChannel channel = new ExecutorSubscribableChannel();
+			channel.addInterceptor(stompChannelInterceptor);
+			channel.addInterceptor(stompMessageContextChannelInterceptor);
+			return channel;
+		}
+
+		@Bean
+		public AlwaysSampler alwaysSampler() {
+			return new AlwaysSampler();
+		}
+	}
+}