Fix checkpoint-restore with replaced or wrapped HikariDataSource

Closes gh-37580

Author: wilkinsona

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

@@ -126,7 +126,7 @@ HikariDataSource dataSource(DataSourceProperties properties, JdbcConnectionDetai
 
 		@Bean
 		@ConditionalOnCheckpointRestore
-		HikariCheckpointRestoreLifecycle hikariCheckpointRestoreLifecycle(HikariDataSource hikariDataSource) {
+		HikariCheckpointRestoreLifecycle hikariCheckpointRestoreLifecycle(DataSource hikariDataSource) {
 			return new HikariCheckpointRestoreLifecycle(hikariDataSource);
 		}
 

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

@@ -22,12 +22,15 @@
 import org.assertj.core.api.InstanceOfAssertFactories;
 import org.junit.jupiter.api.Test;
 
+import org.springframework.beans.BeansException;
+import org.springframework.beans.factory.config.BeanPostProcessor;
 import org.springframework.boot.autoconfigure.AutoConfigurations;
 import org.springframework.boot.jdbc.HikariCheckpointRestoreLifecycle;
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
 import org.springframework.boot.testsupport.classpath.ClassPathOverrides;
 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;
 
@@ -131,6 +134,14 @@ void whenCheckpointRestoreIsAvailableHikariAutoConfigRegistersLifecycleBean() {
 			.run((context) -> assertThat(context).hasSingleBean(HikariCheckpointRestoreLifecycle.class));
 	}
 
+	@Test
+	@ClassPathOverrides("org.crac:crac:1.3.0")
+	void whenCheckpointRestoreIsAvailableAndDataSourceHasBeenWrappedHikariAutoConfigRegistersLifecycleBean() {
+		this.contextRunner.withUserConfiguration(DataSourceWrapperConfiguration.class)
+			.withPropertyValues("spring.datasource.type=" + HikariDataSource.class.getName())
+			.run((context) -> assertThat(context).hasSingleBean(HikariCheckpointRestoreLifecycle.class));
+	}
+
 	@Test
 	void whenCheckpointRestoreIsNotAvailableHikariAutoConfigDoesNotRegisterLifecycleBean() {
 		this.contextRunner.withPropertyValues("spring.datasource.type=" + HikariDataSource.class.getName())
@@ -147,4 +158,24 @@ JdbcConnectionDetails sqlConnectionDetails() {
 
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	static class DataSourceWrapperConfiguration {
+
+		@Bean
+		static BeanPostProcessor dataSourceWrapper() {
+			return new BeanPostProcessor() {
+
+				@Override
+				public Object postProcessBeforeInitialization(Object bean, String beanName) throws BeansException {
+					if (bean instanceof DataSource dataSource) {
+						return new DelegatingDataSource(dataSource);
+					}
+					return bean;
+				}
+
+			};
+		}
+
+	}
+
 }

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

@@ -25,6 +25,8 @@
 import java.util.concurrent.TimeoutException;
 import java.util.function.Function;
 
+import javax.sql.DataSource;
+
 import com.zaxxer.hikari.HikariConfigMXBean;
 import com.zaxxer.hikari.HikariDataSource;
 import com.zaxxer.hikari.HikariPoolMXBean;
@@ -71,10 +73,12 @@ public class HikariCheckpointRestoreLifecycle implements Lifecycle {
 
 	/**
 	 * Creates a new {@code HikariCheckpointRestoreLifecycle} that will allow the given
-	 * {@code dataSource} to participate in checkpoint-restore.
+	 * {@code dataSource} to participate in checkpoint-restore. The {@code dataSource} is
+	 * {@link DataSourceUnwrapper#unwrap unwrapped} to a {@link HikariDataSource}. If such
+	 * unwrapping is not possible, the lifecycle will have no effect.
 	 * @param dataSource the checkpoint-restore participant
 	 */
-	public HikariCheckpointRestoreLifecycle(HikariDataSource dataSource) {
+	public HikariCheckpointRestoreLifecycle(DataSource dataSource) {
 		this.dataSource = DataSourceUnwrapper.unwrap(dataSource, HikariConfigMXBean.class, HikariDataSource.class);
 		this.hasOpenConnections = (pool) -> {
 			ThreadPoolExecutor closeConnectionExecutor = (ThreadPoolExecutor) ReflectionUtils
@@ -86,7 +90,7 @@ public HikariCheckpointRestoreLifecycle(HikariDataSource dataSource) {
 
 	@Override
 	public void start() {
-		if (this.dataSource.isRunning()) {
+		if (this.dataSource == null || this.dataSource.isRunning()) {
 			return;
 		}
 		Assert.state(!this.dataSource.isClosed(), "DataSource has been closed and cannot be restarted");
@@ -98,7 +102,7 @@ public void start() {
 
 	@Override
 	public void stop() {
-		if (!this.dataSource.isRunning()) {
+		if (this.dataSource == null || !this.dataSource.isRunning()) {
 			return;
 		}
 		if (this.dataSource.isAllowPoolSuspension()) {
@@ -143,7 +147,7 @@ private void waitForConnectionsToClose() {
 
 	@Override
 	public boolean isRunning() {
-		return this.dataSource.isRunning();
+		return this.dataSource != null && this.dataSource.isRunning();
 	}
 
 }

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

@@ -18,13 +18,16 @@
 
 import java.util.UUID;
 
+import javax.sql.DataSource;
+
 import com.zaxxer.hikari.HikariConfig;
 import com.zaxxer.hikari.HikariDataSource;
 import org.junit.jupiter.api.Test;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
 import static org.assertj.core.api.Assertions.assertThatNoException;
+import static org.mockito.Mockito.mock;
 
 /**
  * Tests for {@link HikariCheckpointRestoreLifecycle}.
@@ -82,4 +85,22 @@ void whenDataSourceIsClosedThenStartShouldThrow() {
 		assertThatExceptionOfType(RuntimeException.class).isThrownBy(this.lifecycle::start);
 	}
 
+	@Test
+	void startHasNoEffectWhenDataSourceIsNotAHikariDataSource() {
+		HikariCheckpointRestoreLifecycle nonHikariLifecycle = new HikariCheckpointRestoreLifecycle(
+				mock(DataSource.class));
+		assertThat(nonHikariLifecycle.isRunning()).isFalse();
+		nonHikariLifecycle.start();
+		assertThat(nonHikariLifecycle.isRunning()).isFalse();
+	}
+
+	@Test
+	void stopHasNoEffectWhenDataSourceIsNotAHikariDataSource() {
+		HikariCheckpointRestoreLifecycle nonHikariLifecycle = new HikariCheckpointRestoreLifecycle(
+				mock(DataSource.class));
+		assertThat(nonHikariLifecycle.isRunning()).isFalse();
+		nonHikariLifecycle.stop();
+		assertThat(nonHikariLifecycle.isRunning()).isFalse();
+	}
+
 }