From 7da26be608513878ff2edf898aaa5290871af29c Mon Sep 17 00:00:00 2001 From: Spencer Gibb Date: Thu, 6 Aug 2015 15:05:34 -0600 Subject: [PATCH] Instrument spring integration. Via a global channel interceptor that adds headers to messages. fixes gh-14 --- spring-cloud-sleuth-core/pom.xml | 5 +- ...eContextPropagationChannelInterceptor.java | 203 ++++++++++++++++++ ...aceSpringIntegrationAutoConfiguration.java | 38 ++++ .../main/resources/META-INF/spring.factories | 1 + ...extPropagationChannelInterceptorTests.java | 90 ++++++++ 5 files changed, 334 insertions(+), 3 deletions(-) create mode 100644 spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceContextPropagationChannelInterceptor.java create mode 100644 spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceSpringIntegrationAutoConfiguration.java create mode 100644 spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceContextPropagationChannelInterceptorTests.java diff --git a/spring-cloud-sleuth-core/pom.xml b/spring-cloud-sleuth-core/pom.xml index 060e8db7e..c57d9aec3 100644 --- a/spring-cloud-sleuth-core/pom.xml +++ b/spring-cloud-sleuth-core/pom.xml @@ -44,9 +44,8 @@ true - org.springframework.integration - spring-integration-core - ${spring-integration.version} + org.springframework.boot + spring-boot-starter-integration true diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceContextPropagationChannelInterceptor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceContextPropagationChannelInterceptor.java new file mode 100644 index 000000000..b491e0a06 --- /dev/null +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceContextPropagationChannelInterceptor.java @@ -0,0 +1,203 @@ +/* + * Copyright 2013-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.cloud.sleuth.Trace.PARENT_ID_NAME; +import static org.springframework.cloud.sleuth.Trace.PROCESS_ID_NAME; +import static org.springframework.cloud.sleuth.Trace.SPAN_ID_NAME; +import static org.springframework.cloud.sleuth.Trace.SPAN_NAME_NAME; +import static org.springframework.cloud.sleuth.Trace.TRACE_ID_NAME; +import static org.springframework.cloud.sleuth.TraceContextHolder.getCurrentSpan; +import static org.springframework.cloud.sleuth.TraceContextHolder.setCurrentSpan; + +import java.util.HashMap; +import java.util.Map; + +import org.springframework.aop.support.AopUtils; +import org.springframework.cloud.sleuth.Span; +import org.springframework.cloud.sleuth.event.ClientReceivedEvent; +import org.springframework.cloud.sleuth.event.ClientSentEvent; +import org.springframework.context.ApplicationEvent; +import org.springframework.context.ApplicationEventPublisher; +import org.springframework.context.ApplicationEventPublisherAware; +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; + +/** + * The {@link ExecutorChannelInterceptor} implementation responsible for + * the {@link Span} propagation from one message flow's thread to another + * through the {@link MessageChannel}s involved in the flow. + *

+ * In addition this interceptor cleans up (restores) the {@link Span} + * in the containers Threads for channels like + * {@link org.springframework.integration.channel.ExecutorChannel} + * and {@link org.springframework.integration.channel.QueueChannel}. + * @author Spencer Gibb + * @since 1.0 + */ +public class TraceContextPropagationChannelInterceptor + extends ChannelInterceptorAdapter implements ExecutorChannelInterceptor, ApplicationEventPublisherAware { + + private final static ThreadLocal ORIGINAL_CONTEXT = new ThreadLocal<>(); + + private ApplicationEventPublisher publisher; + + @Override + public void setApplicationEventPublisher(ApplicationEventPublisher publisher) { + this.publisher = publisher; + } + + @Override + public final Message preSend(Message message, MessageChannel channel) { + if (DirectChannel.class.isAssignableFrom(AopUtils.getTargetClass(channel))) { + return message; + } + + //TODO: start span from headers? + Span span = getCurrentSpan(); + + if (span != null) { + publish(new ClientSentEvent(this, span)); + return new MessageWithSpan(message, span); + } + else { + return message; + } + } + + @Override + @SuppressWarnings("unchecked") + 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); + + publish(new ClientReceivedEvent(this, messageWithSpan.span)); + return message; + } + return message; + } + + @Override + public void afterMessageHandled(Message message, MessageChannel channel, MessageHandler handler, Exception ex) { + Span originalContext = ORIGINAL_CONTEXT.get(); + try { + if (originalContext == null) { + setCurrentSpan(null); + ORIGINAL_CONTEXT.remove(); + } + else { + setCurrentSpan(originalContext); + } + } + catch (Throwable t) {//NOSONAR + setCurrentSpan(null); + } + } + + @Override + public final Message beforeHandle(Message message, MessageChannel channel, MessageHandler handler) { + return postReceive(message, channel); + } + + private void publish(ApplicationEvent event) { + if (this.publisher !=null) { + this.publisher.publishEvent(event); + } + } + + private String getParentId(Span span) { + return span.getParents() != null && !span.getParents().isEmpty() ? span + .getParents().get(0) : null; + } + + protected void populatePropagatedContext(Span span, Message message, + MessageChannel channel) { + if (span != null) { + Span currentContext = getCurrentSpan(); + + ORIGINAL_CONTEXT.set(currentContext); + + setCurrentSpan(span); + } + } + + private class MessageWithSpan implements Message { + + private final Message message; + + private final Span span; + + private final MessageHeaders messageHeaders; + + public MessageWithSpan(Message message, Span span) { + Assert.notNull(message, "message can not be null"); + Assert.notNull(span, "span can not be null"); + this.message = message; + this.span = span; + + Map headers = new HashMap<>(); + headers.putAll(message.getHeaders()); + + setHeader(headers, SPAN_ID_NAME, this.span.getSpanId()); + setHeader(headers, TRACE_ID_NAME, this.span.getTraceId()); + setHeader(headers, SPAN_NAME_NAME, this.span.getName()); + String parentId = getParentId(getCurrentSpan()); + if (parentId != null) { + setHeader(headers, PARENT_ID_NAME, parentId); + } + String processId = this.span.getProcessId(); + if (processId != null) { + setHeader(headers, PROCESS_ID_NAME, processId); + } + this.messageHeaders = new MessageHeaders(headers); + } + + public void setHeader(Map headers, String name, String value) { + if (!headers.containsKey(name)) { + headers.put(name, value); + } + } + + @Override + public Object getPayload() { + return this.message.getPayload(); + } + + @Override + public MessageHeaders getHeaders() { + return this.messageHeaders; + } + + @Override + public String toString() { + return "MessageWithThreadState{" + + "message=" + message + + ", span=" + span + + ", messageHeaders=" + messageHeaders + + '}'; + } + + } +} \ No newline at end of file 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 new file mode 100644 index 000000000..a649e3560 --- /dev/null +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/integration/TraceSpringIntegrationAutoConfiguration.java @@ -0,0 +1,38 @@ +/* + * Copyright 2013-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 org.springframework.boot.autoconfigure.condition.ConditionalOnClass; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.integration.config.GlobalChannelInterceptor; + +/** + * @author Spencer Gibb + */ +@Configuration +@ConditionalOnClass(GlobalChannelInterceptor.class) +public class TraceSpringIntegrationAutoConfiguration { + + @Bean + @GlobalChannelInterceptor + @ConditionalOnProperty(value = "spring.cloud.sleuth.instrument.integration.globalChannelInterceptor.enabled", matchIfMissing = true) + public TraceContextPropagationChannelInterceptor traceContextPropagationChannelInterceptor() { + return new TraceContextPropagationChannelInterceptor(); + } +} 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 00d6ba6eb..32b4b9e5e 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,6 +2,7 @@ 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.scheduling.TraceSchedulingAutoConfiguration,\ org.springframework.cloud.sleuth.instrument.web.TraceWebAutoConfiguration,\ org.springframework.cloud.sleuth.instrument.web.client.TraceWebClientAutoConfiguration 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/integration/TraceContextPropagationChannelInterceptorTests.java new file mode 100644 index 000000000..cebbe3314 --- /dev/null +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/integration/TraceContextPropagationChannelInterceptorTests.java @@ -0,0 +1,90 @@ +/* + * Copyright 2013-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.junit.Assert.assertNotNull; +import static org.springframework.cloud.sleuth.Trace.SPAN_ID_NAME; +import static org.springframework.cloud.sleuth.Trace.TRACE_ID_NAME; + +import org.junit.After; +import org.junit.Test; +import org.springframework.boot.SpringApplication; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.cloud.sleuth.Trace; +import org.springframework.cloud.sleuth.TraceScope; +import org.springframework.cloud.sleuth.sampler.AlwaysSampler; +import org.springframework.context.ConfigurableApplicationContext; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.integration.annotation.MessageEndpoint; +import org.springframework.integration.channel.QueueChannel; +import org.springframework.integration.config.EnableIntegration; +import org.springframework.integration.support.MessageBuilder; +import org.springframework.messaging.Message; +import org.springframework.messaging.PollableChannel; +import org.springframework.test.context.ContextConfiguration; + +/** + * @author Spencer Gibb + */ +@ContextConfiguration(classes = TraceContextPropagationChannelInterceptorTests.App.class) +public class TraceContextPropagationChannelInterceptorTests { + private ConfigurableApplicationContext context; + + @After + public void close() { + if (context != null) { + context.close(); + } + } + + @Test + public void testSpanPropagation() { + context = SpringApplication.run(App.class); + + PollableChannel channel = context.getBean("channel", PollableChannel.class); + + Trace trace = context.getBean(Trace.class); + + TraceScope traceScope = trace.startSpan("testSendMessage", new AlwaysSampler(), null); + channel.send(MessageBuilder.withPayload("hi").build()); + traceScope.close(); + + Message message = channel.receive(0); + + assertNotNull("message was null", message); + + String spanId = message.getHeaders().get(SPAN_ID_NAME, String.class); + assertNotNull("spanId was null", spanId); + + String traceId = message.getHeaders().get(TRACE_ID_NAME, String.class); + assertNotNull("traceId was null", traceId); + } + + @Configuration + @EnableAutoConfiguration + @MessageEndpoint + @EnableIntegration + static class App { + + @Bean + public QueueChannel channel() { + return new QueueChannel(); + } + + } +}