Register metrics for wrapped R2DBC ConnectionPools

Closes gh-30090
spring-projects · Mar 8, 2022 · 708e57e · 708e57e
1 parent e7705f4
commit 708e57e
Show file tree

Hide file tree

Showing 2 changed files with 69 additions and 4 deletions.
diff --git a/...work/boot/actuate/autoconfigure/metrics/r2dbc/ConnectionPoolMetricsAutoConfiguration.java b/...work/boot/actuate/autoconfigure/metrics/r2dbc/ConnectionPoolMetricsAutoConfiguration.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2020 the original author or authors.
+ * Copyright 2012-2022 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -22,6 +22,7 @@
 import io.micrometer.core.instrument.Tags;
 import io.r2dbc.pool.ConnectionPool;
 import io.r2dbc.spi.ConnectionFactory;
+import io.r2dbc.spi.Wrapped;
 
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.actuate.autoconfigure.metrics.MetricsAutoConfiguration;
@@ -53,10 +54,21 @@ public class ConnectionPoolMetricsAutoConfiguration {
 	public void bindConnectionPoolsToRegistry(Map<String, ConnectionFactory> connectionFactories,
 			MeterRegistry registry) {
 		connectionFactories.forEach((beanName, connectionFactory) -> {
-			if (connectionFactory instanceof ConnectionPool) {
-				new ConnectionPoolMetrics((ConnectionPool) connectionFactory, beanName, Tags.empty()).bindTo(registry);
+			ConnectionPool pool = extractPool(connectionFactory);
+			if (pool != null) {
+				new ConnectionPoolMetrics(pool, beanName, Tags.empty()).bindTo(registry);
 			}
 		});
 	}
 
+	private ConnectionPool extractPool(Object candidate) {
+		if (candidate instanceof ConnectionPool) {
+			return (ConnectionPool) candidate;
+		}
+		if (candidate instanceof Wrapped) {
+			return extractPool(((Wrapped<?>) candidate).unwrap());
+		}
+		return null;
+	}
+
 }
diff --git a/...boot/actuate/autoconfigure/metrics/r2dbc/ConnectionPoolMetricsAutoConfigurationTests.java b/...boot/actuate/autoconfigure/metrics/r2dbc/ConnectionPoolMetricsAutoConfigurationTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2020 the original author or authors.
+ * Copyright 2012-2022 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -27,8 +27,12 @@
 import io.r2dbc.h2.H2ConnectionOption;
 import io.r2dbc.pool.ConnectionPool;
 import io.r2dbc.pool.ConnectionPoolConfiguration;
+import io.r2dbc.spi.Connection;
 import io.r2dbc.spi.ConnectionFactory;
+import io.r2dbc.spi.ConnectionFactoryMetadata;
+import io.r2dbc.spi.Wrapped;
 import org.junit.jupiter.api.Test;
+import org.reactivestreams.Publisher;
 
 import org.springframework.boot.actuate.autoconfigure.metrics.test.MetricsRun;
 import org.springframework.boot.autoconfigure.AutoConfigurations;
@@ -89,6 +93,15 @@ void connectionPoolExposedAsConnectionFactoryTypeIsInstrumented() {
 		});
 	}
 
+	@Test
+	void wrappedConnectionPoolExposedAsConnectionFactoryTypeIsInstrumented() {
+		this.contextRunner.withUserConfiguration(WrappedConnectionPoolConfiguration.class).run((context) -> {
+			MeterRegistry registry = context.getBean(MeterRegistry.class);
+			assertThat(registry.find("r2dbc.pool.acquired").gauges()).extracting(Meter::getId)
+					.extracting((id) -> id.getTag("name")).containsExactly("wrappedConnectionPool");
+		});
+	}
+
 	@Test
 	void allConnectionPoolsCanBeInstrumented() {
 		this.contextRunner.withUserConfiguration(TwoConnectionPoolsConfiguration.class).run((context) -> {
@@ -120,6 +133,46 @@ ConnectionFactory testConnectionPool() {
 
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	static class WrappedConnectionPoolConfiguration {
+
+		@Bean
+		ConnectionFactory wrappedConnectionPool() {
+			return new Wrapper(
+					new ConnectionPool(
+							ConnectionPoolConfiguration
+									.builder(H2ConnectionFactory.inMemory("db-" + UUID.randomUUID(), "sa", "",
+											Collections.singletonMap(H2ConnectionOption.DB_CLOSE_DELAY, "-1")))
+									.build()));
+		}
+
+		static class Wrapper implements ConnectionFactory, Wrapped<ConnectionFactory> {
+
+			private final ConnectionFactory delegate;
+
+			Wrapper(ConnectionFactory delegate) {
+				this.delegate = delegate;
+			}
+
+			@Override
+			public ConnectionFactory unwrap() {
+				return this.delegate;
+			}
+
+			@Override
+			public Publisher<? extends Connection> create() {
+				return this.delegate.create();
+			}
+
+			@Override
+			public ConnectionFactoryMetadata getMetadata() {
+				return this.delegate.getMetadata();
+			}
+
+		}
+
+	}
+
 	@Configuration(proxyBeanMethods = false)
 	static class TwoConnectionPoolsConfiguration {