Added JMS support

fixes gh-695

docs/src/main/asciidoc/spring-cloud-sleuth.adoc

@@ -1275,6 +1275,13 @@ To block this feature, set `spring.sleuth.messaging.kafka.enabled` to `false`.
 NOTE: We do not support context propagation via `@KafkaListener` annotation.
 Check https://github.com/spring-cloud/spring-cloud-sleuth/issues/1001[this issue for more information].
 
+==== Spring JMS
+
+We instrument the `JmsTemplate` so that tracing headers get injected
+into the message. We also support `@JmsListener` annotated methods on the consumer side.
+
+To block this feature, set `spring.sleuth.messaging.jms.enabled` to `false`.
+
 === Zuul
 
 We instrument the Zuul Ribbon integration by enriching the Ribbon requests with tracing information.

pom.xml

@@ -273,7 +273,7 @@
 		<spring-cloud-stream.version>Fishtown.BUILD-SNAPSHOT</spring-cloud-stream.version>
 		<spring-cloud-netflix.version>2.1.0.BUILD-SNAPSHOT</spring-cloud-netflix.version>
 		<spring-cloud-openfeign.version>2.1.0.BUILD-SNAPSHOT</spring-cloud-openfeign.version>
-		<brave.version>5.2.0</brave.version>
+		<brave.version>5.3.0-SNAPSHOT</brave.version>
 		<spring-security-boot-autoconfigure.version>2.0.4.RELEASE</spring-security-boot-autoconfigure.version>
 	</properties>
 

spring-cloud-sleuth-core/pom.xml

@@ -204,6 +204,15 @@
 			<groupId>io.zipkin.brave</groupId>
 			<artifactId>brave-instrumentation-spring-webmvc</artifactId>
 		</dependency>
+		<dependency>
+			<groupId>io.zipkin.brave</groupId>
+			<artifactId>brave-instrumentation-jms</artifactId>
+		</dependency>
+		<dependency>
+			<groupId>javax.jms</groupId>
+			<artifactId>javax.jms-api</artifactId>
+			<optional>true</optional>
+		</dependency>
 		<dependency>
 			<groupId>io.opentracing.brave</groupId>
 			<artifactId>brave-opentracing</artifactId>
@@ -219,7 +228,11 @@
 			<artifactId>spring-security-oauth2-autoconfigure</artifactId>
 			<optional>true</optional>
 		</dependency>
-		<!-- BRAVE -->
+		<dependency>
+			<groupId>org.springframework</groupId>
+			<artifactId>spring-jms</artifactId>
+			<optional>true</optional>
+		</dependency>
 		<dependency>
 			<groupId>org.springframework.boot</groupId>
 			<artifactId>spring-boot-autoconfigure-processor</artifactId>
@@ -281,6 +294,23 @@
 			<artifactId>spring-cloud-starter-netflix-eureka-client</artifactId>
 			<scope>test</scope>
 		</dependency>
+		<dependency>
+			<groupId>org.springframework.boot</groupId>
+			<artifactId>spring-boot-starter-activemq</artifactId>
+			<scope>test</scope>
+		</dependency>
+		<!-- to test DefaultJcaListenerContainerFactory -->
+		<dependency>
+			<groupId>javax.resource</groupId>
+			<artifactId>javax.resource-api</artifactId>
+			<version>1.7.1</version>
+			<scope>test</scope>
+		</dependency>
+		<dependency>
+			<groupId>org.apache.activemq</groupId>
+			<artifactId>activemq-ra</artifactId>
+			<scope>test</scope>
+		</dependency>
 	</dependencies>
 
 	<profiles>

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/SleuthMessagingProperties.java

@@ -82,6 +82,8 @@ public class SleuthMessagingProperties {
 
 		private Kafka kafka = new Kafka();
 
+		private Jms jms = new Jms();
+
 		public boolean isEnabled() {
 			return this.enabled;
 		}
@@ -105,6 +107,14 @@ public class SleuthMessagingProperties {
 		public void setKafka(Kafka kafka) {
 			this.kafka = kafka;
 		}
+
+		public Jms getJms() {
+			return this.jms;
+		}
+
+		public void setJms(Jms jms) {
+			this.jms = jms;
+		}
 	}
 
 	public static class Rabbit {
@@ -150,4 +160,26 @@ public class SleuthMessagingProperties {
 			this.remoteServiceName = remoteServiceName;
 		}
 	}
+
+	public static class Jms {
+		private boolean enabled;
+
+		private String remoteServiceName = "jms";
+
+		public boolean isEnabled() {
+			return this.enabled;
+		}
+
+		public void setEnabled(boolean enabled) {
+			this.enabled = enabled;
+		}
+
+		public String getRemoteServiceName() {
+			return this.remoteServiceName;
+		}
+
+		public void setRemoteServiceName(String remoteServiceName) {
+			this.remoteServiceName = remoteServiceName;
+		}
+	}
 }

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TraceMessagingAutoConfiguration.java

@@ -23,7 +23,9 @@ import java.util.Optional;
 import brave.Span;
 import brave.Tracer;
 import brave.Tracing;
+import brave.jms.JmsTracing;
 import brave.kafka.clients.KafkaTracing;
+import brave.propagation.CurrentTraceContext;
 import brave.spring.rabbit.SpringRabbitTracing;
 import org.aopalliance.intercept.MethodInterceptor;
 import org.aopalliance.intercept.MethodInvocation;
@@ -51,12 +53,12 @@ import org.springframework.boot.context.properties.EnableConfigurationProperties
 import org.springframework.cloud.sleuth.autoconfig.TraceAutoConfiguration;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.jms.annotation.JmsListenerConfigurer;
 import org.springframework.kafka.core.ProducerFactory;
 import org.springframework.kafka.listener.AbstractMessageListenerContainer;
 import org.springframework.kafka.listener.MessageListener;
 import org.springframework.kafka.listener.MessageListenerContainer;
 import org.springframework.kafka.listener.adapter.MessagingMessageListenerAdapter;
-import org.springframework.kafka.support.converter.RecordMessageConverter;
 import org.springframework.util.ReflectionUtils;
 
 /**
@@ -116,6 +118,39 @@ public class TraceMessagingAutoConfiguration {
 			return new SleuthKafkaAspect(kafkaTracing, tracer);
 		}
 	}
+
+	@Configuration
+	@ConditionalOnProperty(value = "spring.sleuth.messaging.jms.enabled", matchIfMissing = true)
+	@ConditionalOnClass(JmsListenerConfigurer.class)
+	protected static class SleuthJmsConfiguration {
+
+		@Bean
+		@ConditionalOnMissingBean
+		JmsTracing jmsTracing(Tracing tracing, SleuthMessagingProperties properties) {
+			return JmsTracing.newBuilder(tracing)
+					.remoteServiceName(properties.getMessaging().getJms().getRemoteServiceName())
+					.build();
+		}
+
+		@Bean
+		// for tests
+		@ConditionalOnMissingBean
+		TracingConnectionFactoryBeanPostProcessor tracingConnectionFactoryBeanPostProcessor(BeanFactory beanFactory) {
+			return new TracingConnectionFactoryBeanPostProcessor(beanFactory);
+		}
+
+		/** Choose the tracing endpoint registry */
+		@Bean
+		TracingJmsListenerEndpointRegistry tracingJmsListenerEndpointRegistry(JmsTracing jmsTracing, CurrentTraceContext current) {
+			return new TracingJmsListenerEndpointRegistry(jmsTracing, current);
+		}
+
+		/** Setup the tracing endpoint registry */
+		@Bean
+		JmsListenerConfigurer configureTracing(TracingJmsListenerEndpointRegistry registry) {
+			return registrar -> registrar.setEndpointRegistry(registry);
+		}
+	}
 }
 
 class SleuthRabbitBeanPostProcessor implements BeanPostProcessor {
@@ -208,14 +243,6 @@ class SleuthKafkaAspect {
 		return listener;
 	}
 
-	private RecordMessageConverter currentRecordMessageConverter(MessagingMessageListenerAdapter adapter)
-			throws IllegalAccessException {
-		if (this.recordMessageConverter != null) {
-			return (RecordMessageConverter) this.recordMessageConverter.get(adapter);
-		}
-		return null;
-	}
-
 	@SuppressWarnings("unchecked")
 	Object createProxy(Object bean) {
 		ProxyFactoryBean factory = new ProxyFactoryBean();

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TracingConnectionFactoryBeanPostProcessor.java

@@ -0,0 +1,349 @@
+/*
+ * Copyright 2013-2018 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 java.lang.reflect.Field;
+import javax.jms.Connection;
+import javax.jms.ConnectionFactory;
+import javax.jms.JMSContext;
+import javax.jms.JMSException;
+import javax.jms.Message;
+import javax.jms.MessageConsumer;
+import javax.jms.MessageListener;
+import javax.jms.Session;
+import javax.jms.XAConnection;
+import javax.jms.XAConnectionFactory;
+import javax.jms.XAJMSContext;
+
+import brave.Span;
+import brave.jms.JmsTracing;
+import brave.propagation.CurrentTraceContext;
+import org.springframework.beans.BeansException;
+import org.springframework.beans.factory.BeanFactory;
+import org.springframework.beans.factory.config.BeanPostProcessor;
+import org.springframework.jms.config.JmsListenerContainerFactory;
+import org.springframework.jms.config.JmsListenerEndpoint;
+import org.springframework.jms.config.JmsListenerEndpointRegistry;
+import org.springframework.jms.config.MethodJmsListenerEndpoint;
+import org.springframework.jms.config.SimpleJmsListenerEndpoint;
+import org.springframework.jms.connection.CachingConnectionFactory;
+import org.springframework.jms.listener.MessageListenerContainer;
+import org.springframework.jms.listener.adapter.MessagingMessageListenerAdapter;
+import org.springframework.jms.listener.endpoint.JmsMessageEndpointManager;
+import org.springframework.lang.Nullable;
+
+/**
+ * {@link BeanPostProcessor} wrapping around JMS {@link ConnectionFactory}
+ * @author Adrian Cole
+ * @since 2.1.0
+ */
+class TracingConnectionFactoryBeanPostProcessor implements BeanPostProcessor {
+
+	private final BeanFactory beanFactory;
+
+	TracingConnectionFactoryBeanPostProcessor(BeanFactory beanFactory) {
+		this.beanFactory = beanFactory;
+	}
+
+	@Override public Object postProcessAfterInitialization(Object bean, String beanName)
+			throws BeansException {
+		// Wrap the caching connection factories instead of its target, because it catches callbacks
+		// such as ExceptionListener. If we don't wrap, cached callbacks like this won't be traced.
+		if (bean instanceof CachingConnectionFactory) {
+			return new LazyConnectionFactory(this.beanFactory,
+					(CachingConnectionFactory) bean);
+		}
+		if (bean instanceof JmsMessageEndpointManager) {
+			JmsMessageEndpointManager manager = (JmsMessageEndpointManager) bean;
+			MessageListener listener = manager.getMessageListener();
+			if (listener != null) {
+				manager.setMessageListener(
+						new LazyMessageListener(this.beanFactory, listener));
+			}
+			return bean;
+		}
+		// We check XA first in case the ConnectionFactory also implements XAConnectionFactory
+		if (bean instanceof XAConnectionFactory) {
+			return new LazyXAConnectionFactory(this.beanFactory,
+					(XAConnectionFactory) bean);
+		}
+		else if (bean instanceof ConnectionFactory) {
+			return new LazyConnectionFactory(this.beanFactory, (ConnectionFactory) bean);
+		}
+		return bean;
+	}
+}
+
+class LazyXAConnectionFactory implements XAConnectionFactory {
+
+	private final BeanFactory beanFactory;
+	private final XAConnectionFactory delegate;
+	private JmsTracing jmsTracing;
+	private XAConnectionFactory wrappedDelegate;
+
+	LazyXAConnectionFactory(BeanFactory beanFactory, XAConnectionFactory delegate) {
+		this.beanFactory = beanFactory;
+		this.delegate = delegate;
+	}
+
+	@Override public XAConnection createXAConnection() throws JMSException {
+		return wrappedDelegate().createXAConnection();
+	}
+
+	@Override public XAConnection createXAConnection(String s, String s1)
+			throws JMSException {
+		return wrappedDelegate().createXAConnection(s, s1);
+	}
+
+	@Override public XAJMSContext createXAContext() {
+		return wrappedDelegate().createXAContext();
+	}
+
+	@Override public XAJMSContext createXAContext(String s, String s1) {
+		return wrappedDelegate().createXAContext(s, s1);
+	}
+
+	private JmsTracing jmsTracing() {
+		if (this.jmsTracing != null) {
+			return this.jmsTracing;
+		}
+		return this.jmsTracing = this.beanFactory.getBean(JmsTracing.class);
+	}
+
+	private XAConnectionFactory wrappedDelegate() {
+		if (this.wrappedDelegate != null) {
+			return this.wrappedDelegate;
+		}
+		return this.wrappedDelegate = jmsTracing().xaConnectionFactory(this.delegate);
+	}
+}
+
+class LazyConnectionFactory implements ConnectionFactory {
+
+	private final BeanFactory beanFactory;
+	private final ConnectionFactory delegate;
+	private JmsTracing jmsTracing;
+	private ConnectionFactory wrappedDelegate;
+
+	LazyConnectionFactory(BeanFactory beanFactory, ConnectionFactory delegate) {
+		this.beanFactory = beanFactory;
+		this.delegate = delegate;
+	}
+
+	@Override public Connection createConnection() throws JMSException {
+		return wrappedDelegate().createConnection();
+	}
+
+	@Override public Connection createConnection(String s, String s1)
+			throws JMSException {
+		return wrappedDelegate().createConnection(s, s1);
+	}
+
+	@Override public JMSContext createContext() {
+		return wrappedDelegate().createContext();
+	}
+
+	@Override public JMSContext createContext(String s, String s1) {
+		return wrappedDelegate().createContext(s, s1);
+	}
+
+	@Override public JMSContext createContext(String s, String s1, int i) {
+		return wrappedDelegate().createContext(s, s1, i);
+	}
+
+	@Override public JMSContext createContext(int i) {
+		return wrappedDelegate().createContext(i);
+	}
+
+	private JmsTracing jmsTracing() {
+		if (this.jmsTracing != null) {
+			return this.jmsTracing;
+		}
+		return this.jmsTracing = this.beanFactory.getBean(JmsTracing.class);
+	}
+
+	private ConnectionFactory wrappedDelegate() {
+		if (this.wrappedDelegate != null) {
+			return this.wrappedDelegate;
+		}
+		return this.wrappedDelegate = jmsTracing().connectionFactory(this.delegate);
+	}
+}
+
+class LazyMessageListener implements MessageListener {
+
+	private final BeanFactory beanFactory;
+	private final MessageListener delegate;
+	private JmsTracing jmsTracing;
+
+	LazyMessageListener(BeanFactory beanFactory, MessageListener delegate) {
+		this.beanFactory = beanFactory;
+		this.delegate = delegate;
+	}
+
+	@Override public void onMessage(Message message) {
+		wrappedDelegate().onMessage(message);
+	}
+
+	private JmsTracing jmsTracing() {
+		if (this.jmsTracing != null) {
+			return this.jmsTracing;
+		}
+		return this.jmsTracing = this.beanFactory.getBean(JmsTracing.class);
+	}
+
+	private MessageListener wrappedDelegate() {
+		// Adds a consumer span as we have no visibility into JCA's implementation of messaging
+		return jmsTracing().messageListener(this.delegate, true);
+	}
+}
+
+/**
+ * This ensures listeners end up continuing the trace from {@link MessageConsumer#receive()}
+ */
+class TracingJmsListenerEndpointRegistry extends JmsListenerEndpointRegistry {
+	final JmsTracing jmsTracing;
+	final CurrentTraceContext current;
+	// Not all state can be copied without using reflection
+	final Field messageHandlerMethodFactoryField;
+	final Field embeddedValueResolverField;
+
+	TracingJmsListenerEndpointRegistry(JmsTracing jmsTracing,
+			CurrentTraceContext current) {
+		this.jmsTracing = jmsTracing;
+		this.current = current;
+		this.messageHandlerMethodFactoryField = tryField("messageHandlerMethodFactory");
+		this.embeddedValueResolverField = tryField("embeddedValueResolver");
+	}
+
+	@Override public void registerListenerContainer(JmsListenerEndpoint endpoint,
+			JmsListenerContainerFactory<?> factory, boolean startImmediately) {
+		if (endpoint instanceof MethodJmsListenerEndpoint) {
+			endpoint = trace((MethodJmsListenerEndpoint) endpoint);
+		}
+		else if (endpoint instanceof SimpleJmsListenerEndpoint) {
+			endpoint = trace((SimpleJmsListenerEndpoint) endpoint);
+		}
+		super.registerListenerContainer(endpoint, factory, startImmediately);
+	}
+
+	/**
+	 * This wraps the {@link SimpleJmsListenerEndpoint#getMessageListener()} delegate in a new span.
+	 */
+	SimpleJmsListenerEndpoint trace(SimpleJmsListenerEndpoint source) {
+		MessageListener delegate = source.getMessageListener();
+		if (delegate == null)
+			return source;
+		source.setMessageListener(this.jmsTracing.messageListener(delegate, false));
+		return source;
+	}
+
+	/**
+	 * It would be better to trace by wrapping, but {@link MethodJmsListenerEndpoint#createMessageListener(MessageListenerContainer)},
+	 * is protected so we can't call it from outside code. In other words, a forwarding pattern can't
+	 * be used. Instead, we copy state from the input.
+	 * <p>
+	 * NOTE: As {@linkplain MethodJmsListenerEndpoint} is neither final, nor effectively final. For
+	 * this reason we can't ensure copying will get all state. For example, a subtype could hold state
+	 * we aren't aware of, or change behavior. We can consider checking that input is not a subtype,
+	 * and most conservatively leaving unknown subtypes untraced.
+	 */
+	MethodJmsListenerEndpoint trace(MethodJmsListenerEndpoint source) {
+		// Skip out rather than incompletely copying the source
+		if (this.messageHandlerMethodFactoryField == null
+				|| this.embeddedValueResolverField == null) {
+			return source;
+		}
+
+		// We want the stock implementation, except we want to wrap the message listener in a new span
+		MethodJmsListenerEndpoint dest = new MethodJmsListenerEndpoint() {
+			@Override protected MessagingMessageListenerAdapter createMessageListenerInstance() {
+				return new TracingMessagingMessageListenerAdapter(
+						TracingJmsListenerEndpointRegistry.this.jmsTracing,
+						TracingJmsListenerEndpointRegistry.this.current);
+			}
+		};
+
+		// set state from AbstractJmsListenerEndpoint
+		dest.setId(source.getId());
+		dest.setDestination(source.getDestination());
+		dest.setSubscription(source.getSubscription());
+		dest.setSelector(source.getSelector());
+		dest.setConcurrency(source.getConcurrency());
+
+		// set state from MethodJmsListenerEndpoint
+		dest.setBean(source.getBean());
+		dest.setMethod(source.getMethod());
+		dest.setMostSpecificMethod(source.getMostSpecificMethod());
+
+		try {
+			dest.setMessageHandlerMethodFactory(
+					get(source, this.messageHandlerMethodFactoryField));
+			dest.setEmbeddedValueResolver(get(source, this.embeddedValueResolverField));
+		}
+		catch (IllegalAccessException e) {
+			return source; // skip out rather than incompletely copying the source
+		}
+		return dest;
+	}
+
+	@Nullable static Field tryField(String name) {
+		try {
+			Field field = MethodJmsListenerEndpoint.class.getDeclaredField(name);
+			field.setAccessible(true);
+			return field;
+		}
+		catch (NoSuchFieldException e) {
+			return null;
+		}
+	}
+
+	@Nullable static <T> T get(Object object, Field field) throws IllegalAccessException {
+		return (T) field.get(object);
+	}
+}
+
+/**
+ * This wraps the message listener in a child span
+ */
+final class TracingMessagingMessageListenerAdapter
+		extends MessagingMessageListenerAdapter {
+
+	final JmsTracing jmsTracing;
+	final CurrentTraceContext current;
+
+	TracingMessagingMessageListenerAdapter(JmsTracing jmsTracing,
+			CurrentTraceContext current) {
+		this.jmsTracing = jmsTracing;
+		this.current = current;
+	}
+
+	@Override public void onMessage(Message message, Session session)
+			throws JMSException {
+		Span span = this.jmsTracing.nextSpan(message).name("on-message").start();
+		try (CurrentTraceContext.Scope ws = this.current.newScope(span.context())) {
+			super.onMessage(message, session);
+		}
+		catch (JMSException | RuntimeException | Error e) {
+			span.error(e);
+			throw e;
+		}
+		finally {
+			span.finish();
+		}
+	}
+}

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/JmsTracingConfigurationTest.java

@@ -0,0 +1,289 @@
+/*
+ * Copyright 2013-2018 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 java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.Callable;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
+import javax.jms.Connection;
+import javax.jms.ConnectionFactory;
+import javax.jms.JMSException;
+import javax.jms.MessageListener;
+import javax.jms.XAConnection;
+import javax.jms.XAConnectionFactory;
+import javax.resource.spi.ResourceAdapter;
+
+import brave.Tracing;
+import brave.internal.HexCodec;
+import brave.propagation.CurrentTraceContext;
+import brave.propagation.TraceContext;
+import org.apache.activemq.ra.ActiveMQActivationSpec;
+import org.apache.activemq.ra.ActiveMQResourceAdapter;
+import org.junit.Test;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.autoconfigure.AutoConfigurations;
+import org.springframework.boot.autoconfigure.AutoConfigureBefore;
+import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
+import org.springframework.boot.autoconfigure.jms.JmsAutoConfiguration;
+import org.springframework.boot.autoconfigure.jms.activemq.ActiveMQAutoConfiguration;
+import org.springframework.boot.jms.XAConnectionFactoryWrapper;
+import org.springframework.boot.test.context.assertj.AssertableApplicationContext;
+import org.springframework.boot.test.context.runner.ApplicationContextRunner;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.jca.support.ResourceAdapterFactoryBean;
+import org.springframework.jca.work.SimpleTaskWorkManager;
+import org.springframework.jms.annotation.EnableJms;
+import org.springframework.jms.annotation.JmsListener;
+import org.springframework.jms.annotation.JmsListenerConfigurer;
+import org.springframework.jms.config.JmsListenerEndpointRegistrar;
+import org.springframework.jms.config.SimpleJmsListenerEndpoint;
+import org.springframework.jms.core.JmsTemplate;
+import org.springframework.jms.listener.endpoint.JmsMessageEndpointManager;
+import zipkin2.Annotation;
+import zipkin2.Span;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+// inspired by org.springframework.boot.autoconfigure.jms.JmsAutoConfigurationTests
+/**
+ * @author Adrian Cole
+ */
+public class JmsTracingConfigurationTest {
+	final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
+			.withConfiguration(AutoConfigurations
+					.of(AnnotationJmsListenerConfiguration.class,
+							XAConfiguration.class,
+							SimpleJmsListenerConfiguration.class,
+							JcaJmsListenerConfiguration.class,
+							JmsTestTracingConfiguration.class));
+
+	@Test public void tracesConnectionFactory() {
+		contextRunner.run(JmsTracingConfigurationTest::checkConnection);
+	}
+
+	@Test public void tracesXAConnectionFactories() {
+		contextRunner.withUserConfiguration(XAConfiguration.class).run(ctx -> {
+			checkConnection(ctx);
+			checkXAConnection(ctx);
+		});
+	}
+
+	@AutoConfigureBefore(ActiveMQAutoConfiguration.class)
+	static class XAConfiguration {
+		@Bean XAConnectionFactoryWrapper xaConnectionFactoryWrapper() {
+			return connectionFactory -> (ConnectionFactory) connectionFactory;
+		}
+	}
+
+	@Test public void tracesListener_jmsMessageListener() {
+		contextRunner.withUserConfiguration(SimpleJmsListenerConfiguration.class)
+				.run(ctx -> {
+					ctx.getBean(JmsTemplate.class).convertAndSend("myQueue", "foo");
+
+					Callable<Span> takeSpan = ctx.getBean("takeSpan", Callable.class);
+					List<Span> trace = Arrays
+							.asList(takeSpan.call(), takeSpan.call(), takeSpan.call());
+
+					assertThat(trace).allSatisfy(s -> assertThat(s.traceId())
+							.isEqualTo(trace.get(0).traceId()));
+					assertThat(trace).extracting(Span::name)
+							.containsExactly("send", "receive", "on-message");
+				});
+	}
+
+	@Configuration
+	@EnableJms
+	static class SimpleJmsListenerConfiguration implements JmsListenerConfigurer {
+		@Autowired CurrentTraceContext current;
+
+		@Override public void configureJmsListeners(
+				JmsListenerEndpointRegistrar registrar) {
+			SimpleJmsListenerEndpoint endpoint = new SimpleJmsListenerEndpoint();
+			endpoint.setId("myCustomEndpointId");
+			endpoint.setDestination("myQueue");
+			endpoint.setMessageListener(simpleMessageListener(current));
+			registrar.registerEndpoint(endpoint);
+		}
+
+		@Bean MessageListener simpleMessageListener(CurrentTraceContext current) {
+			return message -> {
+				// Didn't restart the trace
+				assertThat(current.get()).extracting(TraceContext::parentIdAsLong)
+						.isNotEqualTo(0L);
+			};
+		}
+	}
+
+	@Test public void tracesListener_annotationMessageListener() {
+		contextRunner.withUserConfiguration(AnnotationJmsListenerConfiguration.class)
+				.run(ctx -> {
+					ctx.getBean(JmsTemplate.class).convertAndSend("myQueue", "foo");
+
+					Callable<Span> takeSpan = ctx.getBean("takeSpan", Callable.class);
+					List<Span> trace = Arrays
+							.asList(takeSpan.call(), takeSpan.call(), takeSpan.call());
+
+					assertThat(trace).allSatisfy(s -> assertThat(s.traceId())
+							.isEqualTo(trace.get(0).traceId()));
+					assertThat(trace).extracting(Span::name)
+							.containsExactly("send", "receive", "on-message");
+				});
+	}
+
+	@Configuration
+	@EnableJms
+	static class AnnotationJmsListenerConfiguration {
+		@Autowired CurrentTraceContext current;
+
+		@JmsListener(destination = "myQueue") public void onMessage() {
+			assertThat(current.get()).extracting(TraceContext::parentIdAsLong)
+					.isNotEqualTo(0L);
+		}
+	}
+
+	@Test public void tracesListener_jcaMessageListener() {
+		contextRunner.withUserConfiguration(JcaJmsListenerConfiguration.class)
+				.run(ctx -> {
+					ctx.getBean(JmsTemplate.class).convertAndSend("myQueue", "foo");
+
+					Callable<Span> takeSpan = ctx.getBean("takeSpan", Callable.class);
+					List<Span> trace = Arrays
+							.asList(takeSpan.call(), takeSpan.call(), takeSpan.call());
+
+					assertThat(trace).allSatisfy(s -> assertThat(s.traceId())
+							.isEqualTo(trace.get(0).traceId()));
+					assertThat(trace).extracting(Span::name)
+							.containsExactly("send", "receive", "on-message");
+				});
+	}
+
+	@Configuration static class JcaJmsListenerConfiguration {
+		@Autowired CurrentTraceContext current;
+
+		@Bean ResourceAdapterFactoryBean resourceAdapter() {
+			ResourceAdapterFactoryBean resourceAdapter = new ResourceAdapterFactoryBean();
+			ActiveMQResourceAdapter real = new ActiveMQResourceAdapter();
+			real.setServerUrl("vm://localhost?broker.persistent=false");
+			resourceAdapter.setResourceAdapter(real);
+			resourceAdapter.setWorkManager(new SimpleTaskWorkManager());
+			return resourceAdapter;
+		}
+
+		@Bean MessageListener simpleMessageListener(CurrentTraceContext current) {
+			return message -> {
+				// Didn't restart the trace
+				assertThat(current.get()).extracting(TraceContext::parentIdAsLong)
+						.isNotEqualTo(0L);
+			};
+		}
+
+		@Bean JmsMessageEndpointManager endpointManager(ResourceAdapter resourceAdapter,
+				MessageListener simpleMessageListener) {
+			JmsMessageEndpointManager endpointManager = new JmsMessageEndpointManager();
+			endpointManager.setResourceAdapter(resourceAdapter);
+
+			ActiveMQActivationSpec spec = new ActiveMQActivationSpec();
+			spec.setUseJndi(false);
+			spec.setDestinationType("javax.jms.Queue");
+			spec.setDestination("myQueue");
+
+			endpointManager.setActivationSpec(spec);
+			endpointManager.setMessageListener(simpleMessageListener);
+			return endpointManager;
+		}
+	}
+
+	static void checkConnection(AssertableApplicationContext ctx) throws JMSException {
+		// Not using try-with-resources as that doesn't exist in JMS 1.1
+		Connection con = ctx.getBean(ConnectionFactory.class).createConnection();
+		try {
+			con.setExceptionListener(exception -> {
+			});
+			assertThat(con.getExceptionListener().getClass().getName())
+					.startsWith("brave.jms.TracingExceptionListener");
+		}
+		finally {
+			con.close();
+		}
+	}
+
+	static void checkXAConnection(AssertableApplicationContext ctx) throws JMSException {
+		// Not using try-with-resources as that doesn't exist in JMS 1.1
+		XAConnection con = ctx.getBean(XAConnectionFactory.class).createXAConnection();
+		try {
+			con.setExceptionListener(exception -> {
+			});
+			assertThat(con.getExceptionListener().getClass().getName())
+					.startsWith("brave.jms.TracingExceptionListener");
+		}
+		finally {
+			con.close();
+		}
+	}
+}
+
+@Configuration
+@EnableAutoConfiguration
+class JmsTestTracingConfiguration {
+	static final String CONTEXT_LEAK = "context.leak";
+
+	/**
+	 * When testing servers or asynchronous clients, spans are reported on a worker thread. In order
+	 * to read them on the main thread, we use a concurrent queue. As some implementations report
+	 * after a response is sent, we use a blocking queue to prevent race conditions in tests.
+	 */
+	BlockingQueue<Span> spans = new LinkedBlockingQueue<>();
+
+	/**
+	 * Call this to block until a span was reported
+	 */
+	@Bean Callable<Span> takeSpan() {
+		return () -> {
+			Span result = spans.poll(3, TimeUnit.SECONDS);
+			assertThat(result).withFailMessage("Span was not reported").isNotNull();
+			assertThat(result.annotations()).extracting(Annotation::value)
+					.doesNotContain(CONTEXT_LEAK);
+			return result;
+		};
+	}
+
+	@Bean Tracing tracing(CurrentTraceContext currentTraceContext) {
+		return Tracing.newBuilder().spanReporter(s -> {
+			// make sure the context was cleared prior to finish.. no leaks!
+			TraceContext current = currentTraceContext.get();
+			boolean contextLeak = false;
+			if (current != null) {
+				// add annotation in addition to throwing, in case we are off the main thread
+				if (HexCodec.toLowerHex(current.spanId()).equals(s.id())) {
+					s = s.toBuilder().addAnnotation(s.timestampAsLong(), CONTEXT_LEAK)
+							.build();
+					contextLeak = true;
+				}
+			}
+			spans.add(s);
+			// throw so that we can see the path to the code that leaked the context
+			if (contextLeak) {
+				throw new AssertionError(
+						CONTEXT_LEAK + " on " + Thread.currentThread().getName());
+			}
+		}).currentTraceContext(currentTraceContext).build();
+	}
+}

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/messaging/TraceMessagingAutoConfigurationTests.java

@@ -28,6 +28,7 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mockito;
 import org.springframework.amqp.rabbit.core.RabbitTemplate;
+import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.BeanFactory;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
@@ -54,6 +55,7 @@ public class TraceMessagingAutoConfigurationTests {
 	@Autowired RabbitTemplate rabbitTemplate;
 	@Autowired ArrayListSpanReporter reporter;
 	@Autowired TestSleuthRabbitBeanPostProcessor postProcessor;
+	@Autowired TestSleuthJmsBeanPostProcessor jmsBeanPostProcessor;
 	@Autowired MySleuthKafkaAspect mySleuthKafkaAspect;
 	@Autowired ProducerFactory producerFactory;
 	@Autowired ConsumerFactory consumerFactory;
@@ -64,6 +66,12 @@ public class TraceMessagingAutoConfigurationTests {
 		then(this.postProcessor.rabbitTracingCalled).isTrue();
 	}
 
+	@Test
+	public void should_wrap_jms() {
+		then(this.jmsBeanPostProcessor).isNotNull();
+		then(this.jmsBeanPostProcessor.tracingCalled).isTrue();
+	}
+
 	@Test
 	public void should_wrap_kafka() {
 		this.producerFactory.createProducer();
@@ -86,12 +94,15 @@ public class TraceMessagingAutoConfigurationTests {
 			return new ArrayListSpanReporter();
 		}
 
-		@Bean SleuthRabbitBeanPostProcessor postProcessor(BeanFactory beanFactory) {
+		@Bean SleuthRabbitBeanPostProcessor sleuthRabbitBeanPostProcessor(BeanFactory beanFactory) {
 			return new TestSleuthRabbitBeanPostProcessor(beanFactory);
 		}
 		@Bean SleuthKafkaAspect sleuthKafkaAspect(KafkaTracing kafkaTracing, Tracer tracer) {
 			return new MySleuthKafkaAspect(kafkaTracing, tracer);
 		}
+		@Bean TestSleuthJmsBeanPostProcessor sleuthJmsBeanPostProcessor(BeanFactory beanFactory) {
+			return new TestSleuthJmsBeanPostProcessor(beanFactory);
+		}
 
 		@KafkaListener(topics = "backend", groupId = "foo")
 		public void onMessage(ConsumerRecord<?, ?> message) {
@@ -100,7 +111,7 @@ public class TraceMessagingAutoConfigurationTests {
 	}
 }
 
-class TestSleuthRabbitBeanPostProcessor  extends SleuthRabbitBeanPostProcessor {
+class TestSleuthRabbitBeanPostProcessor extends SleuthRabbitBeanPostProcessor {
 
 	boolean rabbitTracingCalled = false;
 
@@ -141,4 +152,19 @@ class MySleuthKafkaAspect extends SleuthKafkaAspect {
 		this.adapterWrapped = true;
 		return Mockito.mock(MessageListenerContainer.class);
 	}
+}
+
+class TestSleuthJmsBeanPostProcessor extends TracingConnectionFactoryBeanPostProcessor {
+
+	boolean tracingCalled = false;
+
+	TestSleuthJmsBeanPostProcessor(BeanFactory beanFactory) {
+		super(beanFactory);
+	}
+
+	@Override public Object postProcessAfterInitialization(Object bean, String beanName)
+			throws BeansException {
+		this.tracingCalled = true;
+		return super.postProcessAfterInitialization(bean, beanName);
+	}
 }

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/opentracing/BraveTracerTest.java

@@ -316,9 +316,5 @@ public class BraveTracerTest {
     @Bean ArrayListSpanReporter reporter() {
       return new ArrayListSpanReporter();
     }
-
-    @Bean CurrentTraceContext currentTraceContext() {
-      return new StrictCurrentTraceContext();
-    }
   }
 }