open DataSourcePublicMetrics for easy extension (#10288)

spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/metrics/PublicMetricsAutoConfigurationTests.java

@@ -16,19 +16,23 @@
 
 package org.springframework.boot.actuate.autoconfigure.metrics;
 
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.Assert.fail;
+import static org.mockito.BDDMockito.given;
+import static org.mockito.Mockito.mock;
+
 import java.sql.SQLException;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
 
 import javax.sql.DataSource;
 
-import com.zaxxer.hikari.HikariDataSource;
 import org.apache.commons.dbcp2.BasicDataSource;
 import org.junit.After;
 import org.junit.Test;
-
 import org.springframework.boot.actuate.autoconfigure.cache.CacheStatisticsAutoConfiguration;
 import org.springframework.boot.actuate.autoconfigure.web.servlet.MockServletWebServerFactory;
 import org.springframework.boot.actuate.cache.CachePublicMetrics;
@@ -44,6 +48,8 @@
 import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
 import org.springframework.boot.autoconfigure.jdbc.metadata.DataSourcePoolMetadataProvidersConfiguration;
 import org.springframework.boot.jdbc.DataSourceBuilder;
+import org.springframework.boot.jdbc.metadata.DataSourcePoolMetadata;
+import org.springframework.boot.jdbc.metadata.TomcatDataSourcePoolMetadata;
 import org.springframework.boot.web.embedded.tomcat.TomcatServletWebServerFactory;
 import org.springframework.boot.web.servlet.context.AnnotationConfigServletWebServerApplicationContext;
 import org.springframework.cache.CacheManager;
@@ -57,10 +63,7 @@
 import org.springframework.jdbc.core.ConnectionCallback;
 import org.springframework.jdbc.core.JdbcTemplate;
 
-import static org.assertj.core.api.Assertions.assertThat;
-import static org.junit.Assert.fail;
-import static org.mockito.BDDMockito.given;
-import static org.mockito.Mockito.mock;
+import com.zaxxer.hikari.HikariDataSource;
 
 /**
  * Tests for {@link PublicMetricsAutoConfiguration}.
@@ -168,6 +171,15 @@ public void multipleDataSourcesWithCustomPrimary() {
 				"datasource.dataSource.active", "datasource.dataSource.usage");
 	}
 
+	@Test
+	public void customDataSourcePublicMetrics() {
+		load(CustomDataSourcePublicMetricsConfig.class);
+		PublicMetrics bean = this.context.getBean(DataSourcePublicMetrics.class);
+		Collection<Metric<?>> metrics = bean.metrics();
+		assertMetrics(metrics, "datasource.primary.active", "datasource.primary.usage",
+				"datasource.primary.idle", "datasource.primary.poolSize");
+	}
+
 	@Test
 	public void customPrefix() {
 		load(MultipleDataSourcesWithPrimaryConfig.class,
@@ -306,6 +318,36 @@ public DataSource dataSource() {
 
 	}
 
+	@Configuration
+	static class CustomDataSourcePublicMetricsConfig {
+
+		@Bean
+		public DataSource myDataSource() {
+			return InitializedBuilder.create()
+					.type(org.apache.tomcat.jdbc.pool.DataSource.class).build();
+		}
+
+		@Bean
+		public TomcatDataSourcePublicMetrics tomcatDataSourcePublicMetrics() {
+			return new TomcatDataSourcePublicMetrics();
+		}
+
+		static class TomcatDataSourcePublicMetrics extends DataSourcePublicMetrics {
+
+			@Override
+			protected void addMetricsForDataSourcePoolMetadata(Set<Metric<?>> metrics, DataSourcePoolMetadata metadata, String prefix) {
+				super.addMetricsForDataSourcePoolMetadata(metrics, metadata, prefix);
+				if (metadata instanceof TomcatDataSourcePoolMetadata) {
+					TomcatDataSourcePoolMetadata tomcatMetadata = (TomcatDataSourcePoolMetadata) metadata;
+					addMetric(metrics, prefix + "idle", tomcatMetadata.getIdle());
+					addMetric(metrics, prefix + "poolSize", tomcatMetadata.getPoolSize());
+				}
+			}
+		}
+
+	}
+
+
 	@Configuration
 	static class CustomDataSourcePublicMetrics {
 

spring-boot-actuator/src/main/java/org/springframework/boot/actuate/metrics/DataSourcePublicMetrics.java

@@ -43,6 +43,9 @@ public class DataSourcePublicMetrics implements PublicMetrics {
 
 	private static final String DATASOURCE_SUFFIX = "dataSource";
 
+	protected static final String ACTIVE_METRIC_SUFFIX = "active";
+	protected static final String USAGE_METRIC_SUFFIX = "usage";
+
 	@Autowired
 	private ApplicationContext applicationContext;
 
@@ -77,13 +80,17 @@ public Collection<Metric<?>> metrics() {
 			String prefix = entry.getKey();
 			prefix = (prefix.endsWith(".") ? prefix : prefix + ".");
 			DataSourcePoolMetadata metadata = entry.getValue();
-			addMetric(metrics, prefix + "active", metadata.getActive());
-			addMetric(metrics, prefix + "usage", metadata.getUsage());
+			addMetricsForDataSourcePoolMetadata(metrics, metadata, prefix);
 		}
 		return metrics;
 	}
 
-	private <T extends Number> void addMetric(Set<Metric<?>> metrics, String name,
+	protected void addMetricsForDataSourcePoolMetadata(Set<Metric<?>> metrics, DataSourcePoolMetadata metadata, String prefix) {
+		addMetric(metrics, prefix + ACTIVE_METRIC_SUFFIX, metadata.getActive());
+		addMetric(metrics, prefix + USAGE_METRIC_SUFFIX, metadata.getUsage());
+	}
+
+	protected  <T extends Number> void addMetric(Set<Metric<?>> metrics, String name,
 			T value) {
 		if (value != null) {
 			metrics.add(new Metric<>(name, value));

spring-boot/src/main/java/org/springframework/boot/jdbc/metadata/TomcatDataSourcePoolMetadata.java

@@ -53,4 +53,11 @@ public String getValidationQuery() {
 		return getDataSource().getValidationQuery();
 	}
 
+	public Integer getIdle() {
+		return getDataSource().getIdle();
+	}
+
+	public Integer getPoolSize() {
+		return getDataSource().getPoolSize();
+	}
 }