Find DataSourceProxy even when it is wrapped or proxied

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/jdbc/DataSourceJmxConfiguration.java

@@ -26,19 +26,23 @@
 import org.apache.commons.logging.LogFactory;
 import org.apache.tomcat.jdbc.pool.DataSourceProxy;
 
+import org.springframework.aop.framework.AopProxyUtils;
+import org.springframework.aop.support.AopUtils;
 import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnSingleCandidate;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.jdbc.datasource.DelegatingDataSource;
 import org.springframework.jmx.export.MBeanExporter;
 
 /**
  * Configures DataSource related MBeans.
  *
  * @author Stephane Nicoll
+ * @author Tadaya Tsuyukubo
  */
 @Configuration
 @ConditionalOnProperty(prefix = "spring.jmx", name = "enabled", havingValue = "true", matchIfMissing = true)
@@ -89,9 +93,10 @@ static class TomcatDataSourceJmxConfiguration {
 		@Bean
 		@ConditionalOnMissingBean(name = "dataSourceMBean")
 		public Object dataSourceMBean(DataSource dataSource) {
-			if (dataSource instanceof DataSourceProxy) {
+			DataSourceProxy dataSourceProxy = extractDataSourceProxy(dataSource);
+			if (dataSourceProxy != null) {
 				try {
-					return ((DataSourceProxy) dataSource).createPool().getJmxPool();
+					return dataSourceProxy.createPool().getJmxPool();
 				}
 				catch (SQLException ex) {
 					logger.warn("Cannot expose DataSource to JMX (could not connect)");
@@ -100,6 +105,36 @@ public Object dataSourceMBean(DataSource dataSource) {
 			return null;
 		}
 
+		/**
+		 * Perform best effort to retrieve tomcat's {@link DataSourceProxy}.
+		 *
+		 * Since {@link DataSourceProxy#unwrap(Class)} always return {@code null}, it
+		 * cannot directly retrieve {@link DataSourceProxy}. This method tries best effort
+		 * to find {@link DataSourceProxy} if the given {@link DataSource} is wrapped or
+		 * proxied by spring.
+		 * @param dataSource candidate datasource
+		 * @return found DataSourceProxy or null
+		 */
+		private DataSourceProxy extractDataSourceProxy(DataSource dataSource) {
+			if (dataSource instanceof DataSourceProxy) {
+				return (DataSourceProxy) dataSource; // found
+			}
+			else if (dataSource instanceof DelegatingDataSource) {
+				// check delegating target
+				return extractDataSourceProxy(
+						((DelegatingDataSource) dataSource).getTargetDataSource());
+			}
+			else if (AopUtils.isAopProxy(dataSource)) {
+				// for proxy by spring, try target(advised) instance
+				Object target = AopProxyUtils.getSingletonTarget(dataSource);
+				if (target instanceof DataSource) {
+					return extractDataSourceProxy((DataSource) target);
+				}
+			}
+
+			return null;
+		}
+
 	}
 
 }

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/jdbc/DataSourceJmxConfigurationTests.java

@@ -31,20 +31,24 @@
 import org.apache.tomcat.jdbc.pool.jmx.ConnectionPool;
 import org.junit.Test;
 
+import org.springframework.aop.framework.AopProxyUtils;
 import org.springframework.aop.framework.ProxyFactory;
+import org.springframework.aop.support.AopUtils;
 import org.springframework.beans.factory.config.BeanPostProcessor;
 import org.springframework.boot.autoconfigure.AutoConfigurations;
 import org.springframework.boot.autoconfigure.jmx.JmxAutoConfiguration;
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.jdbc.datasource.DelegatingDataSource;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
 /**
  * Tests for {@link DataSourceJmxConfiguration}.
  *
  * @author Stephane Nicoll
+ * @author Tadaya Tsuyukubo
  */
 public class DataSourceJmxConfigurationTests {
 
@@ -162,6 +166,78 @@ public void tomcatAutoConfiguredCanExposeMBeanPool() {
 				});
 	}
 
+	@Test
+	public void tomcatProxiedCanExposeMBeanPool() {
+		this.contextRunner.withUserConfiguration(DataSourceProxyConfiguration.class)
+				.withPropertyValues(
+						"spring.datasource.type=" + DataSource.class.getName(),
+						"spring.datasource.jmx-enabled=true")
+				.run((context) -> {
+					assertThat(context).hasSingleBean(ConnectionPool.class);
+					DataSourceProxy dataSourceProxy = (DataSourceProxy) AopProxyUtils
+							.getSingletonTarget(
+									context.getBean(javax.sql.DataSource.class));
+					assertThat(dataSourceProxy.createPool().getJmxPool())
+							.isSameAs(context.getBean(ConnectionPool.class));
+				});
+	}
+
+	@Test
+	public void tomcatDelegateCanExposeMBeanPool() {
+		this.contextRunner.withUserConfiguration(DataSourceDelegateConfiguration.class)
+				.withPropertyValues(
+						"spring.datasource.type=" + DataSource.class.getName(),
+						"spring.datasource.jmx-enabled=true")
+				.run((context) -> {
+					assertThat(context).hasSingleBean(ConnectionPool.class);
+					DataSourceProxy dataSourceProxy = (DataSourceProxy) context
+							.getBean(DelegatingDataSource.class).getTargetDataSource();
+					assertThat(dataSourceProxy.createPool().getJmxPool())
+							.isSameAs(context.getBean(ConnectionPool.class));
+				});
+	}
+
+	@Test
+	public void tomcatProxyAndDelegateCanExposeMBeanPool() {
+		this.contextRunner
+				.withUserConfiguration(DataSourceMixWrapAndProxyConfiguration.class)
+				.withPropertyValues(
+						"spring.datasource.type=" + DataSource.class.getName(),
+						"spring.datasource.jmx-enabled=true")
+				.run((context) -> {
+					assertThat(context).hasSingleBean(ConnectionPool.class);
+					DataSourceProxy dataSourceProxy = extractTomcatDataSource(
+							context.getBean(javax.sql.DataSource.class));
+					assertThat(dataSourceProxy.createPool().getJmxPool())
+							.isSameAs(context.getBean(ConnectionPool.class));
+				});
+	}
+
+	private static javax.sql.DataSource wrap(javax.sql.DataSource dataSource) {
+		return (javax.sql.DataSource) new ProxyFactory(dataSource).getProxy();
+	}
+
+	private static javax.sql.DataSource delegate(javax.sql.DataSource dataSource) {
+		return new DelegatingDataSource(dataSource);
+	}
+
+	private static DataSource extractTomcatDataSource(javax.sql.DataSource dataSource) {
+		if (dataSource instanceof DataSource) {
+			return (DataSource) dataSource;
+		}
+		else if (dataSource instanceof DelegatingDataSource) {
+			return extractTomcatDataSource(
+					((DelegatingDataSource) dataSource).getTargetDataSource());
+		}
+		else if (AopUtils.isAopProxy(dataSource)) {
+			return extractTomcatDataSource(
+					(javax.sql.DataSource) AopProxyUtils.getSingletonTarget(dataSource));
+		}
+
+		throw new RuntimeException(
+				"Not proxied or delegated tomcat DataSource: " + dataSource);
+	}
+
 	@Configuration
 	static class DataSourceProxyConfiguration {
 
@@ -182,8 +258,53 @@ public Object postProcessAfterInitialization(Object bean, String beanName) {
 			return bean;
 		}
 
-		private static javax.sql.DataSource wrap(javax.sql.DataSource dataSource) {
-			return (javax.sql.DataSource) new ProxyFactory(dataSource).getProxy();
+	}
+
+	@Configuration
+	static class DataSourceDelegateConfiguration {
+
+		@Bean
+		public static DataSourceBeanPostProcessor dataSourceBeanPostProcessor() {
+			return new DataSourceBeanPostProcessor() {
+				@Override
+				public Object postProcessAfterInitialization(Object bean,
+						String beanName) {
+					if (bean instanceof javax.sql.DataSource) {
+						return new DelegatingDataSource((javax.sql.DataSource) bean);
+					}
+					return bean;
+				}
+			};
+		}
+
+	}
+
+	@Configuration
+	static class DataSourceMixWrapAndProxyConfiguration {
+
+		@Bean
+		public static DataSourceBeanPostProcessor dataSourceBeanPostProcessor() {
+			return new DataSourceBeanPostProcessor() {
+				@Override
+				public Object postProcessAfterInitialization(Object bean,
+						String beanName) {
+					if (bean instanceof javax.sql.DataSource) {
+						javax.sql.DataSource dataSource = (javax.sql.DataSource) bean;
+						// delegate/wrap multiple times
+						for (int i = 0; i < 10; i++) {
+							if (i % 2 == 0) {
+								dataSource = wrap(dataSource);
+							}
+							else {
+								dataSource = delegate(dataSource);
+							}
+						}
+
+						return dataSource;
+					}
+					return bean;
+				}
+			};
 		}
 
 	}