Prevent cache from consuming too much memory

Change the cache in `CachingOperationInvoker` to be a reference based
map and also clean stale entries when a specific threshold is met.

Fixes gh-28313

Author: philwebb

spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/endpoint/invoker/cache/CachingOperationInvoker.java

@@ -18,9 +18,10 @@
 
 import java.security.Principal;
 import java.time.Duration;
+import java.util.Iterator;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Objects;
-import java.util.concurrent.ConcurrentHashMap;
 
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -30,6 +31,7 @@
 import org.springframework.boot.actuate.endpoint.invoke.OperationInvoker;
 import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
+import org.springframework.util.ConcurrentReferenceHashMap;
 import org.springframework.util.ObjectUtils;
 
 /**
@@ -45,6 +47,8 @@ public class CachingOperationInvoker implements OperationInvoker {
 
 	private static final boolean IS_REACTOR_PRESENT = ClassUtils.isPresent("reactor.core.publisher.Mono", null);
 
+	private static final int CACHE_CLEANUP_THRESHOLD = 40;
+
 	private final OperationInvoker invoker;
 
 	private final long timeToLive;
@@ -61,7 +65,7 @@ public class CachingOperationInvoker implements OperationInvoker {
 		Assert.isTrue(timeToLive > 0, "TimeToLive must be strictly positive");
 		this.invoker = invoker;
 		this.timeToLive = timeToLive;
-		this.cachedResponses = new ConcurrentHashMap<>();
+		this.cachedResponses = new ConcurrentReferenceHashMap<>();
 	}
 
 	/**
@@ -78,6 +82,9 @@ public Object invoke(InvocationContext context) {
 			return this.invoker.invoke(context);
 		}
 		long accessTime = System.currentTimeMillis();
+		if (this.cachedResponses.size() > CACHE_CLEANUP_THRESHOLD) {
+			cleanExpiredCachedResponses(accessTime);
+		}
 		ApiVersion contextApiVersion = context.getApiVersion();
 		CacheKey cacheKey = new CacheKey(contextApiVersion, context.getSecurityContext().getPrincipal());
 		CachedResponse cached = this.cachedResponses.get(cacheKey);
@@ -89,6 +96,20 @@ public Object invoke(InvocationContext context) {
 		return cached.getResponse();
 	}
 
+	private void cleanExpiredCachedResponses(long accessTime) {
+		try {
+			Iterator<Entry<CacheKey, CachedResponse>> iterator = this.cachedResponses.entrySet().iterator();
+			while (iterator.hasNext()) {
+				Entry<CacheKey, CachedResponse> entry = iterator.next();
+				if (entry.getValue().isStale(accessTime, this.timeToLive)) {
+					iterator.remove();
+				}
+			}
+		}
+		catch (Exception ex) {
+		}
+	}
+
 	private boolean hasInput(InvocationContext context) {
 		Map<String, Object> arguments = context.getArguments();
 		if (!ObjectUtils.isEmpty(arguments)) {

spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/endpoint/invoker/cache/CachingOperationInvokerTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2020 the original author or authors.
+ * 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.
@@ -23,6 +23,7 @@
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import org.assertj.core.api.InstanceOfAssertFactories;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -101,6 +102,31 @@ void cacheInTtlWithFluxResponse() {
 		assertThat(response).isSameAs(cachedResponse);
 	}
 
+	@Test // gh-28313
+	void cacheWhenEachPrincipalIsUniqueDoesNotConsumeTooMuchMemory() throws Exception {
+		MonoOperationInvoker target = new MonoOperationInvoker();
+		CachingOperationInvoker invoker = new CachingOperationInvoker(target, 50L);
+		int count = 1000;
+		for (int i = 0; i < count; i++) {
+			invokeWithUniquePrincipal(invoker);
+		}
+		long expired = System.currentTimeMillis() + 50;
+		while (System.currentTimeMillis() < expired) {
+			Thread.sleep(10);
+		}
+		invokeWithUniquePrincipal(invoker);
+		assertThat(invoker).extracting("cachedResponses").asInstanceOf(InstanceOfAssertFactories.MAP)
+				.hasSizeLessThan(count);
+	}
+
+	private void invokeWithUniquePrincipal(CachingOperationInvoker invoker) {
+		SecurityContext securityContext = mock(SecurityContext.class);
+		Principal principal = mock(Principal.class);
+		given(securityContext.getPrincipal()).willReturn(principal);
+		InvocationContext context = new InvocationContext(securityContext, Collections.emptyMap());
+		((Mono<?>) invoker.invoke(context)).block();
+	}
+
 	private void assertCacheIsUsed(Map<String, Object> parameters) {
 		assertCacheIsUsed(parameters, null);
 	}