Create a LocalCacheProvider SPI (#44950)

Closes #42223

Signed-off-by: Ryan Emerson <remerson@ibm.com>

server-spi-private/src/main/java/org/keycloak/cache/AlternativeLookupProviderFactory.java

@@ -1,6 +1,13 @@
 package org.keycloak.cache;
 
+import java.util.Set;
+
+import org.keycloak.provider.Provider;
 import org.keycloak.provider.ProviderFactory;
 
 public interface AlternativeLookupProviderFactory extends ProviderFactory<AlternativeLookupProvider> {
+    @Override
+    default Set<Class<? extends Provider>> dependsOn() {
+        return Set.of(LocalCacheProvider.class);
+    }
 }

server-spi-private/src/main/java/org/keycloak/cache/LocalCache.java

@@ -0,0 +1,43 @@
+package org.keycloak.cache;
+
+/**
+ * A {@link LocalCache} should be used when a local, non-clustered, cache is required to optimise data access.
+ *
+ * @param <K> the type of the cache Keys used for lookup
+ * @param <V> the type of the cache Values to be stored
+ */
+public interface LocalCache<K, V> extends AutoCloseable {
+
+    /**
+     * Returns the value associated with the {@code key}, or {@code null} if there is no
+     * cached value for the {@code key}.
+     *
+     * @param key the key whose associated value is to be returned
+     * @return the value associated with the specified key or {@code null} if no value exists
+     * @throws NullPointerException if the specified key is null
+     */
+    V get(K key);
+
+    /**
+     * Associates the value with the key in this cache.
+     * If the cache previously contained a value associated with the key, the old value is replaced by the new value.
+     *
+     * @param key the key with which the specified value is to be associated
+     * @param value value to be associated with the specified key
+     * @throws NullPointerException if the specified key or value is null
+     */
+    void put(K key, V value);
+
+    /**
+     * Removes the cached value for the specified {@code key}.
+     *
+     * @param key the key whose mapping is to be removed from the cache
+     * @throws NullPointerException if the specified key is null
+     */
+    void invalidate(K key);
+
+    /**
+     * Closes all resources associated with the cache.
+     */
+    void close();
+}

server-spi-private/src/main/java/org/keycloak/cache/LocalCacheConfiguration.java

@@ -0,0 +1,72 @@
+package org.keycloak.cache;
+
+import java.time.Duration;
+import java.util.Objects;
+import java.util.function.Function;
+
+/**
+ * Configuration {@code record} that encapsulates all configuration required in order to create a {@link LocalCache}.
+ *
+ * @param name the name of the cache
+ * @param maxSize the maximum size of the cache, or -1 if an unbounded cache is desired.
+ * @param expiration the {@link Duration} to wait before entries are expired. A {@code null} value indicates no expiration.
+ * @param loader an optional {@link Function} which is used to retrieve cache values on cache miss.
+ * @param <K> the type of the cache Keys used for lookup
+ * @param <V> the type of the cache Values to be stored
+ */
+public record LocalCacheConfiguration<K, V>(String name, int maxSize, Duration expiration, Function<K, V> loader) {
+
+    public LocalCacheConfiguration {
+        Objects.requireNonNull(name, "A cache name must be configured");
+    }
+
+    public boolean hasExpiration() {
+        return expiration != null;
+    }
+
+    public boolean hasLoader() {
+        return loader != null;
+    }
+
+    public static <K, V> Builder<K, V> builder() {
+        return new Builder<>();
+    }
+
+    /**
+     * A builder class to simplify the creation of {@link LocalCacheConfiguration} objects.
+     *
+     * @param <K> the type of the cache Keys used for lookup
+     * @param <V> the type of the cache Values to be stored
+     */
+    public static class Builder<K, V> {
+
+        private String name;
+        private int maxSize = -1;
+        private Duration expiration;
+        private Function<K, V> loader;
+
+        public Builder<K, V> name(String name) {
+            this.name = name;
+            return this;
+        }
+
+        public Builder<K, V> maxSize(int maxSize) {
+            this.maxSize = maxSize;
+            return this;
+        }
+
+        public Builder<K, V> expiration(Duration duration) {
+            this.expiration = duration;
+            return this;
+        }
+
+        public Builder<K, V> loader(Function<K, V> loader) {
+            this.loader = loader;
+            return this;
+        }
+
+        public LocalCacheConfiguration<K, V> build() {
+            return new LocalCacheConfiguration<>(name, maxSize, expiration, loader);
+        }
+    }
+}

server-spi-private/src/main/java/org/keycloak/cache/LocalCacheProvider.java

@@ -0,0 +1,21 @@
+package org.keycloak.cache;
+
+import org.keycloak.provider.Provider;
+
+/**
+ * A {@link Provider} to abstract the creation of local, non-clustered, in-memory caches from the underlying cache implementation.
+ */
+public interface LocalCacheProvider extends Provider {
+    /**
+     * Creates a new {@link LocalCache} instance for local caching. {@link LocalCacheProvider} implementations
+     * are not responsible for managing the lifecycle of created {@link LocalCache} instances. It is the responsibility
+     * of {@link LocalCache} consumers to ensure that {@link LocalCache#close()} is called when the cache is no longer
+     * required.
+     *
+     * @param configuration the desired cache configuration
+     * @return {@link LocalCache} a newly created cache
+     * @param <K> the type of the cache Keys used for lookup
+     * @param <V> the type of the cache Values to be stored
+     */
+    <K, V> LocalCache<K, V> create(LocalCacheConfiguration<K, V> configuration);
+}

server-spi-private/src/main/java/org/keycloak/cache/LocalCacheProviderFactory.java

@@ -0,0 +1,6 @@
+package org.keycloak.cache;
+
+import org.keycloak.provider.ProviderFactory;
+
+public interface LocalCacheProviderFactory extends ProviderFactory<LocalCacheProvider> {
+}

server-spi-private/src/main/java/org/keycloak/cache/LocalCacheSPI.java

@@ -0,0 +1,27 @@
+package org.keycloak.cache;
+
+import org.keycloak.provider.Provider;
+import org.keycloak.provider.ProviderFactory;
+import org.keycloak.provider.Spi;
+
+public class LocalCacheSPI implements Spi {
+    @Override
+    public boolean isInternal() {
+        return true;
+    }
+
+    @Override
+    public String getName() {
+        return "localCache";
+    }
+
+    @Override
+    public Class<? extends Provider> getProviderClass() {
+        return LocalCacheProvider.class;
+    }
+
+    @Override
+    public Class<? extends ProviderFactory<?>> getProviderFactoryClass() {
+        return LocalCacheProviderFactory.class;
+    }
+}

server-spi-private/src/main/java/org/keycloak/device/DeviceRepresentationProviderFactory.java

@@ -1,7 +1,11 @@
 package org.keycloak.device;
 
+import java.util.Set;
+
 import org.keycloak.Config;
+import org.keycloak.cache.LocalCacheProvider;
 import org.keycloak.models.KeycloakSessionFactory;
+import org.keycloak.provider.Provider;
 import org.keycloak.provider.ProviderFactory;
 
 public interface DeviceRepresentationProviderFactory extends ProviderFactory<DeviceRepresentationProvider> {
@@ -17,4 +21,9 @@ public interface DeviceRepresentationProviderFactory extends ProviderFactory<Dev
     @Override
     default void close() {
     }
+
+    @Override
+    default Set<Class<? extends Provider>> dependsOn() {
+        return Set.of(LocalCacheProvider.class);
+    }
 }

server-spi-private/src/main/resources/META-INF/services/org.keycloak.provider.Spi

@@ -110,4 +110,5 @@ org.keycloak.models.workflow.WorkflowSpi
 org.keycloak.models.workflow.WorkflowConditionSpi
 org.keycloak.protocol.oidc.rar.AuthorizationDetailsProcessorSpi
 org.keycloak.cache.AlternativeLookupSPI
+org.keycloak.cache.LocalCacheSPI
 org.keycloak.models.workflow.WorkflowStateSpi

services/src/main/java/org/keycloak/cache/DefaultAlternativeLookupProvider.java

@@ -8,20 +8,18 @@ import org.keycloak.models.IdentityProviderModel;
 import org.keycloak.models.IdentityProviderQuery;
 import org.keycloak.models.KeycloakSession;
 
-import com.github.benmanes.caffeine.cache.Cache;
-
 public class DefaultAlternativeLookupProvider implements AlternativeLookupProvider {
 
-    private final Cache<String, String> lookupCache;
+    private final LocalCache<String, String> lookupCache;
 
-    public DefaultAlternativeLookupProvider(Cache<String, String> lookupCache) {
+    public DefaultAlternativeLookupProvider(LocalCache<String, String> lookupCache) {
         this.lookupCache = lookupCache;
     }
 
     public IdentityProviderModel lookupIdentityProviderFromIssuer(KeycloakSession session, String issuerUrl) {
         String alternativeKey = ComputedKey.computeKey(session.getContext().getRealm().getId(), "idp", issuerUrl);
 
-        String cachedIdpAlias = lookupCache.getIfPresent(alternativeKey);
+        String cachedIdpAlias = lookupCache.get(alternativeKey);
         if (cachedIdpAlias != null) {
             IdentityProviderModel idp = session.identityProviders().getByAlias(cachedIdpAlias);
             if (idp != null && issuerUrl.equals(idp.getConfig().get(IdentityProviderModel.ISSUER))) {
@@ -43,7 +41,7 @@ public class DefaultAlternativeLookupProvider implements AlternativeLookupProvid
     public ClientModel lookupClientFromClientAttributes(KeycloakSession session, Map<String, String> attributes) {
         String alternativeKey = ComputedKey.computeKey(session.getContext().getRealm().getId(), "client", attributes);
 
-        String cachedClientId = lookupCache.getIfPresent(alternativeKey);
+        String cachedClientId = lookupCache.get(alternativeKey);
         if (cachedClientId != null) {
             ClientModel client = session.clients().getClientByClientId(session.getContext().getRealm(), cachedClientId);
             boolean match = client != null;

services/src/main/java/org/keycloak/cache/DefaultAlternativeLookupProviderFactory.java

@@ -1,19 +1,15 @@
 package org.keycloak.cache;
 
-import java.util.concurrent.TimeUnit;
+import java.time.Duration;
 
 import org.keycloak.Config;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.KeycloakSessionFactory;
 
-import com.github.benmanes.caffeine.cache.Cache;
-import com.github.benmanes.caffeine.cache.Caffeine;
-import io.micrometer.core.instrument.Metrics;
-import io.micrometer.core.instrument.binder.cache.CaffeineStatsCounter;
-
 public class DefaultAlternativeLookupProviderFactory implements AlternativeLookupProviderFactory {
 
-    private Cache<String, String> lookupCache;
+    private LocalCacheConfiguration<String, String> cacheConfig;
+    private LocalCache<String, String> lookupCache;
 
     @Override
     public String getId() {
@@ -30,25 +26,26 @@ public class DefaultAlternativeLookupProviderFactory implements AlternativeLooku
         Integer maximumSize = config.getInt("maximumSize", 1000);
         Integer expireAfter = config.getInt("expireAfter", 60);
 
-        CaffeineStatsCounter metrics = new CaffeineStatsCounter(Metrics.globalRegistry, "lookup");
-
-        this.lookupCache = Caffeine.newBuilder()
-                .maximumSize(maximumSize)
-                .expireAfterAccess(expireAfter, TimeUnit.MINUTES)
-                .recordStats(() -> metrics)
-                .build();
-
-        metrics.registerSizeMetric(lookupCache);
+        cacheConfig = LocalCacheConfiguration.<String, String>builder()
+              .name("lookup")
+              .expiration(Duration.ofMinutes(expireAfter))
+              .maxSize(maximumSize)
+              .build();
     }
 
     @Override
     public void postInit(KeycloakSessionFactory factory) {
+        try (KeycloakSession session = factory.create()) {
+            lookupCache = session.getProvider(LocalCacheProvider.class).create(cacheConfig);
+            cacheConfig = null;
+        }
     }
 
     @Override
     public void close() {
-        lookupCache.cleanUp();
-        lookupCache = null;
+        if (lookupCache != null) {
+            lookupCache.close();
+            lookupCache = null;
+        }
     }
-
 }

services/src/main/java/org/keycloak/cache/DefaultLocalCacheProviderFactory.java

@@ -0,0 +1,114 @@
+package org.keycloak.cache;
+
+import java.util.Objects;
+
+import org.keycloak.Config;
+import org.keycloak.models.KeycloakSession;
+import org.keycloak.models.KeycloakSessionFactory;
+
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
+import com.github.benmanes.caffeine.cache.LoadingCache;
+import io.micrometer.core.instrument.Metrics;
+import io.micrometer.core.instrument.binder.cache.CaffeineStatsCounter;
+
+/**
+ * The default implementation for {@link LocalCacheProvider} and {@link LocalCacheProviderFactory}.
+ */
+public class DefaultLocalCacheProviderFactory implements LocalCacheProvider, LocalCacheProviderFactory {
+
+    @Override
+    public LocalCacheProvider create(KeycloakSession session) {
+        return this;
+    }
+
+    @Override
+    public void init(Config.Scope config) {
+    }
+
+    @Override
+    public void postInit(KeycloakSessionFactory factory) {
+    }
+
+    @Override
+    public String getId() {
+        return "default";
+    }
+
+    @Override
+    public <K, V> LocalCache<K, V> create(LocalCacheConfiguration<K, V> configuration) {
+        CaffeineStatsCounter metrics = new CaffeineStatsCounter(Metrics.globalRegistry, configuration.name());
+        Caffeine<Object, Object> builder = Caffeine.newBuilder().recordStats(() -> metrics);
+
+        if (configuration.maxSize() > 0) {
+            builder.maximumSize(configuration.maxSize());
+        }
+
+        if (configuration.hasExpiration()) {
+            builder.expireAfterAccess(configuration.expiration());
+        }
+
+        if (configuration.hasLoader()) {
+            LoadingCache<K, V> cache = builder.build(k -> configuration.loader().apply(k));
+            metrics.registerSizeMetric(cache);
+            return new LoadingCaffeineWrapper<>(cache);
+        } else {
+            Cache<K, V> cache = builder.build();
+            metrics.registerSizeMetric(cache);
+            return new CaffeineWrapper<>(cache);
+        }
+    }
+
+    @Override
+    public void close() {
+    }
+
+    private static class CaffeineWrapper<K, V> implements LocalCache<K, V> {
+
+        final Cache<K, V> cache;
+
+        CaffeineWrapper(Cache<K, V> cache) {
+            this.cache = cache;
+        }
+
+        @Override
+        public V get(K key) {
+            Objects.requireNonNull(key);
+            return cache.getIfPresent(key);
+        }
+
+        @Override
+        public void put(K key, V value) {
+            Objects.requireNonNull(key);
+            Objects.requireNonNull(value);
+            cache.put(key, value);
+        }
+
+        @Override
+        public void invalidate(K key) {
+            Objects.requireNonNull(key);
+            cache.invalidate(key);
+        }
+
+        @Override
+        public void close() {
+            cache.cleanUp();
+        }
+    }
+
+    private static class LoadingCaffeineWrapper<K, V> extends CaffeineWrapper<K, V> {
+
+        final LoadingCache<K, V> cache;
+
+        LoadingCaffeineWrapper(LoadingCache<K, V> cache) {
+            super(cache);
+            this.cache = cache;
+        }
+
+        @Override
+        public V get(K key) {
+            Objects.requireNonNull(key);
+            return cache.get(key);
+        }
+    }
+}

services/src/main/java/org/keycloak/device/DeviceRepresentationProviderFactoryImpl.java

@@ -3,14 +3,14 @@ package org.keycloak.device;
 import java.util.List;
 
 import org.keycloak.Config;
+import org.keycloak.cache.LocalCache;
+import org.keycloak.cache.LocalCacheConfiguration;
+import org.keycloak.cache.LocalCacheProvider;
 import org.keycloak.models.KeycloakSession;
+import org.keycloak.models.KeycloakSessionFactory;
 import org.keycloak.provider.ProviderConfigProperty;
 import org.keycloak.provider.ProviderConfigurationBuilder;
 
-import com.github.benmanes.caffeine.cache.Caffeine;
-import com.github.benmanes.caffeine.cache.LoadingCache;
-import io.micrometer.core.instrument.Metrics;
-import io.micrometer.core.instrument.binder.cache.CaffeineStatsCounter;
 import ua_parser.Client;
 import ua_parser.Parser;
 
@@ -21,11 +21,11 @@ public class DeviceRepresentationProviderFactoryImpl implements DeviceRepresenta
     // The max user agent size is 512 bytes and it will take 1024 bytes per cache entry.
     // Using 2MB for caching.
     private static final int DEFAULT_CACHE_SIZE = 2048;
-
-    private volatile LoadingCache<String, Client> cache;
-
     public static final String PROVIDER_ID = "deviceRepresentation";
 
+    private LocalCacheConfiguration<String, Client> cacheConfig;
+    private LocalCache<String, Client> cache;
+
     @Override
     public String getId() {
         return PROVIDER_ID;
@@ -33,13 +33,19 @@ public class DeviceRepresentationProviderFactoryImpl implements DeviceRepresenta
 
     @Override
     public void init(Config.Scope config) {
-        CaffeineStatsCounter metrics = new CaffeineStatsCounter(Metrics.globalRegistry, "userAgent");
-        cache = Caffeine.newBuilder()
-                .maximumSize(config.getInt(CACHE_SIZE, DEFAULT_CACHE_SIZE))
-                .recordStats(() -> metrics)
-                .softValues()
-                .build(UA_PARSER::parse);
-        metrics.registerSizeMetric(cache);
+        cacheConfig = LocalCacheConfiguration.<String, Client>builder()
+              .name("userAgent")
+              .maxSize(config.getInt(CACHE_SIZE, DEFAULT_CACHE_SIZE))
+              .loader(UA_PARSER::parse)
+              .build();
+    }
+
+    @Override
+    public void postInit(KeycloakSessionFactory factory) {
+        try (KeycloakSession session = factory.create()) {
+            cache = session.getProvider(LocalCacheProvider.class).create(cacheConfig);
+            cacheConfig = null;
+        }
     }
 
     @Override
@@ -58,4 +64,12 @@ public class DeviceRepresentationProviderFactoryImpl implements DeviceRepresenta
                 .add()
                 .build();
     }
+
+    @Override
+    public void close() {
+        if (cache != null) {
+            cache.close();
+            cache = null;
+        }
+    }
 }

services/src/main/java/org/keycloak/device/DeviceRepresentationProviderImpl.java

@@ -2,11 +2,11 @@ package org.keycloak.device;
 
 import jakarta.ws.rs.core.HttpHeaders;
 
+import org.keycloak.cache.LocalCache;
 import org.keycloak.models.KeycloakContext;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.representations.account.DeviceRepresentation;
 
-import com.github.benmanes.caffeine.cache.LoadingCache;
 import org.jboss.logging.Logger;
 import ua_parser.Client;
 
@@ -14,10 +14,10 @@ public class DeviceRepresentationProviderImpl implements DeviceRepresentationPro
     private static final Logger logger = Logger.getLogger(DeviceActivityManager.class);
     private static final int USER_AGENT_MAX_LENGTH = 512;
 
-    private final LoadingCache<String, Client> cache;
+    private final LocalCache<String, Client> cache;
     private final KeycloakSession session;
 
-    DeviceRepresentationProviderImpl(KeycloakSession session, LoadingCache<String, Client> cache) {
+    DeviceRepresentationProviderImpl(KeycloakSession session, LocalCache<String, Client> cache) {
         this.session = session;
         this.cache = cache;
     }

services/src/main/resources/META-INF/services/org.keycloak.cache.LocalCacheProviderFactory

@@ -0,0 +1 @@
+org.keycloak.cache.DefaultLocalCacheProviderFactory