Shutdown in-memory R2DBC databases before devtools restart

Add `DevToolsR2dbcAutoConfiguration` to automatically shutdown in-memory
R2DBC databases before restarting. Prior to this commit, restarts that
involved SQL initialization scripts could fail due to dirty database
content.

The `DevToolsR2dbcAutoConfiguration` class is similar in design to
`DevToolsDataSourceAutoConfiguration`, but it applies to both pooled
and non-pooled connection factories. The `DataSource` variant does not
need to deal with non-pooled connections due to the fact that
`EmbeddedDataSourceConfiguration` calls `EmbeddedDatabase.shutdown`
as a `destroyMethod`. With R2DB we don't have an `EmbeddedDatabase`
equivalent so we can always trigger a shutdown for devtools.

Fixes gh-28345

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

@@ -32,6 +32,8 @@ dependencies {
 
 	intTestRuntimeOnly("org.springframework:spring-web")
 
+	optional("io.projectreactor:reactor-core")
+	optional("io.r2dbc:r2dbc-spi")
 	optional("javax.servlet:javax.servlet-api")
 	optional("org.apache.derby:derby")
 	optional("org.hibernate:hibernate-core")
@@ -72,6 +74,7 @@ dependencies {
 
 	testRuntimeOnly("org.aspectj:aspectjweaver")
 	testRuntimeOnly("org.yaml:snakeyaml")
+	testRuntimeOnly("io.r2dbc:r2dbc-h2")
 }
 
 task syncIntTestDependencies(type: Sync) {

spring-boot-project/spring-boot-devtools/src/main/java/org/springframework/boot/devtools/autoconfigure/DevToolsR2dbcAutoConfiguration.java

@@ -0,0 +1,150 @@
+/*
+ * Copyright 2012-2021 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.devtools.autoconfigure;
+
+import io.r2dbc.spi.Connection;
+import io.r2dbc.spi.ConnectionFactory;
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Mono;
+
+import org.springframework.beans.factory.DisposableBean;
+import org.springframework.beans.factory.annotation.AnnotatedBeanDefinition;
+import org.springframework.beans.factory.config.BeanDefinition;
+import org.springframework.boot.autoconfigure.AutoConfigureAfter;
+import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
+import org.springframework.boot.autoconfigure.condition.ConditionMessage;
+import org.springframework.boot.autoconfigure.condition.ConditionOutcome;
+import org.springframework.boot.autoconfigure.condition.SpringBootCondition;
+import org.springframework.boot.autoconfigure.r2dbc.R2dbcAutoConfiguration;
+import org.springframework.boot.devtools.autoconfigure.DevToolsR2dbcAutoConfiguration.DevToolsConnectionFactoryCondition;
+import org.springframework.boot.r2dbc.EmbeddedDatabaseConnection;
+import org.springframework.context.ApplicationEventPublisher;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.ConditionContext;
+import org.springframework.context.annotation.Conditional;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.context.annotation.ConfigurationCondition;
+import org.springframework.core.type.AnnotatedTypeMetadata;
+import org.springframework.core.type.MethodMetadata;
+
+/**
+ * {@link EnableAutoConfiguration Auto-configuration} for DevTools-specific R2DBC
+ * configuration.
+ *
+ * @author Phillip Webb
+ * @since 2.5.6
+ */
+@AutoConfigureAfter(R2dbcAutoConfiguration.class)
+@Conditional({ OnEnabledDevToolsCondition.class, DevToolsConnectionFactoryCondition.class })
+@Configuration(proxyBeanMethods = false)
+public class DevToolsR2dbcAutoConfiguration {
+
+	@Bean
+	InMemoryR2dbcDatabaseShutdownExecutor inMemoryR2dbcDatabaseShutdownExecutor(
+			ApplicationEventPublisher eventPublisher, ConnectionFactory connectionFactory) {
+		return new InMemoryR2dbcDatabaseShutdownExecutor(eventPublisher, connectionFactory);
+	}
+
+	final class InMemoryR2dbcDatabaseShutdownExecutor implements DisposableBean {
+
+		private final ApplicationEventPublisher eventPublisher;
+
+		private final ConnectionFactory connectionFactory;
+
+		InMemoryR2dbcDatabaseShutdownExecutor(ApplicationEventPublisher eventPublisher,
+				ConnectionFactory connectionFactory) {
+			this.eventPublisher = eventPublisher;
+			this.connectionFactory = connectionFactory;
+		}
+
+		@Override
+		public void destroy() throws Exception {
+			if (shouldShutdown()) {
+				Mono.usingWhen(this.connectionFactory.create(), this::executeShutdown, this::closeConnection,
+						this::closeConnection, this::closeConnection).block();
+				this.eventPublisher.publishEvent(new R2dbcDatabaseShutdownEvent(this.connectionFactory));
+			}
+		}
+
+		private boolean shouldShutdown() {
+			try {
+				return EmbeddedDatabaseConnection.isEmbedded(this.connectionFactory);
+			}
+			catch (Exception ex) {
+				return false;
+			}
+		}
+
+		private Mono<?> executeShutdown(Connection connection) {
+			return Mono.from(connection.createStatement("SHUTDOWN").execute());
+		}
+
+		private Publisher<Void> closeConnection(Connection connection) {
+			return closeConnection(connection, null);
+		}
+
+		private Publisher<Void> closeConnection(Connection connection, Throwable ex) {
+			return connection.close();
+		}
+
+	}
+
+	static class DevToolsConnectionFactoryCondition extends SpringBootCondition implements ConfigurationCondition {
+
+		@Override
+		public ConfigurationPhase getConfigurationPhase() {
+			return ConfigurationPhase.REGISTER_BEAN;
+		}
+
+		@Override
+		public ConditionOutcome getMatchOutcome(ConditionContext context, AnnotatedTypeMetadata metadata) {
+			ConditionMessage.Builder message = ConditionMessage.forCondition("DevTools ConnectionFactory Condition");
+			String[] beanNames = context.getBeanFactory().getBeanNamesForType(ConnectionFactory.class, true, false);
+			if (beanNames.length != 1) {
+				return ConditionOutcome.noMatch(message.didNotFind("a single ConnectionFactory bean").atAll());
+			}
+			BeanDefinition beanDefinition = context.getRegistry().getBeanDefinition(beanNames[0]);
+			if (beanDefinition instanceof AnnotatedBeanDefinition
+					&& isAutoConfigured((AnnotatedBeanDefinition) beanDefinition)) {
+				return ConditionOutcome.match(message.foundExactly("auto-configured ConnectionFactory"));
+			}
+			return ConditionOutcome.noMatch(message.didNotFind("an auto-configured ConnectionFactory").atAll());
+		}
+
+		private boolean isAutoConfigured(AnnotatedBeanDefinition beanDefinition) {
+			MethodMetadata methodMetadata = beanDefinition.getFactoryMethodMetadata();
+			return methodMetadata != null && methodMetadata.getDeclaringClassName()
+					.startsWith(R2dbcAutoConfiguration.class.getPackage().getName());
+		}
+
+	}
+
+	static class R2dbcDatabaseShutdownEvent {
+
+		private final ConnectionFactory connectionFactory;
+
+		R2dbcDatabaseShutdownEvent(ConnectionFactory connectionFactory) {
+			this.connectionFactory = connectionFactory;
+		}
+
+		ConnectionFactory getConnectionFactory() {
+			return this.connectionFactory;
+		}
+
+	}
+
+}

spring-boot-project/spring-boot-devtools/src/main/resources/META-INF/spring.factories

@@ -10,6 +10,7 @@ org.springframework.boot.devtools.logger.DevToolsLogFactory.Listener
 # Auto Configure
 org.springframework.boot.autoconfigure.EnableAutoConfiguration=\
 org.springframework.boot.devtools.autoconfigure.DevToolsDataSourceAutoConfiguration,\
+org.springframework.boot.devtools.autoconfigure.DevToolsR2dbcAutoConfiguration,\
 org.springframework.boot.devtools.autoconfigure.LocalDevToolsAutoConfiguration,\
 org.springframework.boot.devtools.autoconfigure.RemoteDevToolsAutoConfiguration
 

spring-boot-project/spring-boot-devtools/src/test/java/org/springframework/boot/devtools/autoconfigure/DevToolsR2dbcAutoConfigurationTests.java

@@ -0,0 +1,203 @@
+/*
+ * Copyright 2012-2021 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.devtools.autoconfigure;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
+
+import io.r2dbc.spi.Connection;
+import io.r2dbc.spi.ConnectionFactory;
+import io.r2dbc.spi.ConnectionFactoryMetadata;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Nested;
+import org.junit.jupiter.api.Test;
+import org.reactivestreams.Publisher;
+
+import org.springframework.beans.factory.annotation.AnnotatedGenericBeanDefinition;
+import org.springframework.boot.autoconfigure.r2dbc.R2dbcAutoConfiguration;
+import org.springframework.boot.devtools.autoconfigure.DevToolsR2dbcAutoConfiguration.R2dbcDatabaseShutdownEvent;
+import org.springframework.boot.test.util.TestPropertyValues;
+import org.springframework.boot.testsupport.classpath.ClassPathExclusions;
+import org.springframework.context.ApplicationListener;
+import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.util.ObjectUtils;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Tests for {@link DevToolsR2dbcAutoConfiguration}.
+ *
+ * @author Phillip Webb
+ */
+class DevToolsR2dbcAutoConfigurationTests {
+
+	static List<ConnectionFactory> shutdowns = Collections.synchronizedList(new ArrayList<>());
+
+	abstract static class Common {
+
+		@BeforeEach
+		void reset() {
+			shutdowns.clear();
+		}
+
+		@Test
+		void autoConfiguredInMemoryConnectionFactoryIsShutdown() throws Exception {
+			ConfigurableApplicationContext context = getContext(() -> createContext());
+			ConnectionFactory connectionFactory = context.getBean(ConnectionFactory.class);
+			context.close();
+			assertThat(shutdowns).contains(connectionFactory);
+		}
+
+		@Test
+		void nonEmbeddedConnectionFactoryIsNotShutdown() throws Exception {
+			ConfigurableApplicationContext context = getContext(() -> createContext("r2dbc:h2:file:///testdb"));
+			ConnectionFactory connectionFactory = context.getBean(ConnectionFactory.class);
+			context.close();
+			assertThat(shutdowns).doesNotContain(connectionFactory);
+		}
+
+		@Test
+		void singleManuallyConfiguredConnectionFactoryIsNotClosed() throws Exception {
+			ConfigurableApplicationContext context = getContext(
+					() -> createContext(SingleConnectionFactoryConfiguration.class));
+			ConnectionFactory connectionFactory = context.getBean(ConnectionFactory.class);
+			context.close();
+			assertThat(shutdowns).doesNotContain(connectionFactory);
+		}
+
+		@Test
+		void multipleConnectionFactoriesAreIgnored() throws Exception {
+			ConfigurableApplicationContext context = getContext(
+					() -> createContext(MultipleConnectionFactoriesConfiguration.class));
+			Collection<ConnectionFactory> connectionFactory = context.getBeansOfType(ConnectionFactory.class).values();
+			context.close();
+			assertThat(shutdowns).doesNotContainAnyElementsOf(connectionFactory);
+		}
+
+		@Test
+		void emptyFactoryMethodMetadataIgnored() throws Exception {
+			ConfigurableApplicationContext context = getContext(this::getEmptyFactoryMethodMetadataIgnoredContext);
+			ConnectionFactory connectionFactory = context.getBean(ConnectionFactory.class);
+			context.close();
+			assertThat(shutdowns).doesNotContain(connectionFactory);
+		}
+
+		private ConfigurableApplicationContext getEmptyFactoryMethodMetadataIgnoredContext() {
+			AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
+			ConnectionFactory connectionFactory = new MockConnectionFactory();
+			AnnotatedGenericBeanDefinition beanDefinition = new AnnotatedGenericBeanDefinition(
+					connectionFactory.getClass());
+			context.registerBeanDefinition("connectionFactory", beanDefinition);
+			context.register(R2dbcAutoConfiguration.class, DevToolsR2dbcAutoConfiguration.class);
+			context.refresh();
+			return context;
+		}
+
+		protected ConfigurableApplicationContext getContext(Supplier<ConfigurableApplicationContext> supplier)
+				throws Exception {
+			AtomicReference<ConfigurableApplicationContext> atomicReference = new AtomicReference<>();
+			Thread thread = new Thread(() -> {
+				ConfigurableApplicationContext context = supplier.get();
+				atomicReference.getAndSet(context);
+			});
+			thread.start();
+			thread.join();
+			return atomicReference.get();
+		}
+
+		protected final ConfigurableApplicationContext createContext(Class<?>... classes) {
+			return createContext(null, classes);
+		}
+
+		protected final ConfigurableApplicationContext createContext(String url, Class<?>... classes) {
+			AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
+			if (!ObjectUtils.isEmpty(classes)) {
+				context.register(classes);
+			}
+			context.register(R2dbcAutoConfiguration.class, DevToolsR2dbcAutoConfiguration.class);
+			if (url != null) {
+				TestPropertyValues.of("spring.r2dbc.url:" + url).applyTo(context);
+			}
+			context.addApplicationListener(ApplicationListener.forPayload(this::onEvent));
+			context.refresh();
+			return context;
+		}
+
+		private void onEvent(R2dbcDatabaseShutdownEvent event) {
+			shutdowns.add(event.getConnectionFactory());
+		}
+
+	}
+
+	@Nested
+	@ClassPathExclusions("r2dbc-pool*.jar")
+	static class Embedded extends Common {
+
+	}
+
+	@Nested
+	static class Pooled extends Common {
+
+	}
+
+	@Configuration(proxyBeanMethods = false)
+	static class SingleConnectionFactoryConfiguration {
+
+		@Bean
+		ConnectionFactory connectionFactory() {
+			return new MockConnectionFactory();
+		}
+
+	}
+
+	@Configuration(proxyBeanMethods = false)
+	static class MultipleConnectionFactoriesConfiguration {
+
+		@Bean
+		ConnectionFactory connectionFactoryOne() {
+			return new MockConnectionFactory();
+		}
+
+		@Bean
+		ConnectionFactory connectionFactoryTwo() {
+			return new MockConnectionFactory();
+		}
+
+	}
+
+	private static class MockConnectionFactory implements ConnectionFactory {
+
+		@Override
+		public Publisher<? extends Connection> create() {
+			return null;
+		}
+
+		@Override
+		public ConnectionFactoryMetadata getMetadata() {
+			return null;
+		}
+
+	}
+
+}