Make HikariDataSource participate in checkpoint-restore

See gh-36422

spring-boot-project/spring-boot-autoconfigure/build.gradle

@@ -67,6 +67,7 @@ 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")

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

@@ -122,6 +122,12 @@ abstract class DataSourceConfiguration {
 			return dataSource;
 		}
 
+		@Bean
+		@ConditionalOnClass(name = "org.crac.Resource")
+		HikariLifecycle dataSourceLifecycle(HikariDataSource hikariDataSource) {
+			return new HikariLifecycle(hikariDataSource);
+		}
+
 	}
 
 	/**

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

@@ -0,0 +1,225 @@
+/*
+ * 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<HikariPool, Boolean> 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<Void> 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);
+		}
+
+	}
+
+}

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

@@ -17,6 +17,7 @@
 package org.springframework.boot.autoconfigure.jdbc;
 
 import java.lang.management.ManagementFactory;
+import java.sql.Connection;
 import java.util.Set;
 import java.util.UUID;
 
@@ -35,6 +36,7 @@ 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;
@@ -107,7 +109,8 @@ class DataSourceJmxConfigurationTests {
 				assertThat(hikariDataSource.isRegisterMbeans()).isTrue();
 				// Ensure that the pool has been initialized, triggering MBean
 				// registration
-				hikariDataSource.getConnection().close();
+				Connection connection = hikariDataSource.getConnection();
+				hikariDataSource.evictConnection(connection);
 				// Hikari can still register mBeans
 				validateHikariMBeansRegistration(ManagementFactory.getPlatformMBeanServer(), poolName, true);
 			});
@@ -132,6 +135,21 @@ 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 + ")")))
@@ -196,7 +214,10 @@ class DataSourceJmxConfigurationTests {
 		@Override
 		public Object postProcessAfterInitialization(Object bean, String beanName) {
 			if (bean instanceof javax.sql.DataSource) {
-				return new ProxyFactory(bean).getProxy();
+
+				ProxyFactory pf = new ProxyFactory(bean);
+				pf.setProxyTargetClass(true);
+				return pf.getProxy();
 			}
 			return bean;
 		}

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

@@ -0,0 +1,76 @@
+/*
+ * 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);
+	}
+
+}