diff --git a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/orm/jpa/DataSourceInitializedPublisher.java b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/orm/jpa/DataSourceInitializedPublisher.java index 50d5e5aefa..415482cada 100644 --- a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/orm/jpa/DataSourceInitializedPublisher.java +++ b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/orm/jpa/DataSourceInitializedPublisher.java @@ -17,6 +17,7 @@ package org.springframework.boot.autoconfigure.orm.jpa; import java.util.Map; +import java.util.concurrent.Future; import java.util.function.Supplier; import javax.persistence.EntityManager; @@ -35,7 +36,11 @@ import org.springframework.beans.factory.support.BeanDefinitionRegistry; import org.springframework.boot.autoconfigure.jdbc.DataSourceSchemaCreatedEvent; import org.springframework.boot.jdbc.EmbeddedDatabaseConnection; import org.springframework.context.ApplicationContext; +import org.springframework.context.ApplicationListener; import org.springframework.context.annotation.ImportBeanDefinitionRegistrar; +import org.springframework.context.event.ContextRefreshedEvent; +import org.springframework.core.Ordered; +import org.springframework.core.task.AsyncTaskExecutor; import org.springframework.core.type.AnnotationMetadata; import org.springframework.orm.jpa.JpaDialect; import org.springframework.orm.jpa.JpaVendorAdapter; @@ -60,6 +65,12 @@ class DataSourceInitializedPublisher implements BeanPostProcessor { private DataSourceSchemaCreatedPublisher schemaCreatedPublisher; + private DataSourceInitializationCompletionListener initializationCompletionListener; + + DataSourceInitializedPublisher(DataSourceInitializationCompletionListener completionListener) { + this.initializationCompletionListener = completionListener; + } + @Override public Object postProcessBeforeInitialization(Object bean, String beanName) throws BeansException { if (bean instanceof LocalContainerEntityManagerFactoryBean) { @@ -120,6 +131,35 @@ class DataSourceInitializedPublisher implements BeanPostProcessor { return hibernate.containsKey("hibernate.hbm2ddl.auto"); } + /** + * {@link ApplicationListener} that, upon receiving {@link ContextRefreshedEvent}, + * blocks until any asynchronous DataSource initialization has completed. + */ + static class DataSourceInitializationCompletionListener + implements ApplicationListener, Ordered { + + private volatile Future dataSourceInitialization; + + @Override + public void onApplicationEvent(ContextRefreshedEvent event) { + Future dataSourceInitialization = this.dataSourceInitialization; + if (dataSourceInitialization != null) { + try { + dataSourceInitialization.get(); + } + catch (Exception ex) { + throw new RuntimeException(ex); + } + } + } + + @Override + public int getOrder() { + return Ordered.HIGHEST_PRECEDENCE; + } + + } + /** * {@link ImportBeanDefinitionRegistrar} to register the * {@link DataSourceInitializedPublisher} without causing early bean instantiation @@ -127,19 +167,32 @@ class DataSourceInitializedPublisher implements BeanPostProcessor { */ static class Registrar implements ImportBeanDefinitionRegistrar { - private static final String BEAN_NAME = "dataSourceInitializedPublisher"; + private static final String PUBLISHER_BEAN_NAME = "dataSourceInitializedPublisher"; + + private static final String COMPLETION_LISTENER_BEAN_BEAN = DataSourceInitializationCompletionListener.class + .getName(); @Override public void registerBeanDefinitions(AnnotationMetadata importingClassMetadata, BeanDefinitionRegistry registry) { - if (!registry.containsBeanDefinition(BEAN_NAME)) { - AbstractBeanDefinition beanDefinition = BeanDefinitionBuilder.genericBeanDefinition( - DataSourceInitializedPublisher.class, DataSourceInitializedPublisher::new).getBeanDefinition(); - beanDefinition.setRole(BeanDefinition.ROLE_INFRASTRUCTURE); + if (!registry.containsBeanDefinition(PUBLISHER_BEAN_NAME)) { + DataSourceInitializationCompletionListener completionListener = new DataSourceInitializationCompletionListener(); + DataSourceInitializedPublisher publisher = new DataSourceInitializedPublisher(completionListener); + AbstractBeanDefinition publisherDefinition = BeanDefinitionBuilder + .genericBeanDefinition(DataSourceInitializedPublisher.class, () -> publisher) + .getBeanDefinition(); + publisherDefinition.setRole(BeanDefinition.ROLE_INFRASTRUCTURE); // We don't need this one to be post processed otherwise it can cause a // cascade of bean instantiation that we would rather avoid. - beanDefinition.setSynthetic(true); - registry.registerBeanDefinition(BEAN_NAME, beanDefinition); + publisherDefinition.setSynthetic(true); + registry.registerBeanDefinition(PUBLISHER_BEAN_NAME, publisherDefinition); + AbstractBeanDefinition listenerDefinition = BeanDefinitionBuilder.genericBeanDefinition( + DataSourceInitializationCompletionListener.class, () -> completionListener).getBeanDefinition(); + listenerDefinition.setRole(BeanDefinition.ROLE_INFRASTRUCTURE); + // We don't need this one to be post processed otherwise it can cause a + // cascade of bean instantiation that we would rather avoid. + listenerDefinition.setSynthetic(true); + registry.registerBeanDefinition(COMPLETION_LISTENER_BEAN_BEAN, listenerDefinition); } } @@ -194,7 +247,12 @@ class DataSourceInitializedPublisher implements BeanPostProcessor { @Override public void postProcessEntityManagerFactory(EntityManagerFactory entityManagerFactory) { this.delegate.postProcessEntityManagerFactory(entityManagerFactory); - publishEventIfRequired(this.factoryBean, entityManagerFactory); + AsyncTaskExecutor bootstrapExecutor = this.factoryBean.getBootstrapExecutor(); + if (bootstrapExecutor != null) { + DataSourceInitializedPublisher.this.initializationCompletionListener.dataSourceInitialization = bootstrapExecutor + .submit(() -> DataSourceInitializedPublisher.this.publishEventIfRequired(this.factoryBean, + entityManagerFactory)); + } } }