From e099b6d9f40e7ac5afab58f77ccdf7773b1116a8 Mon Sep 17 00:00:00 2001
From: Marcin Grzejszczak <marcin@grzejszczak.pl>
Date: Thu, 11 Apr 2019 22:36:55 +0400
Subject: [PATCH] Added support for JMS classes implementing Connection and
 XAConnection Factory interfaces

fixes gh-1284
---
 ...ingConnectionFactoryBeanPostProcessor.java | 77 ++++++++++++++++++-
 1 file changed, 76 insertions(+), 1 deletion(-)

diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TracingConnectionFactoryBeanPostProcessor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TracingConnectionFactoryBeanPostProcessor.java
index 300513665..4209809f2 100644
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TracingConnectionFactoryBeanPostProcessor.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/messaging/TracingConnectionFactoryBeanPostProcessor.java
@@ -70,9 +70,13 @@ class TracingConnectionFactoryBeanPostProcessor implements BeanPostProcessor {
 			}
 			return bean;
 		}
+		if (bean instanceof XAConnectionFactory && bean instanceof ConnectionFactory) {
+			return new LazyConnectionAndXaConnectionFactory(this.beanFactory,
+					(ConnectionFactory) bean, (XAConnectionFactory) bean);
+		}
 		// We check XA first in case the ConnectionFactory also implements
 		// XAConnectionFactory
-		if (bean instanceof XAConnectionFactory) {
+		else if (bean instanceof XAConnectionFactory) {
 			return new LazyXAConnectionFactory(this.beanFactory,
 					(XAConnectionFactory) bean);
 		}
@@ -271,6 +275,77 @@ class LazyConnectionFactory implements ConnectionFactory {
 
 }
 
+class LazyConnectionAndXaConnectionFactory
+		implements ConnectionFactory, XAConnectionFactory {
+
+	private final ConnectionFactory connectionFactoryDelegate;
+
+	private final XAConnectionFactory xaConnectionFactoryDelegate;
+
+	LazyConnectionAndXaConnectionFactory(BeanFactory beanFactory,
+			ConnectionFactory connectionFactoryDelegate,
+			XAConnectionFactory xaConnectionFactoryDelegate) {
+		this.connectionFactoryDelegate = new LazyConnectionFactory(beanFactory,
+				connectionFactoryDelegate);
+		this.xaConnectionFactoryDelegate = new LazyXAConnectionFactory(beanFactory,
+				xaConnectionFactoryDelegate);
+	}
+
+	@Override
+	public Connection createConnection() throws JMSException {
+		return this.connectionFactoryDelegate.createConnection();
+	}
+
+	@Override
+	public Connection createConnection(String userName, String password)
+			throws JMSException {
+		return this.connectionFactoryDelegate.createConnection(userName, password);
+	}
+
+	@Override
+	public JMSContext createContext() {
+		return this.connectionFactoryDelegate.createContext();
+	}
+
+	@Override
+	public JMSContext createContext(String userName, String password) {
+		return this.connectionFactoryDelegate.createContext(userName, password);
+	}
+
+	@Override
+	public JMSContext createContext(String userName, String password, int sessionMode) {
+		return this.connectionFactoryDelegate.createContext(userName, password,
+				sessionMode);
+	}
+
+	@Override
+	public JMSContext createContext(int sessionMode) {
+		return this.connectionFactoryDelegate.createContext(sessionMode);
+	}
+
+	@Override
+	public XAConnection createXAConnection() throws JMSException {
+		return this.xaConnectionFactoryDelegate.createXAConnection();
+	}
+
+	@Override
+	public XAConnection createXAConnection(String userName, String password)
+			throws JMSException {
+		return this.xaConnectionFactoryDelegate.createXAConnection(userName, password);
+	}
+
+	@Override
+	public XAJMSContext createXAContext() {
+		return this.xaConnectionFactoryDelegate.createXAContext();
+	}
+
+	@Override
+	public XAJMSContext createXAContext(String userName, String password) {
+		return this.xaConnectionFactoryDelegate.createXAContext(userName, password);
+	}
+
+}
+
 class LazyMessageListener implements MessageListener {
 
 	private final BeanFactory beanFactory;