Evict context from ContextCache before loading a new context

Since Spring Framework 4.2, DefaultContextCache supported an LRU (least
recently used) eviction policy via a custom LruCache which extended
LinkedHashMap. The LruCache reacted to LinkedHashMap's
removeEldestEntry() callback to remove the LRU context if the maxSize
of the cache was exceeded.

Due to the nature of the implementation in LinkedHashMap, the
removeEldestEntry() callback is invoked after a new entry has been
stored to the map.

Consequently, a Spring ApplicationContext (C1) was evicted from the
cache after a new context (C2) was loaded and added to the cache,
leading to failure scenarios such as the following.

- C1 and C2 share an external resource -- for example, a database.
- C2 initializes the external resource with test data when C2 is loaded.
- C1 cleans up the external resource when C1 is closed.
- C1 is loaded and added to the cache.
- C2 is loaded and added to the cache before C1 is evicted.
- C1 is evicted and closed.
- C2 tests fail, because C1 removed test data required for C2.

To address such scenarios, this commit replaces the custom LruCache
with custom LRU eviction logic in DefaultContextCache and revises
the put(MergedContextConfiguration, ApplicationContext) method to
delegate to a new evictLruContextIfNecessary() method.

This commit also introduces a new put(MergedContextConfiguration,
LoadFunction) method in the ContextCache API which is overridden by
DefaultContextCache to ensure that an evicted context is removed and
closed before a new context is loaded to take its place in the cache.

In addition, DefaultCacheAwareContextLoaderDelegate has been revised to
make use of the new put(MergedContextConfiguration, LoadFunction) API.

Closes gh-21007

Author: sbrannen

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

@@ -21,6 +21,7 @@
 import org.springframework.context.ApplicationContext;
 import org.springframework.test.annotation.DirtiesContext.HierarchyMode;
 import org.springframework.test.context.MergedContextConfiguration;
+import org.springframework.util.Assert;
 
 /**
  * {@code ContextCache} defines the SPI for caching Spring
@@ -102,6 +103,7 @@ public interface ContextCache {
 	 * @param key the context key; never {@code null}
 	 * @return the corresponding {@code ApplicationContext} instance, or {@code null}
 	 * if not found in the cache
+	 * @see #put(MergedContextConfiguration, LoadFunction)
 	 * @see #unregisterContextUsage(MergedContextConfiguration, Class)
 	 * @see #remove(MergedContextConfiguration, HierarchyMode)
 	 */
@@ -113,9 +115,37 @@ public interface ContextCache {
 	 * @param key the context key; never {@code null}
 	 * @param context the {@code ApplicationContext}; never {@code null}
 	 * @see #get(MergedContextConfiguration)
+	 * @see #put(MergedContextConfiguration, LoadFunction)
 	 */
 	void put(MergedContextConfiguration key, ApplicationContext context);
 
+	/**
+	 * Explicitly add an {@link ApplicationContext} to the cache under the given
+	 * key, potentially honoring a custom eviction policy.
+	 * <p>The supplied {@link LoadFunction} will be invoked to load the
+	 * {@code ApplicationContext}.
+	 * <p>Concrete implementations which honor a custom eviction policy must
+	 * override this method to ensure that an evicted context is removed from the
+	 * cache and closed before a new context is loaded via the supplied
+	 * {@code LoadFunction}.
+	 * @param key the context key; never {@code null}
+	 * @param loadFunction a function which loads the context for the supplied key;
+	 * never {@code null}
+	 * @return the {@code ApplicationContext}; never {@code null}
+	 * @since 7.0
+	 * @see #get(MergedContextConfiguration)
+	 * @see #put(MergedContextConfiguration, ApplicationContext)
+	 */
+	default ApplicationContext put(MergedContextConfiguration key, LoadFunction loadFunction) {
+		Assert.notNull(key, "Key must not be null");
+		Assert.notNull(loadFunction, "LoadFunction must not be null");
+
+		ApplicationContext applicationContext = loadFunction.loadContext(key);
+		Assert.state(applicationContext != null, "LoadFunction must return a non-null ApplicationContext");
+		put(key, applicationContext);
+		return applicationContext;
+	}
+
 	/**
 	 * Remove the context with the given key from the cache and explicitly
 	 * {@linkplain org.springframework.context.ConfigurableApplicationContext#close() close}
@@ -281,4 +311,26 @@ default int getContextUsageCount() {
 	 */
 	void logStatistics();
 
+
+	/**
+	 * Represents a function that loads an {@link ApplicationContext}.
+	 *
+	 * @since 7.0
+	 */
+	@FunctionalInterface
+	interface LoadFunction {
+
+		/**
+		 * Load a new {@link ApplicationContext} based on the supplied
+		 * {@link MergedContextConfiguration} and return the context in a fully
+		 * <em>refreshed</em> state.
+		 * @param mergedConfig the merged context configuration to use to load the
+		 * application context
+		 * @return a new application context; never {@code null}
+		 * @see org.springframework.test.context.SmartContextLoader#loadContext(MergedContextConfiguration)
+		 */
+		ApplicationContext loadContext(MergedContextConfiguration mergedConfig);
+
+	}
+
 }

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

@@ -139,34 +139,44 @@ public boolean isContextLoaded(MergedContextConfiguration mergedConfig) {
 	public ApplicationContext loadContext(MergedContextConfiguration mergedConfig) {
 		mergedConfig = replaceIfNecessary(mergedConfig);
 		synchronized (this.contextCache) {
-			ApplicationContext context = this.contextCache.get(mergedConfig);
 			try {
-				if (context == null) {
-					int failureCount = this.contextCache.getFailureCount(mergedConfig);
-					if (failureCount >= this.failureThreshold) {
-						throw new IllegalStateException("""
-								ApplicationContext failure threshold (%d) exceeded: \
-								skipping repeated attempt to load context for %s"""
-									.formatted(this.failureThreshold, mergedConfig));
+				ApplicationContext context = this.contextCache.get(mergedConfig);
+				if (context != null) {
+					if (logger.isTraceEnabled()) {
+						logger.trace("Retrieved ApplicationContext [%s] from cache with key %s".formatted(
+								System.identityHashCode(context), mergedConfig));
 					}
+					return context;
+				}
+
+				int failureCount = this.contextCache.getFailureCount(mergedConfig);
+				if (failureCount >= this.failureThreshold) {
+					throw new IllegalStateException("""
+							ApplicationContext failure threshold (%d) exceeded: \
+							skipping repeated attempt to load context for %s"""
+								.formatted(this.failureThreshold, mergedConfig));
+				}
+
+				return this.contextCache.put(mergedConfig, key -> {
 					try {
-						if (mergedConfig instanceof AotMergedContextConfiguration aotMergedConfig) {
-							context = loadContextInAotMode(aotMergedConfig);
+						ApplicationContext newContext;
+						if (key instanceof AotMergedContextConfiguration aotMergedConfig) {
+							newContext = loadContextInAotMode(aotMergedConfig);
 						}
 						else {
-							context = loadContextInternal(mergedConfig);
+							newContext = loadContextInternal(key);
 						}
 						if (logger.isTraceEnabled()) {
 							logger.trace("Storing ApplicationContext [%s] in cache under key %s".formatted(
-									System.identityHashCode(context), mergedConfig));
+									System.identityHashCode(newContext), key));
 						}
-						this.contextCache.put(mergedConfig, context);
+						return newContext;
 					}
 					catch (Exception ex) {
 						if (logger.isTraceEnabled()) {
-							logger.trace("Incrementing ApplicationContext failure count for " + mergedConfig);
+							logger.trace("Incrementing ApplicationContext failure count for " + key);
 						}
-						this.contextCache.incrementFailureCount(mergedConfig);
+						this.contextCache.incrementFailureCount(key);
 						Throwable cause = ex;
 						if (ex instanceof ContextLoadException cle) {
 							cause = cle.getCause();
@@ -182,22 +192,13 @@ ApplicationContext failure threshold (%d) exceeded: \
 								}
 							}
 						}
-						throw new IllegalStateException(
-								"Failed to load ApplicationContext for " + mergedConfig, cause);
+						throw new IllegalStateException("Failed to load ApplicationContext for " + key, cause);
 					}
-				}
-				else {
-					if (logger.isTraceEnabled()) {
-						logger.trace("Retrieved ApplicationContext [%s] from cache with key %s".formatted(
-								System.identityHashCode(context), mergedConfig));
-					}
-				}
+				});
 			}
 			finally {
 				this.contextCache.logStatistics();
 			}
-
-			return context;
 		}
 	}
 

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

@@ -19,6 +19,7 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashSet;
+import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -62,7 +63,7 @@ public class DefaultContextCache implements ContextCache {
 	 * Map of context keys to Spring {@code ApplicationContext} instances.
 	 */
 	private final Map<MergedContextConfiguration, ApplicationContext> contextMap =
-			Collections.synchronizedMap(new LruCache(32, 0.75f));
+			Collections.synchronizedMap(new LinkedHashMap<>(32, 0.75f, true));
 
 	/**
 	 * Map of parent keys to sets of children keys, representing a top-down <em>tree</em>
@@ -157,7 +158,41 @@ public void put(MergedContextConfiguration key, ApplicationContext context) {
 		Assert.notNull(key, "Key must not be null");
 		Assert.notNull(context, "ApplicationContext must not be null");
 
+		evictLruContextIfNecessary();
+		putInternal(key, context);
+	}
+
+	@Override
+	public ApplicationContext put(MergedContextConfiguration key, LoadFunction loadFunction) {
+		Assert.notNull(key, "Key must not be null");
+		Assert.notNull(loadFunction, "LoadFunction must not be null");
+
+		evictLruContextIfNecessary();
+		ApplicationContext context = loadFunction.loadContext(key);
+		Assert.state(context != null, "LoadFunction must return a non-null ApplicationContext");
+		putInternal(key, context);
+		return context;
+	}
+
+	/**
+	 * Evict the least recently used (LRU) context if necessary.
+	 * @since 7.0
+	 */
+	private void evictLruContextIfNecessary() {
+		if (this.contextMap.size() >= this.maxSize) {
+			Iterator<MergedContextConfiguration> iterator = this.contextMap.keySet().iterator();
+			Assert.state(iterator.hasNext(), "Failed to retrieve LRU context");
+			// The least recently used (LRU) key is the first/head in a LinkedHashMap
+			// configured for access-order iteration order.
+			MergedContextConfiguration lruKey = iterator.next();
+			remove(lruKey, HierarchyMode.CURRENT_LEVEL);
+		}
+	}
+
+	private void putInternal(MergedContextConfiguration key, ApplicationContext context) {
 		this.contextMap.put(key, context);
+
+		// Update context hierarchy map.
 		MergedContextConfiguration child = key;
 		MergedContextConfiguration parent = child.getParent();
 		while (parent != null) {
@@ -357,37 +392,4 @@ public String toString() {
 				.toString();
 	}
 
-
-	/**
-	 * Simple cache implementation based on {@link LinkedHashMap} with a maximum
-	 * size and a <em>least recently used</em> (LRU) eviction policy that
-	 * properly closes application contexts.
-	 * @since 4.3
-	 */
-	@SuppressWarnings("serial")
-	private class LruCache extends LinkedHashMap<MergedContextConfiguration, ApplicationContext> {
-
-		/**
-		 * Create a new {@code LruCache} with the supplied initial capacity
-		 * and load factor.
-		 * @param initialCapacity the initial capacity
-		 * @param loadFactor the load factor
-		 */
-		LruCache(int initialCapacity, float loadFactor) {
-			super(initialCapacity, loadFactor, true);
-		}
-
-		@Override
-		protected boolean removeEldestEntry(Map.Entry<MergedContextConfiguration, ApplicationContext> eldest) {
-			if (this.size() > DefaultContextCache.this.getMaxSize()) {
-				// Do NOT delete "DefaultContextCache.this."; otherwise, we accidentally
-				// invoke java.util.Map.remove(Object, Object).
-				DefaultContextCache.this.remove(eldest.getKey(), HierarchyMode.CURRENT_LEVEL);
-			}
-
-			// Return false since we invoke a custom eviction algorithm.
-			return false;
-		}
-	}
-
 }