Optimize caching and use of DB connections when Organisations are enabled

Closes #33353

Signed-off-by: Alexander Schwartz <aschwart@redhat.com>

model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/organization/InfinispanOrganizationProvider.java

@@ -38,13 +38,12 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
     private static final String ORG_MEMBERS_COUNT_KEY_SUFFIX = ".members.count";
 
     private final KeycloakSession session;
-    private final OrganizationProvider orgDelegate;
+    private OrganizationProvider orgDelegate;
     private final RealmCacheSession realmCache;
     private final Map<String, OrganizationAdapter> managedOrganizations = new HashMap<>();
 
     public InfinispanOrganizationProvider(KeycloakSession session) {
         this.session = session;
-        this.orgDelegate = session.getProvider(OrganizationProvider.class, "jpa");
         this.realmCache = (RealmCacheSession) session.getProvider(CacheRealmProvider.class);
     }
 
@@ -59,14 +58,22 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
     @Override
     public OrganizationModel create(String id, String name, String alias) {
         registerCountInvalidation();
-        return orgDelegate.create(id, name, alias);
+        return getDelegate().create(id, name, alias);
+    }
+
+    private OrganizationProvider getDelegate() {
+        if (orgDelegate == null) {
+            // use lazy initialization to avoid touching the entity manager
+            orgDelegate = session.getProvider(OrganizationProvider.class, "jpa");
+        }
+        return orgDelegate;
     }
 
     @Override
     public boolean remove(OrganizationModel organization) {
         registerOrganizationInvalidation(organization);
         registerCountInvalidation();
-        return orgDelegate.remove(organization);
+        return getDelegate().remove(organization);
     }
 
     @Override
@@ -79,7 +86,7 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
 
         if (cached == null) {
             Long loaded = realmCache.getCache().getCurrentRevision(id);
-            OrganizationModel model = orgDelegate.getById(id);
+            OrganizationModel model = getDelegate().getById(id);
             if (model == null) return null;
             if (isInvalid(id)) return model;
             cached = new CachedOrganization(loaded, getRealm(), model);
@@ -87,11 +94,11 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
 
         // no need to check for realm invalidation as IdP changes are handled by events within InfinispanOrganizationProviderFactory
         } else if (isInvalid(id)) {
-            return orgDelegate.getById(id);
+            return getDelegate().getById(id);
         } else if (managedOrganizations.containsKey(id)) {
             return managedOrganizations.get(id);
         }
-        OrganizationAdapter adapter = new OrganizationAdapter(cached, realmCache, orgDelegate, this);
+        OrganizationAdapter adapter = new OrganizationAdapter(cached, () -> getDelegate(), this);
         managedOrganizations.put(id, adapter);
         return adapter;
     }
@@ -101,18 +108,18 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
         String cacheKey = cacheKeyByDomain(domainName);
 
         if (isInvalid(cacheKey)) {
-            return orgDelegate.getByDomainName(domainName);
+            return getDelegate().getByDomainName(domainName);
         }
 
         CachedOrganizationIds cached = realmCache.getCache().get(cacheKey, CachedOrganizationIds.class);
 
         if (cached == null) {
             Long loaded = realmCache.getCache().getCurrentRevision(cacheKey);
-            OrganizationModel model = orgDelegate.getByDomainName(domainName);
+            OrganizationModel model = getDelegate().getByDomainName(domainName);
             if (model == null) {
                 return null;
             }
-            cached = new CachedOrganizationIds(loaded, cacheKey, getRealm(), model);
+            cached = new CachedOrganizationIds(loaded, cacheKey, getRealm(), Stream.ofNullable(model));
             realmCache.getCache().addRevisioned(cached, realmCache.getStartupRevision());
         }
 
@@ -122,13 +129,13 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
     @Override
     public Stream<OrganizationModel> getAllStream(String search, Boolean exact, Integer first, Integer max) {
         // Return cache delegates to ensure cache invalidation during write operations
-        return getCacheDelegates(orgDelegate.getAllStream(search, exact, first, max));
+        return getCacheDelegates(getDelegate().getAllStream(search, exact, first, max));
     }
 
     @Override
     public Stream<OrganizationModel> getAllStream(Map<String, String> attributes, Integer first, Integer max) {
         // Return cache delegates to ensure cache invalidation during write operations
-        return getCacheDelegates(orgDelegate.getAllStream(attributes, first, max));
+        return getCacheDelegates(getDelegate().getAllStream(attributes, first, max));
     }
 
     @Override
@@ -141,24 +148,24 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
     @Override
     public boolean addManagedMember(OrganizationModel organization, UserModel user) {
         registerMemberInvalidation(organization, user);
-        return orgDelegate.addManagedMember(organization, user);
+        return getDelegate().addManagedMember(organization, user);
     }
 
     @Override
     public boolean addMember(OrganizationModel organization, UserModel user) {
         registerMemberInvalidation(organization, user);
-        return orgDelegate.addMember(organization, user);
+        return getDelegate().addMember(organization, user);
     }
 
     @Override
     public boolean removeMember(OrganizationModel organization, UserModel member) {
         registerMemberInvalidation(organization, member);
-        return orgDelegate.removeMember(organization, member);
+        return getDelegate().removeMember(organization, member);
     }
 
     @Override
     public Stream<UserModel> getMembersStream(OrganizationModel organization, String search, Boolean exact, Integer first, Integer max) {
-        return orgDelegate.getMembersStream(organization, search, exact, first, max);
+        return getDelegate().getMembersStream(organization, search, exact, first, max);
     }
 
     @Override
@@ -172,7 +179,7 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
         }
 
         Long loaded = realmCache.getCache().getCurrentRevision(cacheKey);
-        long membersCount = orgDelegate.getMembersCount(organization);
+        long membersCount = getDelegate().getMembersCount(organization);
         cached = new CachedCount(loaded, getRealm(), cacheKey, membersCount);
         realmCache.getCache().addRevisioned(cached, realmCache.getStartupRevision());
 
@@ -191,15 +198,15 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
         String cacheKey = cacheKeyMembership(realm, organization, user);
 
         if (isInvalid(cacheKey)) {
-            return orgDelegate.getMemberById(organization, user.getId());
+            return getDelegate().getMemberById(organization, user.getId());
         }
 
         CachedMembership cached = realmCache.getCache().get(cacheKey, CachedMembership.class);
 
         if (cached == null) {
-            boolean isManaged = orgDelegate.isManagedMember(organization, user);
+            boolean isManaged = getDelegate().isManagedMember(organization, user);
             Long loaded = realmCache.getCache().getCurrentRevision(cacheKey);
-            UserModel member = orgDelegate.getMemberById(organization, user.getId());
+            UserModel member = getDelegate().getMemberById(organization, user.getId());
             cached = new CachedMembership(loaded, cacheKeyMembership(realm, organization, user), realm, isManaged, member != null);
             realmCache.getCache().addRevisioned(cached, realmCache.getStartupRevision());
         }
@@ -212,14 +219,14 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
         String cacheKey = cacheKeyByMember(member);
 
         if (isInvalid(cacheKey)) {
-            return orgDelegate.getByMember(member);
+            return getDelegate().getByMember(member);
         }
 
         CachedOrganizationIds cached = realmCache.getCache().get(cacheKey, CachedOrganizationIds.class);
 
         if (cached == null) {
             Long loaded = realmCache.getCache().getCurrentRevision(cacheKey);
-            Stream<OrganizationModel> model = orgDelegate.getByMember(member);
+            Stream<OrganizationModel> model = getDelegate().getByMember(member);
             cached = new CachedOrganizationIds(loaded, cacheKey, getRealm(), model);
             realmCache.getCache().addRevisioned(cached, realmCache.getStartupRevision());
         }
@@ -239,7 +246,7 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
         CachedMembership cached = realmCache.getCache().get(cacheKey, CachedMembership.class);
 
         if (cached == null) {
-            return orgDelegate.isManagedMember(organization, user);
+            return getDelegate().isManagedMember(organization, user);
         }
 
         return cached.isManaged();
@@ -248,7 +255,7 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
 
     @Override
     public boolean addIdentityProvider(OrganizationModel organization, IdentityProviderModel identityProvider) {
-        boolean added = orgDelegate.addIdentityProvider(organization, identityProvider);
+        boolean added = getDelegate().addIdentityProvider(organization, identityProvider);
         if (added) {
             registerOrganizationInvalidation(organization);
         }
@@ -257,12 +264,12 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
 
     @Override
     public Stream<IdentityProviderModel> getIdentityProviders(OrganizationModel organization) {
-        return orgDelegate.getIdentityProviders(organization);
+        return getDelegate().getIdentityProviders(organization);
     }
 
     @Override
     public boolean removeIdentityProvider(OrganizationModel organization, IdentityProviderModel identityProvider) {
-        boolean removed = orgDelegate.removeIdentityProvider(organization, identityProvider);
+        boolean removed = getDelegate().removeIdentityProvider(organization, identityProvider);
         if (removed) {
             registerOrganizationInvalidation(organization);
         }
@@ -285,7 +292,7 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
         }
 
         Long loaded = realmCache.getCache().getCurrentRevision(cacheKey);
-        long count = orgDelegate.count();
+        long count = getDelegate().count();
         cached = new CachedCount(loaded, getRealm(), cacheKey, count);
         realmCache.getCache().addRevisioned(cached, realmCache.getStartupRevision());
 
@@ -294,7 +301,9 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
 
     @Override
     public void close() {
-        orgDelegate.close();
+        if (orgDelegate != null) {
+            getDelegate().close();
+        }
     }
 
     void registerOrganizationInvalidation(OrganizationModel organization) {

model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/organization/OrganizationAdapter.java

@@ -25,7 +25,6 @@ import org.keycloak.models.IdentityProviderModel;
 import org.keycloak.models.OrganizationDomainModel;
 import org.keycloak.models.OrganizationModel;
 import org.keycloak.models.UserModel;
-import org.keycloak.models.cache.CacheRealmProvider;
 import org.keycloak.organization.OrganizationProvider;
 
 public class OrganizationAdapter implements OrganizationModel {
@@ -33,14 +32,12 @@ public class OrganizationAdapter implements OrganizationModel {
     private volatile boolean invalidated;
     private volatile OrganizationModel updated;
     private final Supplier<OrganizationModel> modelSupplier;
-    private final CacheRealmProvider realmCache;
     private final CachedOrganization cached;
-    private final OrganizationProvider delegate;
+    private final Supplier<OrganizationProvider> delegate;
     private final InfinispanOrganizationProvider organizationCache;
 
-    public OrganizationAdapter(CachedOrganization cached, CacheRealmProvider realmCache, OrganizationProvider delegate, InfinispanOrganizationProvider organizationCache) {
+    public OrganizationAdapter(CachedOrganization cached, Supplier<OrganizationProvider> delegate, InfinispanOrganizationProvider organizationCache) {
         this.cached = cached;
-        this.realmCache = realmCache;
         this.delegate = delegate;
         this.organizationCache = organizationCache;
         this.modelSupplier = this::getOrganizationModel;
@@ -51,7 +48,7 @@ public class OrganizationAdapter implements OrganizationModel {
     }
 
     private OrganizationModel getOrganizationModel() {
-        return delegate.getById(cached.getId());
+        return delegate.get().getById(cached.getId());
     }
 
     private boolean isUpdated() {
@@ -156,12 +153,12 @@ public class OrganizationAdapter implements OrganizationModel {
 
     @Override
     public boolean isManaged(UserModel user) {
-        return delegate.isManagedMember(this, user);
+        return delegate.get().isManagedMember(this, user);
     }
 
     @Override
     public boolean isMember(UserModel user) {
-        return delegate.isMember(this, user);
+        return delegate.get().isMember(this, user);
     }
 
     @Override

services/src/main/java/org/keycloak/organization/utils/Organizations.java

@@ -26,9 +26,7 @@ import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
-import java.util.Set;
 import java.util.function.Consumer;
-import java.util.stream.Collectors;
 
 import org.keycloak.OAuth2Constants;
 import org.keycloak.TokenVerifier;
@@ -43,17 +41,13 @@ import org.keycloak.models.GroupModel;
 import org.keycloak.models.GroupModel.Type;
 import org.keycloak.models.IdentityProviderModel;
 import org.keycloak.models.KeycloakSession;
-import org.keycloak.models.OrganizationDomainModel;
 import org.keycloak.models.OrganizationModel;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.UserModel;
 import org.keycloak.organization.OrganizationProvider;
 import org.keycloak.organization.protocol.mappers.oidc.OrganizationScope;
-import org.keycloak.representations.idm.OrganizationDomainRepresentation;
-import org.keycloak.representations.idm.OrganizationRepresentation;
 import org.keycloak.services.ErrorResponse;
 import org.keycloak.sessions.AuthenticationSessionModel;
-import org.keycloak.utils.StringUtil;
 
 public class Organizations {
 
@@ -185,6 +179,10 @@ public class Organizations {
     }
 
     public static OrganizationModel resolveOrganization(KeycloakSession session, UserModel user, String domain) {
+        if (!session.getContext().getRealm().isOrganizationsEnabled()) {
+            return null;
+        }
+
         Optional<OrganizationModel> organization = Optional.ofNullable(session.getContext().getOrganization());
 
         if (organization.isPresent()) {