From b476d368dbce9fd6dd8810f0f408ac30bb156aef Mon Sep 17 00:00:00 2001 From: Andy Wilkinson Date: Tue, 25 Jul 2023 16:27:31 +0100 Subject: [PATCH] Polish "Make HikariDataSource participate in checkpoint-restore" See gh-36422 --- .../spring-boot-autoconfigure/build.gradle | 1 - .../jdbc/DataSourceConfiguration.java | 8 +- .../autoconfigure/jdbc/HikariLifecycle.java | 225 ------------------ .../jdbc/DataSourceJmxConfigurationTests.java | 25 +- .../HikariDataSourceConfigurationTests.java | 15 ++ .../jdbc/HikariLifecycleTests.java | 76 ------ .../HikariCheckpointRestoreLifecycle.java | 149 ++++++++++++ ...HikariCheckpointRestoreLifecycleTests.java | 85 +++++++ 8 files changed, 256 insertions(+), 328 deletions(-) delete mode 100644 spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/jdbc/HikariLifecycle.java delete mode 100644 spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/jdbc/HikariLifecycleTests.java create mode 100644 spring-boot-project/spring-boot/src/main/java/org/springframework/boot/jdbc/HikariCheckpointRestoreLifecycle.java create mode 100644 spring-boot-project/spring-boot/src/test/java/org/springframework/boot/jdbc/HikariCheckpointRestoreLifecycleTests.java diff --git a/spring-boot-project/spring-boot-autoconfigure/build.gradle b/spring-boot-project/spring-boot-autoconfigure/build.gradle index 4dba7255b00..a30895a521d 100644 --- a/spring-boot-project/spring-boot-autoconfigure/build.gradle +++ b/spring-boot-project/spring-boot-autoconfigure/build.gradle @@ -67,7 +67,6 @@ dependencies { optional("org.apache.tomcat:tomcat-jdbc") optional("org.apiguardian:apiguardian-api") optional("org.apache.groovy:groovy-templates") - optional("org.crac:crac:1.3.0") optional("org.eclipse.angus:angus-mail") optional("com.github.ben-manes.caffeine:caffeine") optional("com.github.mxab.thymeleaf.extras:thymeleaf-extras-data-attribute") 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 c01d840850e..d0a1098e2e1 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; @@ -123,9 +125,9 @@ abstract class DataSourceConfiguration { } @Bean - @ConditionalOnClass(name = "org.crac.Resource") - HikariLifecycle dataSourceLifecycle(HikariDataSource hikariDataSource) { - return new HikariLifecycle(hikariDataSource); + @ConditionalOnCheckpointRestore + HikariCheckpointRestoreLifecycle hikariCheckpointRestoreLifecycle(HikariDataSource hikariDataSource) { + return new HikariCheckpointRestoreLifecycle(hikariDataSource); } } diff --git a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/jdbc/HikariLifecycle.java b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/jdbc/HikariLifecycle.java deleted file mode 100644 index 611844ce9d6..00000000000 --- a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/jdbc/HikariLifecycle.java +++ /dev/null @@ -1,225 +0,0 @@ -/* - * Copyright 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.autoconfigure.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 javax.sql.DataSource; - -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.boot.jdbc.DataSourceUnwrapper; -import org.springframework.context.Lifecycle; -import org.springframework.util.Assert; -import org.springframework.util.ReflectionUtils; - -/** - * {@link Lifecycle} management for a {@link HikariDataSource} taking care of - * {@link Lifecycle#start() starting}/{@link Lifecycle#stop() stopping} the - * {@link javax.sql.DataSource} by {@link HikariDataSource#isAllowPoolSuspension() - * eventually} suspending/resuming the underlying {@link HikariPool connection pool} and - * {@link HikariPoolMXBean#softEvictConnections() evicting} open & idle connections. - * - * @author Christoph Strobl - */ -class HikariLifecycle implements Lifecycle { - - private final HikariDataSource dataSource; - - private final LifecycleExecutor lifecycleExecutor; - - HikariLifecycle(HikariDataSource dataSource) { - - this.dataSource = dataSource; - this.lifecycleExecutor = new LifecycleExecutor(dataSource); - } - - HikariDataSource getManagedInstance() { - return this.dataSource; - } - - @Override - public void start() { - - if (this.dataSource.isRunning()) { - return; - } - - if (this.dataSource.isClosed()) { - throw new IllegalStateException("DataSource has been closed and cannot be restarted"); - } - - this.lifecycleExecutor.resume(); - } - - @Override - public void stop() { - if (this.dataSource.isRunning()) { - this.lifecycleExecutor.pause(); - } - } - - @Override - public boolean isRunning() { - return this.dataSource.isRunning(); - } - - /** - * Component to help suspend/resume a {@link HikariDataSource} by taking the pool - * suspension flag into account. Will perform best effort to make sure connections - * reported as closed buy the {@link HikariPoolMXBean} have actually been closed by - * the {@link java.util.concurrent.Executor} that is in charge of closing them. - * - * @author Christoph Strobl - */ - private static class LifecycleExecutor { - - private static final Log logger = LogFactory.getLog(LifecycleExecutor.class); - - private static final Field CLOSE_CONNECTION_EXECUTOR; - - private final HikariDataSource dataSource; - - private final Function hasOpenConnections; - - 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; - } - - LifecycleExecutor(HikariDataSource hikariDataSource) { - - this.dataSource = getUltimateTargetObject(hikariDataSource); - this.hasOpenConnections = new Function<>() { - - @Override - public Boolean apply(HikariPool pool) { - - ThreadPoolExecutor closeConnectionExecutor = (ThreadPoolExecutor) ReflectionUtils - .getField(CLOSE_CONNECTION_EXECUTOR, pool); - if (closeConnectionExecutor == null) { - throw new IllegalStateException("CloseConnectionExecutor was null"); - } - return closeConnectionExecutor.getActiveCount() > 0; - } - }; - } - - /** - * Pause the {@link HikariDataSource} and try to suspend obtaining new connections - * from the pool if possible. Will wait for connection to be closed. Default - * timeout is set to {@link HikariDataSource#getConnectionTimeout()} + 250 ms. - */ - void pause() { - pause(Duration.ofMillis(this.dataSource.getConnectionTimeout() + 250)); - } - - /** - * Pause the {@link HikariDataSource} and try to suspend obtaining new connections - * from the pool if possible. Wait at most the given {@literal shutdownTimeout} - * for connections to be closed. - * @param shutdownTimeout max timeout to wait for connections to be closed. - */ - void pause(Duration shutdownTimeout) { - - if (this.dataSource.isAllowPoolSuspension()) { - logger.info("Suspending Hikari pool"); - this.dataSource.getHikariPoolMXBean().suspendPool(); - } - closeConnections(shutdownTimeout); - } - - /** - * Resume the {@link HikariDataSource} by lifting the pool suspension if set. - */ - void resume() { - - if (this.dataSource.isAllowPoolSuspension()) { - logger.info("Resuming Hikari pool"); - this.dataSource.getHikariPoolMXBean().resumePool(); - } - } - - 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.error("Interrupted while waiting for connections to be closed", ex); - Thread.currentThread().interrupt(); - } - catch (TimeoutException ex) { - logger.error("Hikari connections could not be closed within %s".formatted(shutdownTimeout), ex); - } - catch (ExecutionException ex) { - throw new RuntimeException("Failed to close Hikari connections", ex); - } - } - - private void waitForConnectionsToClose() { - - if (!(this.dataSource.getHikariPoolMXBean() instanceof HikariPool pool)) { - throw new IllegalStateException( - "Expected HikariPool instance but was %s".formatted(this.dataSource.getHikariPoolMXBean())); - } - - while (this.hasOpenConnections.apply(pool)) { - try { - TimeUnit.MILLISECONDS.sleep(50); - } - catch (InterruptedException ex) { - logger.error("Interrupted while waiting for datasource connections to be closed", ex); - Thread.currentThread().interrupt(); - } - } - } - - @SuppressWarnings("unchecked") - private static HikariDataSource getUltimateTargetObject(DataSource candidate) { - return DataSourceUnwrapper.unwrap(candidate, HikariConfigMXBean.class, HikariDataSource.class); - } - - } - -} diff --git a/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/jdbc/DataSourceJmxConfigurationTests.java b/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/jdbc/DataSourceJmxConfigurationTests.java index 7bf81e7fe19..b8c365a8f54 100644 --- a/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/jdbc/DataSourceJmxConfigurationTests.java +++ b/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/jdbc/DataSourceJmxConfigurationTests.java @@ -17,7 +17,6 @@ package org.springframework.boot.autoconfigure.jdbc; import java.lang.management.ManagementFactory; -import java.sql.Connection; import java.util.Set; import java.util.UUID; @@ -36,7 +35,6 @@ import org.springframework.aop.framework.ProxyFactory; import org.springframework.beans.factory.config.BeanPostProcessor; import org.springframework.boot.autoconfigure.AutoConfigurations; import org.springframework.boot.autoconfigure.jmx.JmxAutoConfiguration; -import org.springframework.boot.test.context.FilteredClassLoader; import org.springframework.boot.test.context.runner.ApplicationContextRunner; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; @@ -109,8 +107,7 @@ class DataSourceJmxConfigurationTests { assertThat(hikariDataSource.isRegisterMbeans()).isTrue(); // Ensure that the pool has been initialized, triggering MBean // registration - Connection connection = hikariDataSource.getConnection(); - hikariDataSource.evictConnection(connection); + hikariDataSource.getConnection().close(); // Hikari can still register mBeans validateHikariMBeansRegistration(ManagementFactory.getPlatformMBeanServer(), poolName, true); }); @@ -135,21 +132,6 @@ class DataSourceJmxConfigurationTests { }); } - @Test - void hikariAutoConfigRegistersLifecycleBean() { - - this.contextRunner.withPropertyValues("spring.datasource.type=" + HikariDataSource.class.getName()) - .run((context) -> assertThat(context).hasSingleBean(HikariLifecycle.class)); - } - - @Test - void hikariAutoConfigConditionallyRegistersLifecycleBean() { - - this.contextRunner.withPropertyValues("spring.datasource.type=" + HikariDataSource.class.getName()) - .withClassLoader(new FilteredClassLoader("org.crac")) - .run((context) -> assertThat(context).doesNotHaveBean(HikariLifecycle.class)); - } - private void validateHikariMBeansRegistration(MBeanServer mBeanServer, String poolName, boolean expected) throws MalformedObjectNameException { assertThat(mBeanServer.isRegistered(new ObjectName("com.zaxxer.hikari:type=Pool (" + poolName + ")"))) @@ -214,10 +196,7 @@ class DataSourceJmxConfigurationTests { @Override public Object postProcessAfterInitialization(Object bean, String beanName) { if (bean instanceof javax.sql.DataSource) { - - ProxyFactory pf = new ProxyFactory(bean); - pf.setProxyTargetClass(true); - return pf.getProxy(); + return new ProxyFactory(bean).getProxy(); } return bean; } 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 653d59cb8f4..0793585ad38 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-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/jdbc/HikariLifecycleTests.java b/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/jdbc/HikariLifecycleTests.java deleted file mode 100644 index 4fea7255624..00000000000 --- a/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/jdbc/HikariLifecycleTests.java +++ /dev/null @@ -1,76 +0,0 @@ -/* - * Copyright 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.autoconfigure.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; - -/** - * Tests for {@link HikariLifecycle}. - * - * @author Christoph Strobl - */ -class HikariLifecycleTests { - - @Test - void stopStartHikariDataSource() { - - HikariLifecycle hikariLifecycle = createLifecycle(); - - assertThat(hikariLifecycle.isRunning()).isTrue(); - - hikariLifecycle.stop(); - - assertThat(hikariLifecycle.getManagedInstance().isRunning()).isFalse(); - assertThat(hikariLifecycle.getManagedInstance().isClosed()).isFalse(); - assertThat(hikariLifecycle.isRunning()).isFalse(); - assertThat(hikariLifecycle.getManagedInstance().getHikariPoolMXBean().getTotalConnections()).isZero(); - - hikariLifecycle.start(); - - assertThat(hikariLifecycle.getManagedInstance().isRunning()).isTrue(); - assertThat(hikariLifecycle.getManagedInstance().isClosed()).isFalse(); - assertThat(hikariLifecycle.isRunning()).isTrue(); - } - - @Test - void cannotStartClosedDataSource() { - - HikariLifecycle hikariLifecycle = createLifecycle(); - hikariLifecycle.getManagedInstance().close(); - - assertThatExceptionOfType(RuntimeException.class).isThrownBy(hikariLifecycle::start); - } - - HikariLifecycle createLifecycle() { - - HikariConfig config = new HikariConfig(); - config.setAllowPoolSuspension(true); - config.setJdbcUrl("jdbc:hsqldb:mem:test-" + UUID.randomUUID()); - config.setPoolName("lifecycle-tests"); - - HikariDataSource source = new HikariDataSource(config); - return new HikariLifecycle(source); - } - -} 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 00000000000..4f458e9f5af --- /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 00000000000..05811c5d593 --- /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); + } + +}