Extend CRaC checkpoint/restore support to Oracle UCP

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

@@ -25,15 +25,11 @@
 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.jspecify.annotations.Nullable;
 
 import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.context.Lifecycle;
@@ -52,6 +48,7 @@
  * @author Christoph Strobl
  * @author Andy Wilkinson
  * @author Moritz Halbritter
+ * @author Fabio Grassi
  * @since 3.2.0
  */
 public class HikariCheckpointRestoreLifecycle implements Lifecycle {
@@ -72,21 +69,20 @@ public class HikariCheckpointRestoreLifecycle implements Lifecycle {
 
 	private final Function<HikariPool, Boolean> hasOpenConnections;
 
-	private final @Nullable HikariDataSource dataSource;
+	private final HikariDataSource dataSource;
 
 	private final ConfigurableApplicationContext applicationContext;
 
 	/**
 	 * Creates a new {@code HikariCheckpointRestoreLifecycle} that will allow the given
-	 * {@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.
+	 * {@link HikariDataSource} to participate in checkpoint-restore.
 	 * @param dataSource the checkpoint-restore participant
 	 * @param applicationContext the application context
 	 * @since 3.4.0
 	 */
-	public HikariCheckpointRestoreLifecycle(DataSource dataSource, ConfigurableApplicationContext applicationContext) {
-		this.dataSource = DataSourceUnwrapper.unwrap(dataSource, HikariConfigMXBean.class, HikariDataSource.class);
+	public HikariCheckpointRestoreLifecycle(HikariDataSource dataSource,
+			ConfigurableApplicationContext applicationContext) {
+		this.dataSource = dataSource;
 		this.applicationContext = applicationContext;
 		this.hasOpenConnections = (pool) -> {
 			ThreadPoolExecutor closeConnectionExecutor = (ThreadPoolExecutor) ReflectionUtils
@@ -98,7 +94,7 @@ public HikariCheckpointRestoreLifecycle(DataSource dataSource, ConfigurableAppli
 
 	@Override
 	public void start() {
-		if (this.dataSource == null || this.dataSource.isRunning()) {
+		if (this.dataSource.isRunning()) {
 			return;
 		}
 		Assert.state(!this.dataSource.isClosed(), "DataSource has been closed and cannot be restarted");
@@ -110,7 +106,7 @@ public void start() {
 
 	@Override
 	public void stop() {
-		if (this.dataSource == null || !this.dataSource.isRunning()) {
+		if (!this.dataSource.isRunning()) {
 			return;
 		}
 		if (this.dataSource.isAllowPoolSuspension()) {
@@ -164,7 +160,7 @@ private void waitForConnectionsToClose(HikariDataSource dataSource) {
 
 	@Override
 	public boolean isRunning() {
-		return this.dataSource != null && this.dataSource.isRunning();
+		return this.dataSource.isRunning();
 	}
 
 }

module/spring-boot-jdbc/src/main/java/org/springframework/boot/jdbc/OracleUcpCheckpointRestoreLifecycle.java

@@ -0,0 +1,128 @@
+/*
+ * Copyright 2012-present 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.sql.SQLException;
+import java.util.Arrays;
+
+import oracle.ucp.UniversalConnectionPoolException;
+import oracle.ucp.admin.UniversalConnectionPoolManager;
+import oracle.ucp.admin.UniversalConnectionPoolManagerImpl;
+import oracle.ucp.jdbc.JDBCConnectionPool;
+import oracle.ucp.jdbc.PoolDataSourceImpl;
+import org.jspecify.annotations.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.springframework.context.Lifecycle;
+import org.springframework.util.Assert;
+
+/**
+ * A {@link Lifecycle} over the connection pool of a single
+ * {@link oracle.ucp.jdbc.PoolDataSourceImpl}, which lets the pool be started and stopped
+ * along with the application context without being destroyed in between.
+ * <p>
+ * {@link #start()} creates the pool when it does not exist yet, since a pool data source
+ * has no {@code connectionPoolName} until then, and UCP registers a freshly created pool
+ * in the stopped state. Both {@code start()} and {@link #stop()} guard on the current
+ * life cycle state, as UCP rejects a transition that has already happened.
+ *
+ * @author Fabio Grassi
+ * @since 4.1.0
+ */
+public final class OracleUcpCheckpointRestoreLifecycle implements Lifecycle {
+
+	private static final Logger logger = LoggerFactory.getLogger(OracleUcpCheckpointRestoreLifecycle.class);
+
+	private final PoolDataSourceImpl poolDataSource;
+
+	public OracleUcpCheckpointRestoreLifecycle(final PoolDataSourceImpl poolDataSource) {
+		Assert.notNull(poolDataSource, "Non null PoolDataSourceImpl instance expected");
+		this.poolDataSource = poolDataSource;
+	}
+
+	@Override
+	public void start() {
+		JDBCConnectionPool pool = getPool(this.poolDataSource.getConnectionPoolName());
+		if (pool == null) {
+			pool = createPool();
+			logger.info("Created new Oracle Universal Connection Pool named '{}'", pool.getName());
+		}
+		if (!pool.isLifecycleRunning() && !pool.isLifecycleStarting()) {
+			doWithPool(pool::start);
+			logger.info("Oracle Universal Connection Pool '{}' started", pool.getName());
+		}
+	}
+
+	@Override
+	public void stop() {
+		final JDBCConnectionPool pool = getPool(this.poolDataSource.getConnectionPoolName());
+		if (pool != null && !pool.isLifecycleStopped() && !pool.isLifecycleStopping()) {
+			doWithPool(pool::stop);
+			logger.info("Oracle Universal Connection Pool '{}' stopped", pool.getName());
+		}
+	}
+
+	@Override
+	public boolean isRunning() {
+		final JDBCConnectionPool pool = getPool(this.poolDataSource.getConnectionPoolName());
+		final boolean isRunning = pool != null && pool.isLifecycleRunning();
+		logger.info("Oracle Universal Connection Pool '{}' is {}running", this.poolDataSource.getConnectionPoolName(),
+				isRunning ? "" : "not ");
+		return isRunning;
+	}
+
+	private JDBCConnectionPool createPool() {
+		try {
+			return (JDBCConnectionPool) this.poolDataSource.createUniversalConnectionPool();
+		}
+		catch (SQLException sqle) {
+			throw new IllegalStateException("Failed to create new Oracle Universal Connection Pool", sqle);
+		}
+	}
+
+	private static @Nullable JDBCConnectionPool getPool(final @Nullable String poolName) {
+		try {
+			final UniversalConnectionPoolManager mgr = UniversalConnectionPoolManagerImpl
+				.getUniversalConnectionPoolManager();
+			if (Arrays.asList(mgr.getConnectionPoolNames()).contains(poolName)) {
+				return (JDBCConnectionPool) mgr.getConnectionPool(poolName);
+			}
+		}
+		catch (UniversalConnectionPoolException ucpe) {
+			throw new IllegalStateException("Failed to retrieve existing Oracle Universal Connection Pool", ucpe);
+		}
+		return null;
+	}
+
+	private static void doWithPool(final PoolCommand command) {
+		try {
+			command.execute();
+		}
+		catch (UniversalConnectionPoolException ucpe) {
+			throw new IllegalStateException("Oracle Universal Connection Pool command failed", ucpe);
+		}
+	}
+
+	@FunctionalInterface
+	private interface PoolCommand {
+
+		void execute() throws UniversalConnectionPoolException;
+
+	}
+
+}

module/spring-boot-jdbc/src/main/java/org/springframework/boot/jdbc/autoconfigure/DataSourceCheckpointRestoreConfiguration.java

@@ -16,27 +16,43 @@
 
 package org.springframework.boot.jdbc.autoconfigure;
 
+import java.util.Collection;
+import java.util.LinkedList;
+import java.util.function.Function;
+
 import javax.sql.DataSource;
 
+import com.zaxxer.hikari.HikariConfigMXBean;
 import com.zaxxer.hikari.HikariDataSource;
+import oracle.jdbc.OracleConnection;
+import oracle.ucp.jdbc.PoolDataSource;
+import oracle.ucp.jdbc.PoolDataSourceImpl;
 
-import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
+import org.springframework.beans.factory.ObjectProvider;
+import org.springframework.beans.factory.SmartInitializingSingleton;
+import org.springframework.boot.autoconfigure.condition.AnyNestedCondition;
 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.jdbc.DataSourceUnwrapper;
 import org.springframework.boot.jdbc.HikariCheckpointRestoreLifecycle;
+import org.springframework.boot.jdbc.OracleUcpCheckpointRestoreLifecycle;
+import org.springframework.boot.jdbc.autoconfigure.DataSourceCheckpointRestoreConfiguration.CheckpointRestorePoolsAvailableCondition;
 import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.context.Lifecycle;
 import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Conditional;
 import org.springframework.context.annotation.Configuration;
 
 /**
  * Checkpoint-restore specific configuration.
  *
  * @author Olga Maciaszek-Sharma
+ * @author Fabio Grassi
  */
 @Configuration(proxyBeanMethods = false)
 @ConditionalOnCheckpointRestore
-@ConditionalOnBean(DataSource.class)
+@Conditional(CheckpointRestorePoolsAvailableCondition.class)
 class DataSourceCheckpointRestoreConfiguration {
 
 	@Configuration(proxyBeanMethods = false)
@@ -45,9 +61,125 @@ static class Hikari {
 
 		@Bean
 		@ConditionalOnMissingBean
-		HikariCheckpointRestoreLifecycle hikariCheckpointRestoreLifecycle(DataSource dataSource,
-				ConfigurableApplicationContext applicationContext) {
-			return new HikariCheckpointRestoreLifecycle(dataSource, applicationContext);
+		HikariCheckpointRestoreLifecycleRegistry hikariCheckpointRestoreLifecycle(
+				final ObjectProvider<DataSource> dataSources, final ConfigurableApplicationContext applicationContext) {
+			return new HikariCheckpointRestoreLifecycleRegistry(dataSources, applicationContext);
+		}
+
+		static final class HikariCheckpointRestoreLifecycleRegistry
+				extends DataSourceCheckpointRestoreLifecycleRegistry<HikariConfigMXBean, HikariDataSource> {
+
+			HikariCheckpointRestoreLifecycleRegistry(final ObjectProvider<DataSource> dataSources,
+					final ConfigurableApplicationContext applicationContext) {
+				super(dataSources, HikariConfigMXBean.class, HikariDataSource.class,
+						hds -> new HikariCheckpointRestoreLifecycle(hds, applicationContext));
+			}
+
+		}
+
+	}
+
+	@Configuration(proxyBeanMethods = false)
+	@ConditionalOnClass({ PoolDataSourceImpl.class, OracleConnection.class })
+	static class OracleUcp {
+
+		@Bean
+		@ConditionalOnMissingBean
+		OracleUcpCheckpointRestoreLifecycleRegistry oracleUcpCheckpointRestoreLifecycle(
+				final ObjectProvider<DataSource> dataSources) {
+			return new OracleUcpCheckpointRestoreLifecycleRegistry(dataSources);
+		}
+
+		static final class OracleUcpCheckpointRestoreLifecycleRegistry
+				extends DataSourceCheckpointRestoreLifecycleRegistry<PoolDataSource, PoolDataSourceImpl> {
+
+			OracleUcpCheckpointRestoreLifecycleRegistry(final ObjectProvider<DataSource> dataSources) {
+				super(dataSources, PoolDataSource.class, PoolDataSourceImpl.class,
+						OracleUcpCheckpointRestoreLifecycle::new);
+			}
+
+		}
+
+	}
+
+	static class CheckpointRestorePoolsAvailableCondition extends AnyNestedCondition {
+
+		CheckpointRestorePoolsAvailableCondition() {
+			super(ConfigurationPhase.PARSE_CONFIGURATION);
+		}
+
+		@ConditionalOnClass(HikariDataSource.class)
+		static class HickariAvailable {
+
+		}
+
+		@ConditionalOnClass({ PoolDataSourceImpl.class, OracleConnection.class })
+		static class OracleUcpAvailable {
+
+		}
+
+	}
+
+	/**
+	 * A {@link Lifecycle} container that propagates {@code start()} and {@code stop()}
+	 * signals to all its elements and {@code isRunning()} if and only if all its elements
+	 * are running or there are no elements.
+	 * <p>
+	 * This class implements also {@link SmartInitializingSingleton} to hook into the bean
+	 * factory lifecyle after all singleton beans registration and iterate over all
+	 * {@code DataSource}s, including the ones that are neither default nor autowire
+	 * candidates, unwrap each of them to reach the underlying data source, supply it to
+	 * the given factory to create a {@code Lifecycle} and add it its elements.
+	 *
+	 * @author Fabio Grassi
+	 * @since 4.1.0
+	 */
+	static sealed class DataSourceCheckpointRestoreLifecycleRegistry<I, T extends I>
+			implements SmartInitializingSingleton, Lifecycle {
+
+		private final ObjectProvider<DataSource> dataSources;
+
+		private final Class<I> wrappingInterface;
+
+		private final Class<T> targetClass;
+
+		private final Function<T, Lifecycle> lifecycleFactory;
+
+		private final Collection<Lifecycle> lifecycles;
+
+		DataSourceCheckpointRestoreLifecycleRegistry(final ObjectProvider<DataSource> dataSources,
+				final Class<I> wrappingInterface, final Class<T> targetClass,
+				final Function<T, Lifecycle> lifecycleFactory) {
+			this.dataSources = dataSources;
+			this.wrappingInterface = wrappingInterface;
+			this.targetClass = targetClass;
+			this.lifecycleFactory = lifecycleFactory;
+			this.lifecycles = new LinkedList<>();
+		}
+
+		@Override
+		public void afterSingletonsInstantiated() {
+			this.dataSources.stream(ObjectProvider.UNFILTERED, false).forEach(ds -> {
+				final T unwrapped = DataSourceUnwrapper.unwrap(ds, this.wrappingInterface, this.targetClass);
+				if (unwrapped != null) {
+					this.lifecycles.add(this.lifecycleFactory.apply(unwrapped));
+				}
+			});
+		}
+
+		@Override
+		public void start() {
+			this.lifecycles.forEach(Lifecycle::start);
+		}
+
+		@Override
+		public void stop() {
+			this.lifecycles.forEach(Lifecycle::stop);
+		}
+
+		@Override
+		public boolean isRunning() {
+			return this.lifecycles.stream().allMatch(Lifecycle::isRunning);
 		}
 
 	}