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 | 76 ++++++++++++-------
 1 file changed, 49 insertions(+), 27 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 7033b961a2..edafed5c7d 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,38 +75,34 @@ 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 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);
-	}
-
-	private void initializeContainers(Set<String> beanNames) {
-		List<Object> beans = new ArrayList<>(beanNames.size());
-		for (String beanName : beanNames) {
-			try {
-				beans.add(this.beanFactory.getBean(beanName));
-			}
-			catch (BeanCreationException ex) {
-				if (ex.contains(BeanCurrentlyInCreationException.class)) {
-					return;
-				}
-				throw ex;
-			}
+	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;
 		}
-		if (!this.containersInitialized) {
-			this.containersInitialized = true;
-			if (!beanNames.isEmpty()) {
-				logger.debug(LogMessage.format("Initialized container beans '%s'", beanNames));
-			}
-			start(beans);
+		beanNames.add(startableBeanName);
+		beans.add(startableBean);
+		start(beans);
+		if (!beanNames.isEmpty()) {
+			logger.debug(LogMessage.format("Initialized and started startable beans '%s'", beanNames));
 		}
 	}
 
@@ -115,6 +114,29 @@ class TestcontainersLifecycleBeanPostProcessor implements DestructionAwareBeanPo
 		this.startup.start(startables);
 	}
 
+	private void initializeContainers() {
+		List<String> beanNames = List.of(this.beanFactory.getBeanNamesForType(ContainerState.class, false, false));
+		if (getBeans(beanNames) == null) {
+			this.containersInitialized.set(false);
+		}
+	}
+
+	private List<Object> getBeans(List<String> beanNames) {
+		List<Object> beans = new ArrayList<>(beanNames.size());
+		for (String beanName : beanNames) {
+			try {
+				beans.add(this.beanFactory.getBean(beanName));
+			}
+			catch (BeanCreationException ex) {
+				if (ex.contains(BeanCurrentlyInCreationException.class)) {
+					return null;
+				}
+				throw ex;
+			}
+		}
+		return beans;
+	}
+
 	@Override
 	public boolean requiresDestruction(Object bean) {
 		return bean instanceof Startable;