Pause unused application contexts in the TestContext framework

Since the introduction of the Spring TestContext Framework in 2007,
application contexts have always been stored in the context cache in a
"running" state. However, leaving a context running means that
components in the context may continue to run in the background. For
example, JMS listeners may continue to consume messages from a queue;
scheduled tasks may continue to perform active work, etc.; and this can
lead to issues within a test suite.

To address such issues, this commit introduces built-in support for
pausing application contexts when they are not in use and restarting
them if they are needed again.

Specifically, the TestContextManager now marks a test's application
context as "unused" after execution of the test class has ended, and
the underlying ContextCache then "stops" the application context if no
other test class is currently using the context. When a
TestExecutionListener later attempts to obtain a paused application
context -- for example, for a subsequent test class that shares the
same application context -- the ContextCache ensures that context is
restarted before returning it.

See https://github.com/spring-projects/spring-boot/issues/28312
See gh-35171
Closes gh-35168

spring-test/src/main/java/org/springframework/test/context/CacheAwareContextLoaderDelegate.java

@@ -29,6 +29,11 @@ import org.springframework.test.annotation.DirtiesContext.HierarchyMode;
  * {@link org.springframework.test.context.cache.ContextCache ContextCache}
  * behind the scenes.
  *
+ * <p>As of Spring Framework 7.0, this SPI includes optional support for
+ * {@linkplain #registerContextUsage(MergedContextConfiguration, Class) registering} and
+ * {@linkplain #unregisterContextUsage(MergedContextConfiguration, Class) unregistering}
+ * context usage.
+ *
  * <p>Note: {@code CacheAwareContextLoaderDelegate} does not extend the
  * {@link ContextLoader} or {@link SmartContextLoader} interface.
  *
@@ -142,4 +147,38 @@ public interface CacheAwareContextLoaderDelegate {
 	 */
 	void closeContext(MergedContextConfiguration mergedConfig, @Nullable HierarchyMode hierarchyMode);
 
+	/**
+	 * Register usage of the {@linkplain ApplicationContext application context}
+	 * for the supplied {@link MergedContextConfiguration} as well as usage of the
+	 * application context for its {@linkplain MergedContextConfiguration#getParent()
+	 * parent}, recursively.
+	 * <p>This is intended to be invoked whenever a
+	 * {@link org.springframework.test.context.TestExecutionListener TestExecutionListener}
+	 * interacts with the application context(s) on behalf of the supplied test class.
+	 * @param key the context key; never {@code null}
+	 * @param testClass the test class that is using the application context(s)
+	 * @since 7.0
+	 * @see #unregisterContextUsage(MergedContextConfiguration, Class)
+	 */
+	default void registerContextUsage(MergedContextConfiguration key, Class<?> testClass) {
+		/* no-op */
+	}
+
+	/**
+	 * Unregister usage of the {@linkplain ApplicationContext application context}
+	 * for the supplied {@link MergedContextConfiguration} as well as usage of the
+	 * application context for its {@linkplain MergedContextConfiguration#getParent()
+	 * parent}, recursively.
+	 * <p>This informs the {@code ContextCache} that the application context(s) can
+	 * be safely {@linkplain org.springframework.context.Lifecycle#stop() stopped}
+	 * if no other test classes are actively using the same application context(s).
+	 * @param key the context key; never {@code null}
+	 * @param testClass the test class that is no longer using the application context(s)
+	 * @since 7.0
+	 * @see #registerContextUsage(MergedContextConfiguration, Class)
+	 */
+	default void unregisterContextUsage(MergedContextConfiguration key, Class<?> testClass) {
+		/* no-op */
+	}
+
 }

spring-test/src/main/java/org/springframework/test/context/TestContext.java

@@ -53,7 +53,8 @@ public interface TestContext extends AttributeAccessor, Serializable {
 	 * Determine if the {@linkplain ApplicationContext application context} for
 	 * this test context is known to be available.
 	 * <p>If this method returns {@code true}, a subsequent invocation of
-	 * {@link #getApplicationContext()} should succeed.
+	 * {@link #getApplicationContext()} or {@link #markApplicationContextUnused()}
+	 * should succeed.
 	 * <p>The default implementation of this method always returns {@code false}.
 	 * Custom {@code TestContext} implementations are therefore highly encouraged
 	 * to override this method with a more meaningful implementation. Note that
@@ -62,6 +63,7 @@ public interface TestContext extends AttributeAccessor, Serializable {
 	 * @return {@code true} if the application context has already been loaded
 	 * @since 5.2
 	 * @see #getApplicationContext()
+	 * @see #markApplicationContextUnused()
 	 */
 	default boolean hasApplicationContext() {
 		return false;
@@ -77,6 +79,7 @@ public interface TestContext extends AttributeAccessor, Serializable {
 	 * @throws IllegalStateException if an error occurs while retrieving the
 	 * application context
 	 * @see #hasApplicationContext()
+	 * @see #markApplicationContextUnused()
 	 */
 	ApplicationContext getApplicationContext();
 
@@ -128,6 +131,24 @@ public interface TestContext extends AttributeAccessor, Serializable {
 	 */
 	@Nullable Throwable getTestException();
 
+	/**
+	 * Call this method to signal that the {@linkplain #getTestClass() test class}
+	 * is no longer using the {@linkplain ApplicationContext application context}
+	 * associated with this test context.
+	 * <p>This informs the context cache that the application context can be
+	 * safely {@linkplain org.springframework.context.Lifecycle#stop() stopped}
+	 * if no other test classes are actively using the same application context.
+	 * <p>This method is intended to be invoked after execution of the test class
+	 * has ended and should not be invoked unless the application context for this
+	 * test context is known to be {@linkplain #hasApplicationContext() available}.
+	 * <p>This feature is primarily intended for use within the framework.
+	 * @since 7.0
+	 * @see TestContextManager#afterTestClass()
+	 */
+	default void markApplicationContextUnused() {
+		/* no-op */
+	}
+
 	/**
 	 * Call this method to signal that the {@linkplain ApplicationContext application
 	 * context} associated with this test context is <em>dirty</em> and should be

spring-test/src/main/java/org/springframework/test/context/TestContextManager.java

@@ -520,10 +520,14 @@ public class TestContextManager {
 	 * the first exception.
 	 * <p>Note that listeners will be executed in the opposite order in which they
 	 * were registered.
+	 * <p>As of Spring Framework 7.0, this method also ensures that the application
+	 * context for the current {@link #getTestContext() TestContext} is marked as
+	 * {@linkplain TestContext#markApplicationContextUnused() unused}.
 	 * @throws Exception if a registered TestExecutionListener throws an exception
 	 * @since 3.0
 	 * @see #getTestExecutionListeners()
 	 * @see Throwable#addSuppressed(Throwable)
+	 * @see TestContext#markApplicationContextUnused()
 	 */
 	public void afterTestClass() throws Exception {
 		Class<?> testClass = getTestContext().getTestClass();
@@ -550,6 +554,20 @@ public class TestContextManager {
 			}
 		}
 
+		try {
+			if (getTestContext().hasApplicationContext()) {
+				getTestContext().markApplicationContextUnused();
+			}
+		}
+		catch (Throwable ex) {
+			if (afterTestClassException == null) {
+				afterTestClassException = ex;
+			}
+			else {
+				afterTestClassException.addSuppressed(ex);
+			}
+		}
+
 		this.testContextHolder.remove();
 
 		if (afterTestClassException != null) {

spring-test/src/main/java/org/springframework/test/context/cache/ContextCache.java

@@ -35,7 +35,10 @@ import org.springframework.test.context.MergedContextConfiguration;
  * <p>As of Spring Framework 6.1, this SPI includes optional support for
  * {@linkplain #getFailureCount(MergedContextConfiguration) tracking} and
  * {@linkplain #incrementFailureCount(MergedContextConfiguration) incrementing}
- * failure counts.
+ * failure counts. As of Spring Framework 7.0, this SPI includes optional support for
+ * {@linkplain #registerContextUsage(MergedContextConfiguration, Class) registering} and
+ * {@linkplain #unregisterContextUsage(MergedContextConfiguration, Class) unregistering}
+ * context usage.
  *
  * <h3>Rationale</h3>
  * <p>Context caching can have significant performance benefits if context
@@ -88,13 +91,19 @@ public interface ContextCache {
 	boolean contains(MergedContextConfiguration key);
 
 	/**
-	 * Obtain a cached {@code ApplicationContext} for the given key.
-	 * <p>The {@linkplain #getHitCount() hit} and {@linkplain #getMissCount() miss}
-	 * counts must be updated accordingly.
+	 * Obtain a cached {@link ApplicationContext} for the given key.
+	 * <p>If the cached application context was previously
+	 * {@linkplain org.springframework.context.Lifecycle#stop() stopped}, it
+	 * must be
+	 * {@linkplain org.springframework.context.support.AbstractApplicationContext#restart()
+	 * restarted}. This applies to parent contexts as well.
+	 * <p>In addition, the {@linkplain #getHitCount() hit} and
+	 * {@linkplain #getMissCount() miss} counts must be updated accordingly.
 	 * @param key the context key (never {@code null})
 	 * @return the corresponding {@code ApplicationContext} instance, or {@code null}
 	 * if not found in the cache
-	 * @see #remove
+	 * @see #unregisterContextUsage(MergedContextConfiguration, Class)
+	 * @see #remove(MergedContextConfiguration, HierarchyMode)
 	 */
 	@Nullable ApplicationContext get(MergedContextConfiguration key);
 
@@ -151,6 +160,64 @@ public interface ContextCache {
 	 * @see #getFailureCount(MergedContextConfiguration)
 	 */
 	default void incrementFailureCount(MergedContextConfiguration key) {
+		/* no-op */
+	}
+
+	/**
+	 * Register usage of the {@link ApplicationContext} for the supplied
+	 * {@link MergedContextConfiguration} and any of its parents.
+	 * <p>The default implementation of this method does nothing. Concrete
+	 * implementations are therefore highly encouraged to override this
+	 * method, {@link #unregisterContextUsage(MergedContextConfiguration, Class)},
+	 * and {@link #getContextUsageCount()} with appropriate behavior. Note that
+	 * the standard {@code ContextContext} implementation in Spring overrides
+	 * these methods appropriately.
+	 * @param key the context key; never {@code null}
+	 * @param testClass the test class that is using the application context(s)
+	 * @since 7.0
+	 * @see #unregisterContextUsage(MergedContextConfiguration, Class)
+	 * @see #getContextUsageCount()
+	 */
+	default void registerContextUsage(MergedContextConfiguration key, Class<?> testClass) {
+		/* no-op */
+	}
+
+	/**
+	 * Unregister usage of the {@link ApplicationContext} for the supplied
+	 * {@link MergedContextConfiguration} and any of its parents.
+	 * <p>If no other test classes are actively using the same application
+	 * context(s), the application context(s) should be
+	 * {@linkplain org.springframework.context.Lifecycle#stop() stopped}.
+	 * <p>The default implementation of this method does nothing. Concrete
+	 * implementations are therefore highly encouraged to override this
+	 * method, {@link #registerContextUsage(MergedContextConfiguration, Class)},
+	 * and {@link #getContextUsageCount()} with appropriate behavior. Note that
+	 * the standard {@code ContextContext} implementation in Spring overrides
+	 * these methods appropriately.
+	 * @param key the context key; never {@code null}
+	 * @param testClass the test class that is no longer using the application context(s)
+	 * @since 7.0
+	 * @see #registerContextUsage(MergedContextConfiguration, Class)
+	 * @see #getContextUsageCount()
+	 */
+	default void unregisterContextUsage(MergedContextConfiguration key, Class<?> testClass) {
+		/* no-op */
+	}
+
+	/**
+	 * Determine the number of contexts within the cache that are currently in use.
+	 * <p>The default implementation of this method always returns {@code 0}.
+	 * Concrete implementations are therefore highly encouraged to override this
+	 * method, {@link #registerContextUsage(MergedContextConfiguration, Class)},
+	 * and {@link #unregisterContextUsage(MergedContextConfiguration, Class)} with
+	 * appropriate behavior. Note that the standard {@code ContextContext}
+	 * implementation in Spring overrides these methods appropriately.
+	 * @since 7.0
+	 * @see #registerContextUsage(MergedContextConfiguration, Class)
+	 * @see #unregisterContextUsage(MergedContextConfiguration, Class)
+	 */
+	default int getContextUsageCount() {
+		return 0;
 	}
 
 	/**
@@ -204,6 +271,7 @@ public interface ContextCache {
 	 * <ul>
 	 * <li>name of the concrete {@code ContextCache} implementation</li>
 	 * <li>{@linkplain #size}</li>
+	 * <li>{@linkplain #getContextUsageCount() context usage count}</li>
 	 * <li>{@linkplain #getParentContextCount() parent context count}</li>
 	 * <li>{@linkplain #getHitCount() hit count}</li>
 	 * <li>{@linkplain #getMissCount() miss count}</li>

spring-test/src/main/java/org/springframework/test/context/cache/DefaultCacheAwareContextLoaderDelegate.java

@@ -58,6 +58,11 @@ import org.springframework.util.Assert;
  * delegating to {@link ContextCacheUtils#retrieveContextFailureThreshold()} to
  * obtain the threshold value to use.
  *
+ * <p>As of Spring Framework 7.0, this class provides support for
+ * {@linkplain #registerContextUsage(MergedContextConfiguration, Class) registering} and
+ * {@linkplain #unregisterContextUsage(MergedContextConfiguration, Class) unregistering}
+ * context usage.
+ *
  * @author Sam Brannen
  * @since 4.1
  */
@@ -204,6 +209,22 @@ public class DefaultCacheAwareContextLoaderDelegate implements CacheAwareContext
 		}
 	}
 
+	@Override
+	public void registerContextUsage(MergedContextConfiguration mergedConfig, Class<?> testClass) {
+		mergedConfig = replaceIfNecessary(mergedConfig);
+		synchronized (this.contextCache) {
+			this.contextCache.registerContextUsage(mergedConfig, testClass);
+		}
+	}
+
+	@Override
+	public void unregisterContextUsage(MergedContextConfiguration mergedConfig, Class<?> testClass) {
+		mergedConfig = replaceIfNecessary(mergedConfig);
+		synchronized (this.contextCache) {
+			this.contextCache.unregisterContextUsage(mergedConfig, testClass);
+		}
+	}
+
 	/**
 	 * Get the {@link ContextCache} used by this context loader delegate.
 	 */

spring-test/src/main/java/org/springframework/test/context/cache/DefaultContextCache.java

@@ -73,6 +73,13 @@ public class DefaultContextCache implements ContextCache {
 	private final Map<MergedContextConfiguration, Set<MergedContextConfiguration>> hierarchyMap =
 			new ConcurrentHashMap<>(32);
 
+	/**
+	 * Map of context keys to active test classes (i.e., test classes that are actively
+	 * using the corresponding {@link ApplicationContext}).
+	 * @since 7.0
+	 */
+	private final Map<MergedContextConfiguration, Set<Class<?>>> contextUsageMap = new ConcurrentHashMap<>(32);
+
 	/**
 	 * Map of context keys to context load failure counts.
 	 * @since 6.1
@@ -129,10 +136,22 @@ public class DefaultContextCache implements ContextCache {
 		}
 		else {
 			this.hitCount.incrementAndGet();
+			restartContextIfNecessary(context);
 		}
 		return context;
 	}
 
+	private void restartContextIfNecessary(ApplicationContext context) {
+		// Recurse up the context hierarchy first.
+		ApplicationContext parent = context.getParent();
+		if (parent != null) {
+			restartContextIfNecessary(parent);
+		}
+		if (context instanceof ConfigurableApplicationContext cac && !cac.isRunning()) {
+			cac.restart();
+		}
+	}
+
 	@Override
 	public void put(MergedContextConfiguration key, ApplicationContext context) {
 		Assert.notNull(key, "Key must not be null");
@@ -149,6 +168,41 @@ public class DefaultContextCache implements ContextCache {
 		}
 	}
 
+	@Override
+	public void registerContextUsage(MergedContextConfiguration mergedConfig, Class<?> testClass) {
+		// Recurse up the context hierarchy first.
+		MergedContextConfiguration parent = mergedConfig.getParent();
+		if (parent != null) {
+			registerContextUsage(parent, testClass);
+		}
+		getActiveTestClasses(mergedConfig).add(testClass);
+	}
+
+	@Override
+	public void unregisterContextUsage(MergedContextConfiguration mergedConfig, Class<?> testClass) {
+		ApplicationContext context = this.contextMap.get(mergedConfig);
+		Assert.state(context != null, "ApplicationContext must not be null for: " + mergedConfig);
+
+		Set<Class<?>> activeTestClasses = getActiveTestClasses(mergedConfig);
+		activeTestClasses.remove(testClass);
+		if (activeTestClasses.isEmpty()) {
+			if (context instanceof ConfigurableApplicationContext cac && cac.isRunning()) {
+				cac.stop();
+			}
+			this.contextUsageMap.remove(mergedConfig);
+		}
+
+		// Recurse up the context hierarchy last.
+		MergedContextConfiguration parent = mergedConfig.getParent();
+		if (parent != null) {
+			unregisterContextUsage(parent, testClass);
+		}
+	}
+
+	private Set<Class<?>> getActiveTestClasses(MergedContextConfiguration mergedConfig) {
+		return this.contextUsageMap.computeIfAbsent(mergedConfig, mcc -> new HashSet<>());
+	}
+
 	@Override
 	public void remove(MergedContextConfiguration key, @Nullable HierarchyMode hierarchyMode) {
 		Assert.notNull(key, "Key must not be null");
@@ -199,6 +253,7 @@ public class DefaultContextCache implements ContextCache {
 		// Physically remove and close leaf nodes first (i.e., on the way back up the
 		// stack as opposed to prior to the recursive call).
 		ApplicationContext context = this.contextMap.remove(key);
+		this.contextUsageMap.remove(key);
 		if (context instanceof ConfigurableApplicationContext cac) {
 			cac.close();
 		}
@@ -228,6 +283,11 @@ public class DefaultContextCache implements ContextCache {
 		return this.maxSize;
 	}
 
+	@Override
+	public int getContextUsageCount() {
+		return this.contextUsageMap.size();
+	}
+
 	@Override
 	public int getParentContextCount() {
 		return this.hierarchyMap.size();
@@ -258,6 +318,7 @@ public class DefaultContextCache implements ContextCache {
 		synchronized (this.contextMap) {
 			this.contextMap.clear();
 			this.hierarchyMap.clear();
+			this.contextUsageMap.clear();
 		}
 	}
 
@@ -288,6 +349,7 @@ public class DefaultContextCache implements ContextCache {
 		return new ToStringCreator(this)
 				.append("size", size())
 				.append("maxSize", getMaxSize())
+				.append("contextUsageCount", getContextUsageCount())
 				.append("parentContextCount", getParentContextCount())
 				.append("hitCount", getHitCount())
 				.append("missCount", getMissCount())

spring-test/src/main/java/org/springframework/test/context/support/DefaultTestContext.java

@@ -136,9 +136,25 @@ public class DefaultTestContext implements TestContext {
 					from the ContextCache due to a maximum cache size policy."""
 						.formatted(this.mergedConfig));
 		}
+		this.cacheAwareContextLoaderDelegate.registerContextUsage(this.mergedConfig, this.testClass);
 		return context;
 	}
 
+	/**
+	 * Mark the {@linkplain ApplicationContext application context} associated
+	 * with this test context as <em>unused</em> so that it can be safely
+	 * {@linkplain org.springframework.context.Lifecycle#stop() stopped} if no
+	 * other test classes are actively using the same application context.
+	 * <p>The default implementation delegates to the {@link CacheAwareContextLoaderDelegate}
+	 * that was supplied when this {@code TestContext} was constructed.
+	 * @since 7.0
+	 * @see CacheAwareContextLoaderDelegate#unregisterContextUsage(MergedContextConfiguration, Class)
+	 */
+	@Override
+	public void markApplicationContextUnused() {
+		this.cacheAwareContextLoaderDelegate.unregisterContextUsage(this.mergedConfig, this.testClass);
+	}
+
 	/**
 	 * Mark the {@linkplain ApplicationContext application context} associated
 	 * with this test context as <em>dirty</em> (i.e., by removing it from the