From 502b6b1f850772046fb1706ae55db273e818ed24 Mon Sep 17 00:00:00 2001 From: Marcin Grzejszczak Date: Mon, 4 Jan 2016 10:16:12 +0100 Subject: [PATCH] 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 --- spring-cloud-sleuth-core/pom.xml | 10 +- .../integration/StompMessageBuilder.java | 123 +++++++++++++++ ...aceSpringIntegrationAutoConfiguration.java | 19 ++- .../TraceStompMessageChannelInterceptor.java | 108 +++++++++++++ ...eContextPropagationChannelInterceptor.java | 149 ++++++++++++++++++ ...ceStompMessageChannelInterceptorTests.java | 133 ++++++++++++++++ ...extPropagationChannelInterceptorTests.java | 112 +++++++++++++ 7 files changed, 647 insertions(+), 7 deletions(-) create mode 100644 spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/StompMessageBuilder.java create mode 100644 spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceStompMessageChannelInterceptor.java create mode 100644 spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceStompMessageContextPropagationChannelInterceptor.java create mode 100644 spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceStompMessageChannelInterceptorTests.java create mode 100644 spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceStompMessageContextPropagationChannelInterceptorTests.java diff --git a/spring-cloud-sleuth-core/pom.xml b/spring-cloud-sleuth-core/pom.xml index 671f1f478..5ed191851 100644 --- a/spring-cloud-sleuth-core/pom.xml +++ b/spring-cloud-sleuth-core/pom.xml @@ -1,7 +1,6 @@ - + 4.0.0 spring-cloud-sleuth-core @@ -130,6 +129,11 @@ JUnitParams test + + org.springframework + spring-messaging + 4.2.3.RELEASE + diff --git 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 new file mode 100644 index 000000000..077624f54 --- /dev/null +++ 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 headers = new TreeMap(); + 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 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 parents = currentSpan.getParents(); + return parents == null || parents.isEmpty() ? null : parents.get(0); + } +} 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/integration/TraceSpringIntegrationAutoConfiguration.java index 805cbc569..cf70a6c1d 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/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); } } diff --git 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 new file mode 100644 index 000000000..154cf93fe --- /dev/null +++ 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 traceScopeHolder = new ThreadLocal(); + 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 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); + } +} diff --git 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 new file mode 100644 index 000000000..e767a8a07 --- /dev/null +++ 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 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 { + + 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 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 + "}"; + } + + } + +} diff --git 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 new file mode 100644 index 000000000..3bc2d8f21 --- /dev/null +++ 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("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("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("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(); + } + } +} diff --git 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 new file mode 100644 index 000000000..a85d4fc74 --- /dev/null +++ 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("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(); + } + } +}