From ab06d10d642b5bacdf0206b2894394a805876d84 Mon Sep 17 00:00:00 2001 From: Andy Wilkinson Date: Tue, 26 Sep 2023 12:31:12 +0100 Subject: [PATCH] Fix checkpoint-restore with replaced or wrapped HikariDataSource Closes gh-37580 --- .../jdbc/DataSourceConfiguration.java | 2 +- .../HikariDataSourceConfigurationTests.java | 31 +++++++++++++++++++ .../HikariCheckpointRestoreLifecycle.java | 14 ++++++--- ...HikariCheckpointRestoreLifecycleTests.java | 21 +++++++++++++ 4 files changed, 62 insertions(+), 6 deletions(-) diff --git a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/jdbc/DataSourceConfiguration.java b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/jdbc/DataSourceConfiguration.java index d0a1098e2e..6389876a24 100644 --- a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/jdbc/DataSourceConfiguration.java +++ b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/jdbc/DataSourceConfiguration.java @@ -126,7 +126,7 @@ abstract class DataSourceConfiguration { @Bean @ConditionalOnCheckpointRestore - HikariCheckpointRestoreLifecycle hikariCheckpointRestoreLifecycle(HikariDataSource hikariDataSource) { + HikariCheckpointRestoreLifecycle hikariCheckpointRestoreLifecycle(DataSource hikariDataSource) { return new HikariCheckpointRestoreLifecycle(hikariDataSource); } diff --git a/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/jdbc/HikariDataSourceConfigurationTests.java b/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/jdbc/HikariDataSourceConfigurationTests.java index 0793585ad3..617f276456 100644 --- a/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/jdbc/HikariDataSourceConfigurationTests.java +++ b/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/jdbc/HikariDataSourceConfigurationTests.java @@ -22,12 +22,15 @@ import com.zaxxer.hikari.HikariDataSource; 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 @@ class HikariDataSourceConfigurationTests { .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 @@ class HikariDataSourceConfigurationTests { } + @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; + } + + }; + } + + } + } diff --git a/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/jdbc/HikariCheckpointRestoreLifecycle.java b/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/jdbc/HikariCheckpointRestoreLifecycle.java index 4f458e9f5a..40585e4885 100644 --- a/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/jdbc/HikariCheckpointRestoreLifecycle.java +++ b/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/jdbc/HikariCheckpointRestoreLifecycle.java @@ -25,6 +25,8 @@ import java.util.concurrent.TimeUnit; 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 class HikariCheckpointRestoreLifecycle implements Lifecycle { @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 class HikariCheckpointRestoreLifecycle implements Lifecycle { @Override public void stop() { - if (!this.dataSource.isRunning()) { + if (this.dataSource == null || !this.dataSource.isRunning()) { return; } if (this.dataSource.isAllowPoolSuspension()) { @@ -143,7 +147,7 @@ public class HikariCheckpointRestoreLifecycle implements Lifecycle { @Override public boolean isRunning() { - return this.dataSource.isRunning(); + return this.dataSource != null && this.dataSource.isRunning(); } } diff --git a/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/jdbc/HikariCheckpointRestoreLifecycleTests.java b/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/jdbc/HikariCheckpointRestoreLifecycleTests.java index 05811c5d59..119af02e60 100644 --- a/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/jdbc/HikariCheckpointRestoreLifecycleTests.java +++ b/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/jdbc/HikariCheckpointRestoreLifecycleTests.java @@ -18,6 +18,8 @@ package org.springframework.boot.jdbc; import java.util.UUID; +import javax.sql.DataSource; + import com.zaxxer.hikari.HikariConfig; import com.zaxxer.hikari.HikariDataSource; import org.junit.jupiter.api.Test; @@ -25,6 +27,7 @@ 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 @@ class HikariCheckpointRestoreLifecycleTests { 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(); + } + }