Do not block in DSInitializedPublisher when EMF is using async bootstrap

Previously, we would retrieve the EntityManagerFactory’s DataSource as
soon as it was being post-processed. When the native
EntityManagerFactory is being bootstrapped asynchronously, this
retrieval would block until bootstrapping had completed. This negated
some of the benefits of asynchronous bootstrapping.

This commit updates DataSourceInitializedPublisher so that it only
accesses the EntityManagerFactory’s DataSource once its bootstrapping
has completed. This is achieved using a decorated JpaVendorAdapter
that is called one the boostrapping has completed.

Closes gh-14061

Author: wilkinsona

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/orm/jpa/DataSourceInitializedPublisher.java

@@ -19,7 +19,10 @@
 import java.util.Map;
 import java.util.function.Supplier;
 
+import javax.persistence.EntityManager;
 import javax.persistence.EntityManagerFactory;
+import javax.persistence.spi.PersistenceProvider;
+import javax.persistence.spi.PersistenceUnitInfo;
 import javax.sql.DataSource;
 
 import org.springframework.beans.BeansException;
@@ -32,7 +35,11 @@
 import org.springframework.boot.jdbc.EmbeddedDatabaseConnection;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
+import org.springframework.core.task.AsyncTaskExecutor;
 import org.springframework.core.type.AnnotationMetadata;
+import org.springframework.orm.jpa.JpaDialect;
+import org.springframework.orm.jpa.JpaVendorAdapter;
+import org.springframework.orm.jpa.LocalContainerEntityManagerFactoryBean;
 
 /**
  * {@link BeanPostProcessor} used to fire {@link DataSourceSchemaCreatedEvent}s. Should
@@ -55,6 +62,10 @@ class DataSourceInitializedPublisher implements BeanPostProcessor {
 	@Override
 	public Object postProcessBeforeInitialization(Object bean, String beanName)
 			throws BeansException {
+		if (bean instanceof LocalContainerEntityManagerFactoryBean) {
+			LocalContainerEntityManagerFactoryBean factory = (LocalContainerEntityManagerFactoryBean) bean;
+			factory.setJpaVendorAdapter(new DataSourceSchemeCreatedPublisher(factory));
+		}
 		return bean;
 	}
 
@@ -71,9 +82,6 @@ public Object postProcessAfterInitialization(Object bean, String beanName)
 		if (bean instanceof HibernateProperties) {
 			this.hibernateProperties = (HibernateProperties) bean;
 		}
-		if (bean instanceof EntityManagerFactory) {
-			publishEventIfRequired((EntityManagerFactory) bean);
-		}
 		return bean;
 	}
 
@@ -88,8 +96,8 @@ private void publishEventIfRequired(EntityManagerFactory entityManagerFactory) {
 	private DataSource findDataSource(EntityManagerFactory entityManagerFactory) {
 		Object dataSource = entityManagerFactory.getProperties()
 				.get("javax.persistence.nonJtaDataSource");
-		return (dataSource != null && dataSource instanceof DataSource)
-				? (DataSource) dataSource : this.dataSource;
+		return (dataSource instanceof DataSource) ? (DataSource) dataSource
+				: this.dataSource;
 	}
 
 	private boolean isInitializingDatabase(DataSource dataSource) {
@@ -132,4 +140,66 @@ public void registerBeanDefinitions(AnnotationMetadata importingClassMetadata,
 
 	}
 
+	final class DataSourceSchemeCreatedPublisher implements JpaVendorAdapter {
+
+		private final JpaVendorAdapter delegate;
+
+		private final LocalContainerEntityManagerFactoryBean factory;
+
+		private DataSourceSchemeCreatedPublisher(
+				LocalContainerEntityManagerFactoryBean factory) {
+			this.delegate = factory.getJpaVendorAdapter();
+			this.factory = factory;
+		}
+
+		@Override
+		public PersistenceProvider getPersistenceProvider() {
+			return this.delegate.getPersistenceProvider();
+		}
+
+		@Override
+		public String getPersistenceProviderRootPackage() {
+			return this.delegate.getPersistenceProviderRootPackage();
+		}
+
+		@Override
+		public Map<String, ?> getJpaPropertyMap(PersistenceUnitInfo pui) {
+			return this.delegate.getJpaPropertyMap(pui);
+		}
+
+		@Override
+		public Map<String, ?> getJpaPropertyMap() {
+			return this.delegate.getJpaPropertyMap();
+		}
+
+		@Override
+		public JpaDialect getJpaDialect() {
+			return this.delegate.getJpaDialect();
+		}
+
+		@Override
+		public Class<? extends EntityManagerFactory> getEntityManagerFactoryInterface() {
+			return this.delegate.getEntityManagerFactoryInterface();
+		}
+
+		@Override
+		public Class<? extends EntityManager> getEntityManagerInterface() {
+			return this.delegate.getEntityManagerInterface();
+		}
+
+		@Override
+		public void postProcessEntityManagerFactory(EntityManagerFactory emf) {
+			this.delegate.postProcessEntityManagerFactory(emf);
+			AsyncTaskExecutor bootstrapExecutor = this.factory.getBootstrapExecutor();
+			if (bootstrapExecutor != null) {
+				bootstrapExecutor.execute(() -> DataSourceInitializedPublisher.this
+						.publishEventIfRequired(emf));
+			}
+			else {
+				DataSourceInitializedPublisher.this.publishEventIfRequired(emf);
+			}
+		}
+
+	}
+
 }