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 8dd321ee7c..d0a1098e2e 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 @@ -25,12 +25,14 @@ import oracle.jdbc.OracleConnection; import oracle.ucp.jdbc.PoolDataSourceImpl; import org.springframework.beans.factory.ObjectProvider; +import org.springframework.boot.autoconfigure.condition.ConditionalOnCheckpointRestore; import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; import org.springframework.boot.context.properties.ConfigurationProperties; import org.springframework.boot.jdbc.DataSourceBuilder; import org.springframework.boot.jdbc.DatabaseDriver; +import org.springframework.boot.jdbc.HikariCheckpointRestoreLifecycle; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; import org.springframework.util.StringUtils; @@ -122,6 +124,12 @@ abstract class DataSourceConfiguration { return dataSource; } + @Bean + @ConditionalOnCheckpointRestore + HikariCheckpointRestoreLifecycle hikariCheckpointRestoreLifecycle(HikariDataSource 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 653d59cb8f..0793585ad3 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 @@ -23,7 +23,9 @@ import org.assertj.core.api.InstanceOfAssertFactories; import org.junit.jupiter.api.Test; 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; @@ -122,6 +124,19 @@ class HikariDataSourceConfigurationTests { }); } + @Test + @ClassPathOverrides("org.crac:crac:1.3.0") + void whenCheckpointRestoreIsAvailableHikariAutoConfigRegistersLifecycleBean() { + this.contextRunner.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()) + .run((context) -> assertThat(context).doesNotHaveBean(HikariCheckpointRestoreLifecycle.class)); + } + @Configuration(proxyBeanMethods = false) static class ConnectionDetailsConfiguration { 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 new file mode 100644 index 0000000000..4f458e9f5a --- /dev/null +++ b/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/jdbc/HikariCheckpointRestoreLifecycle.java @@ -0,0 +1,149 @@ +/* + * Copyright 2012-2023 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. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.boot.jdbc; + +import java.lang.reflect.Field; +import java.time.Duration; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.function.Function; + +import com.zaxxer.hikari.HikariConfigMXBean; +import com.zaxxer.hikari.HikariDataSource; +import com.zaxxer.hikari.HikariPoolMXBean; +import com.zaxxer.hikari.pool.HikariPool; +import org.apache.commons.logging.Log; +import org.apache.commons.logging.LogFactory; + +import org.springframework.context.Lifecycle; +import org.springframework.core.log.LogMessage; +import org.springframework.util.Assert; +import org.springframework.util.ReflectionUtils; + +/** + * {@link Lifecycle} for a {@link HikariDataSource} allowing it to participate in + * checkpoint-restore. When {@link #stop() stopped}, and the data source + * {@link HikariDataSource#isAllowPoolSuspension() allows it}, its pool is suspended, + * blocking any attempts to borrow connections. Open and idle connections are then + * evicted. When subsequently {@link #start() started}, the pool is + * {@link HikariPoolMXBean#resumePool() resumed} if necessary. + * + * @author Christoph Strobl + * @author Andy Wilkinson + * @since 3.2.0 + */ +public class HikariCheckpointRestoreLifecycle implements Lifecycle { + + private static final Log logger = LogFactory.getLog(HikariCheckpointRestoreLifecycle.class); + + private static final Field CLOSE_CONNECTION_EXECUTOR; + + static { + Field closeConnectionExecutor = ReflectionUtils.findField(HikariPool.class, "closeConnectionExecutor"); + Assert.notNull(closeConnectionExecutor, "Unable to locate closeConnectionExecutor for HikariPool"); + Assert.isAssignable(ThreadPoolExecutor.class, closeConnectionExecutor.getType(), + "Expected ThreadPoolExecutor for closeConnectionExecutor but found %s" + .formatted(closeConnectionExecutor.getType())); + ReflectionUtils.makeAccessible(closeConnectionExecutor); + CLOSE_CONNECTION_EXECUTOR = closeConnectionExecutor; + } + + private final Function hasOpenConnections; + + private final HikariDataSource dataSource; + + /** + * Creates a new {@code HikariCheckpointRestoreLifecycle} that will allow the given + * {@code dataSource} to participate in checkpoint-restore. + * @param dataSource the checkpoint-restore participant + */ + public HikariCheckpointRestoreLifecycle(HikariDataSource dataSource) { + this.dataSource = DataSourceUnwrapper.unwrap(dataSource, HikariConfigMXBean.class, HikariDataSource.class); + this.hasOpenConnections = (pool) -> { + ThreadPoolExecutor closeConnectionExecutor = (ThreadPoolExecutor) ReflectionUtils + .getField(CLOSE_CONNECTION_EXECUTOR, pool); + Assert.notNull(closeConnectionExecutor, "CloseConnectionExecutor was null"); + return closeConnectionExecutor.getActiveCount() > 0; + }; + } + + @Override + public void start() { + if (this.dataSource.isRunning()) { + return; + } + Assert.state(!this.dataSource.isClosed(), "DataSource has been closed and cannot be restarted"); + if (this.dataSource.isAllowPoolSuspension()) { + logger.info("Resuming Hikari pool"); + this.dataSource.getHikariPoolMXBean().resumePool(); + } + } + + @Override + public void stop() { + if (!this.dataSource.isRunning()) { + return; + } + if (this.dataSource.isAllowPoolSuspension()) { + logger.info("Suspending Hikari pool"); + this.dataSource.getHikariPoolMXBean().suspendPool(); + } + closeConnections(Duration.ofMillis(this.dataSource.getConnectionTimeout() + 250)); + } + + private void closeConnections(Duration shutdownTimeout) { + logger.info("Evicting Hikari connections"); + this.dataSource.getHikariPoolMXBean().softEvictConnections(); + logger.debug("Waiting for Hikari connections to be closed"); + CompletableFuture allConnectionsClosed = CompletableFuture.runAsync(this::waitForConnectionsToClose); + try { + allConnectionsClosed.get(shutdownTimeout.toMillis(), TimeUnit.MILLISECONDS); + logger.debug("Hikari connections closed"); + } + catch (InterruptedException ex) { + logger.warn("Interrupted while waiting for connections to be closed", ex); + Thread.currentThread().interrupt(); + } + catch (TimeoutException ex) { + logger.warn(LogMessage.format("Hikari connections could not be closed within %s", shutdownTimeout), ex); + } + catch (ExecutionException ex) { + throw new RuntimeException("Failed to close Hikari connections", ex); + } + } + + private void waitForConnectionsToClose() { + while (this.hasOpenConnections.apply((HikariPool) this.dataSource.getHikariPoolMXBean())) { + try { + TimeUnit.MILLISECONDS.sleep(50); + } + catch (InterruptedException ex) { + logger.error("Interrupted while waiting for datasource connections to be closed", ex); + Thread.currentThread().interrupt(); + } + } + } + + @Override + public boolean isRunning() { + return 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 new file mode 100644 index 0000000000..05811c5d59 --- /dev/null +++ b/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/jdbc/HikariCheckpointRestoreLifecycleTests.java @@ -0,0 +1,85 @@ +/* + * Copyright 2012-2023 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. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.boot.jdbc; + +import java.util.UUID; + +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; + +/** + * Tests for {@link HikariCheckpointRestoreLifecycle}. + * + * @author Christoph Strobl + * @author Andy Wilkinson + */ +class HikariCheckpointRestoreLifecycleTests { + + private final HikariCheckpointRestoreLifecycle lifecycle; + + private final HikariDataSource dataSource; + + HikariCheckpointRestoreLifecycleTests() { + HikariConfig config = new HikariConfig(); + config.setAllowPoolSuspension(true); + config.setJdbcUrl("jdbc:hsqldb:mem:test-" + UUID.randomUUID()); + config.setPoolName("lifecycle-tests"); + this.dataSource = new HikariDataSource(config); + this.lifecycle = new HikariCheckpointRestoreLifecycle(this.dataSource); + } + + @Test + void startedWhenStartedShouldSucceed() { + assertThat(this.lifecycle.isRunning()).isTrue(); + this.lifecycle.start(); + assertThat(this.lifecycle.isRunning()).isTrue(); + } + + @Test + void stopWhenStoppedShouldSucceed() { + assertThat(this.lifecycle.isRunning()).isTrue(); + this.lifecycle.stop(); + assertThat(this.dataSource.isRunning()).isFalse(); + assertThatNoException().isThrownBy(this.lifecycle::stop); + } + + @Test + void whenStoppedAndStartedDataSourceShouldPauseAndResume() { + assertThat(this.lifecycle.isRunning()).isTrue(); + this.lifecycle.stop(); + assertThat(this.dataSource.isRunning()).isFalse(); + assertThat(this.dataSource.isClosed()).isFalse(); + assertThat(this.lifecycle.isRunning()).isFalse(); + assertThat(this.dataSource.getHikariPoolMXBean().getTotalConnections()).isZero(); + this.lifecycle.start(); + assertThat(this.dataSource.isRunning()).isTrue(); + assertThat(this.dataSource.isClosed()).isFalse(); + assertThat(this.lifecycle.isRunning()).isTrue(); + } + + @Test + void whenDataSourceIsClosedThenStartShouldThrow() { + this.dataSource.close(); + assertThatExceptionOfType(RuntimeException.class).isThrownBy(this.lifecycle::start); + } + +}