Fix the usage of CacheIteratorHelper for service account (#75510) (#75765)

CacheIteratorHelper requires lock acquisition for any mutation to the underlying cache. This means it is incorrect to manipulate the cache without invocation of CacheIteratorHelper#acquireUpdateLock. This is OK for caches of simple values, but feels excessive for caches of ListenableFuture. This PR update the cache invalidation code to use cache.forEach instead of CacheInvalidator. It simplifies the code by removing any explicit lockings. The tradeoff is that it needs to build a list of keys to delete in memory. Overall it is a better tradeoff since no explicit locking is required and better leverage of Cache's own methods. Co-authored-by: Yang Wang <yang.wang@elastic.co>

Author: elasticsearchmachine

x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/authc/service/CachingServiceAccountTokenStore.java

@@ -20,12 +20,17 @@
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xpack.core.security.action.service.TokenInfo.TokenSource;
 import org.elasticsearch.xpack.core.security.authc.support.Hasher;
-import org.elasticsearch.xpack.core.security.support.CacheIteratorHelper;
 import org.elasticsearch.xpack.security.support.CacheInvalidatorRegistry;
 
+import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Predicate;
 
 public abstract class CachingServiceAccountTokenStore implements ServiceAccountTokenStore, CacheInvalidatorRegistry.CacheInvalidator {
 
@@ -42,7 +47,6 @@ public abstract class CachingServiceAccountTokenStore implements ServiceAccountT
  private final Settings settings;
  private final ThreadPool threadPool;
  private final Cache<String, ListenableFuture<CachedResult>> cache;
- private CacheIteratorHelper<String, ListenableFuture<CachedResult>> cacheIteratorHelper;
  private final Hasher hasher;
 
  CachingServiceAccountTokenStore(Settings settings, ThreadPool threadPool) {
@@ -54,10 +58,8 @@ public abstract class CachingServiceAccountTokenStore implements ServiceAccountT
  .setExpireAfterWrite(ttl)
  .setMaximumWeight(CACHE_MAX_TOKENS_SETTING.get(settings))
  .build();
- cacheIteratorHelper = new CacheIteratorHelper<>(cache);
  } else {
  cache = null;
- cacheIteratorHelper = null;
  }
  hasher = Hasher.resolve(CACHE_HASH_ALGO_SETTING.get(settings));
  }
@@ -126,16 +128,29 @@ private void authenticateWithCache(ServiceAccountToken token, ActionListener<Sto
  @Override
  public final void invalidate(Collection<String> qualifiedTokenNames) {
  if (cache != null) {
- logger.trace("invalidating cache for service token [{}]",
- Strings.collectionToCommaDelimitedString(qualifiedTokenNames));
- for (String qualifiedTokenName : qualifiedTokenNames) {
- if (qualifiedTokenName.endsWith("/")) {
- // Wildcard case of invalidating all tokens for a service account, e.g. "elastic/fleet-server/"
- cacheIteratorHelper.removeKeysIf(key -> key.startsWith(qualifiedTokenName));
- } else {
- cache.invalidate(qualifiedTokenName);
+ logger.trace("invalidating cache for service token [{}]", Strings.collectionToCommaDelimitedString(qualifiedTokenNames));
+ final Set<String> exacts = new HashSet<>(qualifiedTokenNames);
+ final Set<String> prefixes = new HashSet<>();
+ final Iterator<String> it = exacts.iterator();
+ while (it.hasNext()) {
+ final String name = it.next();
+ if (name.endsWith("/")) {
+ prefixes.add(name);
+ it.remove();
  }
  }
+
+ exacts.forEach(cache::invalidate);
+ if (false == prefixes.isEmpty()) {
+ final Predicate<String> predicate = k -> prefixes.stream().anyMatch(k::startsWith);
+ final List<String> keys = new ArrayList<>();
+ cache.forEach((k, v) -> {
+ if (predicate.test(k)) {
+ keys.add(k);
+ }
+ });
+ keys.forEach(cache::invalidate);
+ }
  }
  }