Polish "Unwrap DataSource target rather than plain instanceof calls"

Closes gh-15227

Author: snicoll

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/metrics/jdbc/HikariDataSourceMetricsPostProcessor.java

@@ -16,13 +16,16 @@
 
 package org.springframework.boot.actuate.autoconfigure.metrics.jdbc;
 
+import javax.sql.DataSource;
+
 import com.zaxxer.hikari.HikariDataSource;
 import com.zaxxer.hikari.metrics.micrometer.MicrometerMetricsTrackerFactory;
 import io.micrometer.core.instrument.MeterRegistry;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
 import org.springframework.beans.factory.config.BeanPostProcessor;
+import org.springframework.boot.jdbc.DataSourceUnwrapper;
 import org.springframework.context.ApplicationContext;
 import org.springframework.core.Ordered;
 
@@ -48,13 +51,21 @@ class HikariDataSourceMetricsPostProcessor implements BeanPostProcessor, Ordered
 
 	@Override
 	public Object postProcessAfterInitialization(Object bean, String beanName) {
-		if (bean instanceof HikariDataSource) {
-			bindMetricsRegistryToHikariDataSource(getMeterRegistry(),
-					(HikariDataSource) bean);
+		HikariDataSource hikariDataSource = determineHikariDataSource(bean);
+		if (hikariDataSource != null) {
+			bindMetricsRegistryToHikariDataSource(getMeterRegistry(), hikariDataSource);
 		}
 		return bean;
 	}
 
+	private HikariDataSource determineHikariDataSource(Object bean) {
+		if (!(bean instanceof DataSource)) {
+			return null;
+		}
+		DataSource dataSource = (DataSource) bean;
+		return DataSourceUnwrapper.unwrap(dataSource, HikariDataSource.class);
+	}
+
 	private void bindMetricsRegistryToHikariDataSource(MeterRegistry registry,
 			HikariDataSource dataSource) {
 		if (dataSource.getMetricRegistry() == null

spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/metrics/jdbc/DataSourcePoolMetricsAutoConfigurationTests.java

@@ -27,6 +27,7 @@
 import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
 import org.junit.Test;
 
+import org.springframework.aop.framework.ProxyFactory;
 import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.config.BeanPostProcessor;
 import org.springframework.boot.actuate.autoconfigure.metrics.test.MetricsRun;
@@ -39,6 +40,7 @@
 import org.springframework.context.annotation.Configuration;
 import org.springframework.core.Ordered;
 import org.springframework.core.PriorityOrdered;
+import org.springframework.jdbc.datasource.DelegatingDataSource;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -185,6 +187,25 @@ public void someHikariDataSourcesCanBeInstrumented() {
 				});
 	}
 
+	@Test
+	public void hikariProxiedDataSourceCanBeInstrumented() {
+		this.contextRunner
+				.withUserConfiguration(ProxiedHikariDataSourcesConfiguration.class)
+				.withConfiguration(
+						AutoConfigurations.of(DataSourceAutoConfiguration.class))
+				.run((context) -> {
+					context.getBean("proxiedDataSource", DataSource.class)
+							.getConnection();
+					context.getBean("delegateDataSource", DataSource.class)
+							.getConnection();
+					MeterRegistry registry = context.getBean(MeterRegistry.class);
+					registry.get("hikaricp.connections").tags("pool", "firstDataSource")
+							.meter();
+					registry.get("hikaricp.connections").tags("pool", "secondOne")
+							.meter();
+				});
+	}
+
 	@Test
 	public void hikariDataSourceIsInstrumentedWithoutMetadataProvider() {
 		this.contextRunner.withUserConfiguration(OneHikariDataSourceConfiguration.class)
@@ -199,6 +220,14 @@ public void hikariDataSourceIsInstrumentedWithoutMetadataProvider() {
 				});
 	}
 
+	private static HikariDataSource createHikariDataSource(String poolName) {
+		String url = "jdbc:hsqldb:mem:test-" + UUID.randomUUID();
+		HikariDataSource hikariDataSource = DataSourceBuilder.create().url(url)
+				.type(HikariDataSource.class).build();
+		hikariDataSource.setPoolName(poolName);
+		return hikariDataSource;
+	}
+
 	@Configuration
 	static class BaseConfiguration {
 
@@ -242,12 +271,20 @@ public DataSource secondOne() {
 			return createHikariDataSource("secondOne");
 		}
 
-		private HikariDataSource createHikariDataSource(String poolName) {
-			String url = "jdbc:hsqldb:mem:test-" + UUID.randomUUID();
-			HikariDataSource hikariDataSource = DataSourceBuilder.create().url(url)
-					.type(HikariDataSource.class).build();
-			hikariDataSource.setPoolName(poolName);
-			return hikariDataSource;
+	}
+
+	@Configuration
+	static class ProxiedHikariDataSourcesConfiguration {
+
+		@Bean
+		public DataSource proxiedDataSource() {
+			return (DataSource) new ProxyFactory(
+					createHikariDataSource("firstDataSource")).getProxy();
+		}
+
+		@Bean
+		public DataSource delegateDataSource() {
+			return new DelegatingDataSource(createHikariDataSource("secondOne"));
 		}
 
 	}
@@ -257,11 +294,7 @@ static class OneHikariDataSourceConfiguration {
 
 		@Bean
 		public DataSource hikariDataSource() {
-			String url = "jdbc:hsqldb:mem:test-" + UUID.randomUUID();
-			HikariDataSource hikariDataSource = DataSourceBuilder.create().url(url)
-					.type(HikariDataSource.class).build();
-			hikariDataSource.setPoolName("hikariDataSource");
-			return hikariDataSource;
+			return createHikariDataSource("hikariDataSource");
 		}
 
 	}