Apply metrics and health check to all DataSources

Before this commit, non default and autowire candidate DataSources are not applied. Closes spring-projectsGH-43481
quaff · Dec 11, 2024 · 7ac6d52 · 7ac6d52
1 parent e72546d
commit 7ac6d52
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 6 deletions.
diff --git a/...amework/boot/actuate/autoconfigure/jdbc/DataSourceHealthContributorAutoConfiguration.java b/...amework/boot/actuate/autoconfigure/jdbc/DataSourceHealthContributorAutoConfiguration.java
@@ -27,6 +27,7 @@
 import javax.sql.DataSource;
 
 import org.springframework.beans.factory.InitializingBean;
+import org.springframework.beans.factory.ListableBeanFactory;
 import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.boot.actuate.autoconfigure.health.ConditionalOnEnabledHealthIndicator;
 import org.springframework.boot.actuate.health.CompositeHealthContributor;
@@ -59,6 +60,7 @@
  * @author Arthur Kalimullin
  * @author Julio Gomez
  * @author Safeer Ansari
+ * @author Yanming Zhou
  * @since 2.0.0
  */
 @AutoConfiguration(after = DataSourceAutoConfiguration.class)
@@ -84,8 +86,9 @@ public void afterPropertiesSet() {
 
 	@Bean
 	@ConditionalOnMissingBean(name = { "dbHealthIndicator", "dbHealthContributor" })
-	public HealthContributor dbHealthContributor(Map<String, DataSource> dataSources,
+	public HealthContributor dbHealthContributor(ListableBeanFactory beanFactory,
 			DataSourceHealthIndicatorProperties dataSourceHealthIndicatorProperties) {
+		Map<String, DataSource> dataSources = beanFactory.getBeansOfType(DataSource.class);
 		if (dataSourceHealthIndicatorProperties.isIgnoreRoutingDataSources()) {
 			Map<String, DataSource> filteredDatasources = dataSources.entrySet()
 				.stream()

diff --git a/...ework/boot/actuate/autoconfigure/metrics/jdbc/DataSourcePoolMetricsAutoConfiguration.java b/...ework/boot/actuate/autoconfigure/metrics/jdbc/DataSourcePoolMetricsAutoConfiguration.java
@@ -31,6 +31,7 @@
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
+import org.springframework.beans.factory.ListableBeanFactory;
 import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.boot.actuate.autoconfigure.metrics.MetricsAutoConfiguration;
 import org.springframework.boot.actuate.autoconfigure.metrics.export.simple.SimpleMetricsExportAutoConfiguration;
@@ -52,6 +53,7 @@
  * {@link DataSource datasources}.
  *
  * @author Stephane Nicoll
+ * @author Yanming Zhou
  * @since 2.0.0
  */
 @AutoConfiguration(after = { MetricsAutoConfiguration.class, DataSourceAutoConfiguration.class,
@@ -67,8 +69,9 @@ static class DataSourcePoolMetadataMetricsConfiguration {
 		private static final String DATASOURCE_SUFFIX = "dataSource";
 
 		@Bean
-		DataSourcePoolMetadataMeterBinder dataSourcePoolMetadataMeterBinder(Map<String, DataSource> dataSources,
+		DataSourcePoolMetadataMeterBinder dataSourcePoolMetadataMeterBinder(ListableBeanFactory beanFactory,
 				ObjectProvider<DataSourcePoolMetadataProvider> metadataProviders) {
+			Map<String, DataSource> dataSources = beanFactory.getBeansOfType(DataSource.class);
 			return new DataSourcePoolMetadataMeterBinder(dataSources, metadataProviders);
 		}
 
@@ -120,17 +123,17 @@ private String getDataSourceName(String beanName) {
 	static class HikariDataSourceMetricsConfiguration {
 
 		@Bean
-		HikariDataSourceMeterBinder hikariDataSourceMeterBinder(ObjectProvider<DataSource> dataSources) {
-			return new HikariDataSourceMeterBinder(dataSources);
+		HikariDataSourceMeterBinder hikariDataSourceMeterBinder(ListableBeanFactory beanFactory) {
+			return new HikariDataSourceMeterBinder(beanFactory.getBeansOfType(DataSource.class).values());
 		}
 
 		static class HikariDataSourceMeterBinder implements MeterBinder {
 
 			private static final Log logger = LogFactory.getLog(HikariDataSourceMeterBinder.class);
 
-			private final ObjectProvider<DataSource> dataSources;
+			private final Collection<DataSource> dataSources;
 
-			HikariDataSourceMeterBinder(ObjectProvider<DataSource> dataSources) {
+			HikariDataSourceMeterBinder(Collection<DataSource> dataSources) {
 				this.dataSources = dataSources;
 			}
 

diff --git a/...rk/boot/actuate/autoconfigure/jdbc/DataSourceHealthContributorAutoConfigurationTests.java b/...rk/boot/actuate/autoconfigure/jdbc/DataSourceHealthContributorAutoConfigurationTests.java
@@ -19,6 +19,7 @@
 import java.sql.SQLException;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.UUID;
 
 import javax.sql.DataSource;
 
@@ -53,6 +54,7 @@
  * @author Phillip Webb
  * @author Julio Gomez
  * @author Safeer Ansari
+ * @author Yanming Zhou
  */
 class DataSourceHealthContributorAutoConfigurationTests {
 
@@ -211,6 +213,16 @@ void runWhenDataSourceHasNullRoutingKeyShouldProduceUnnamedComposedIndicator() {
 		});
 	}
 
+	@Test
+	void runWhenAdditionalDataSourceBeansShouldCreateCompositeIndicator() {
+		this.contextRunner.withUserConfiguration(AdditionalDataSourceConfig.class).run((context) -> {
+			assertThat(context).hasSingleBean(CompositeHealthContributor.class);
+			CompositeHealthContributor contributor = context.getBean(CompositeHealthContributor.class);
+			String[] names = contributor.stream().map(NamedContributor::getName).toArray(String[]::new);
+			assertThat(names).containsExactlyInAnyOrder("dataSource", "additionalDataSource");
+		});
+	}
+
 	@Configuration(proxyBeanMethods = false)
 	@EnableConfigurationProperties
 	static class DataSourceConfig {
@@ -245,6 +257,22 @@ AbstractRoutingDataSource routingDataSource() throws SQLException {
 
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	@EnableConfigurationProperties
+	static class AdditionalDataSourceConfig {
+
+		@Bean(defaultCandidate = false)
+		DataSource additionalDataSource() {
+			return DataSourceBuilder.create()
+				.type(org.apache.tomcat.jdbc.pool.DataSource.class)
+				.driverClassName("org.hsqldb.jdbc.JDBCDriver")
+				.url("jdbc:hsqldb:mem:test-" + UUID.randomUUID())
+				.username("sa")
+				.build();
+		}
+
+	}
+
 	static class ProxyDataSourceBeanPostProcessor implements BeanPostProcessor {
 
 		@Override

diff --git a/.../boot/actuate/autoconfigure/metrics/jdbc/DataSourcePoolMetricsAutoConfigurationTests.java b/.../boot/actuate/autoconfigure/metrics/jdbc/DataSourcePoolMetricsAutoConfigurationTests.java
@@ -51,6 +51,7 @@
  * @author Stephane Nicoll
  * @author Andy Wilkinson
  * @author Tommy Ludwig
+ * @author Yanming Zhou
  */
 class DataSourcePoolMetricsAutoConfigurationTests {
 
@@ -219,6 +220,20 @@ void hikariDataSourceIsInstrumentedWithoutMetadataProvider() {
 		});
 	}
 
+	@Test
+	void additionalHikariDataSourcesCanBeInstrumented() {
+		this.contextRunner.withUserConfiguration(AdditionalHikariDataSourceConfiguration.class)
+			.withConfiguration(AutoConfigurations.of(DataSourceAutoConfiguration.class))
+			.run((context) -> {
+				assertThat(context).getBeanNames(DataSource.class)
+					.containsExactlyInAnyOrder("dataSource", "additionalHikariDataSource");
+				context.getBean("additionalHikariDataSource", DataSource.class).getConnection();
+				MeterRegistry registry = context.getBean(MeterRegistry.class);
+				assertThat(registry.get("hikaricp.connections").meter().getId().getTags())
+					.containsExactly(Tag.of("pool", "additionalHikariDataSource"));
+			});
+	}
+
 	private static HikariDataSource createHikariDataSource(String poolName) {
 		String url = "jdbc:hsqldb:mem:test-" + UUID.randomUUID();
 		HikariDataSource hikariDataSource = DataSourceBuilder.create().url(url).type(HikariDataSource.class).build();
@@ -271,6 +286,16 @@ DataSource secondOne() {
 
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	static class AdditionalHikariDataSourceConfiguration {
+
+		@Bean(defaultCandidate = false)
+		DataSource additionalHikariDataSource() {
+			return createHikariDataSource("additionalHikariDataSource");
+		}
+
+	}
+
 	@Configuration(proxyBeanMethods = false)
 	static class ProxiedHikariDataSourcesConfiguration {