From 0c66db7b1815cfea98e336bccaea85246ba4b4f0 Mon Sep 17 00:00:00 2001
From: Phillip Webb <pwebb@vmware.com>
Date: Wed, 25 Oct 2023 11:53:04 -0700
Subject: [PATCH] Refine container initialization and parallel startup logic

Update `TestcontainersLifecycleBeanPostProcessor` to restore early
container initialization logic and refine startup logic. Initial bean
access now again triggers the creation all container beans. In addition
the first access of a `Startable` bean now attempts to find and start
all other `Startable` beans.

Fixes gh-37989
---
 ...tcontainersLifecycleBeanPostProcessor.java | 68 ++++++++++++-------
 1 file changed, 45 insertions(+), 23 deletions(-)

diff --git a/spring-boot-project/spring-boot-testcontainers/src/main/java/org/springframework/boot/testcontainers/lifecycle/TestcontainersLifecycleBeanPostProcessor.java b/spring-boot-project/spring-boot-testcontainers/src/main/java/org/springframework/boot/testcontainers/lifecycle/TestcontainersLifecycleBeanPostProcessor.java
index 7033b961a29..edafed5c7d0 100644
--- a/spring-boot-project/spring-boot-testcontainers/src/main/java/org/springframework/boot/testcontainers/lifecycle/TestcontainersLifecycleBeanPostProcessor.java
+++ b/spring-boot-project/spring-boot-testcontainers/src/main/java/org/springframework/boot/testcontainers/lifecycle/TestcontainersLifecycleBeanPostProcessor.java
@@ -20,6 +20,7 @@ import java.util.ArrayList;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.stream.Collectors;
 
 import org.apache.commons.logging.Log;
@@ -62,7 +63,9 @@ class TestcontainersLifecycleBeanPostProcessor implements DestructionAwareBeanPo
 
 	private final TestcontainersStartup startup;
 
-	private volatile boolean containersInitialized = false;
+	private final AtomicBoolean startablesInitialized = new AtomicBoolean();
+
+	private final AtomicBoolean containersInitialized = new AtomicBoolean();
 
 	TestcontainersLifecycleBeanPostProcessor(ConfigurableListableBeanFactory beanFactory,
 			TestcontainersStartup startup) {
@@ -72,20 +75,53 @@ class TestcontainersLifecycleBeanPostProcessor implements DestructionAwareBeanPo
 
 	@Override
 	public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
-		if (!this.containersInitialized && this.beanFactory.isConfigurationFrozen()) {
+		if (this.beanFactory.isConfigurationFrozen() && this.containersInitialized.compareAndSet(false, true)) {
 			initializeContainers();
 		}
+		if (bean instanceof Startable startableBean) {
+			if (this.startablesInitialized.compareAndSet(false, true)) {
+				initializeStartables(startableBean, beanName);
+			}
+			else {
+				startableBean.start();
+			}
+		}
 		return bean;
 	}
 
+	private void initializeStartables(Startable startableBean, String startableBeanName) {
+		List<String> beanNames = new ArrayList<>(
+				List.of(this.beanFactory.getBeanNamesForType(Startable.class, false, false)));
+		beanNames.remove(startableBeanName);
+		List<Object> beans = getBeans(beanNames);
+		if (beans == null) {
+			this.startablesInitialized.set(false);
+			return;
+		}
+		beanNames.add(startableBeanName);
+		beans.add(startableBean);
+		start(beans);
+		if (!beanNames.isEmpty()) {
+			logger.debug(LogMessage.format("Initialized and started startable beans '%s'", beanNames));
+		}
+	}
+
+	private void start(List<Object> beans) {
+		Set<Startable> startables = beans.stream()
+			.filter(Startable.class::isInstance)
+			.map(Startable.class::cast)
+			.collect(Collectors.toCollection(LinkedHashSet::new));
+		this.startup.start(startables);
+	}
+
 	private void initializeContainers() {
-		Set<String> beanNames = new LinkedHashSet<>();
-		beanNames.addAll(List.of(this.beanFactory.getBeanNamesForType(ContainerState.class, false, false)));
-		beanNames.addAll(List.of(this.beanFactory.getBeanNamesForType(Startable.class, false, false)));
-		initializeContainers(beanNames);
+		List<String> beanNames = List.of(this.beanFactory.getBeanNamesForType(ContainerState.class, false, false));
+		if (getBeans(beanNames) == null) {
+			this.containersInitialized.set(false);
+		}
 	}
 
-	private void initializeContainers(Set<String> beanNames) {
+	private List<Object> getBeans(List<String> beanNames) {
 		List<Object> beans = new ArrayList<>(beanNames.size());
 		for (String beanName : beanNames) {
 			try {
@@ -93,26 +129,12 @@ class TestcontainersLifecycleBeanPostProcessor implements DestructionAwareBeanPo
 			}
 			catch (BeanCreationException ex) {
 				if (ex.contains(BeanCurrentlyInCreationException.class)) {
-					return;
+					return null;
 				}
 				throw ex;
 			}
 		}
-		if (!this.containersInitialized) {
-			this.containersInitialized = true;
-			if (!beanNames.isEmpty()) {
-				logger.debug(LogMessage.format("Initialized container beans '%s'", beanNames));
-			}
-			start(beans);
-		}
-	}
-
-	private void start(List<Object> beans) {
-		Set<Startable> startables = beans.stream()
-			.filter(Startable.class::isInstance)
-			.map(Startable.class::cast)
-			.collect(Collectors.toCollection(LinkedHashSet::new));
-		this.startup.start(startables);
+		return beans;
 	}
 
 	@Override