libre.sh/keycloak-scim
2
0
Fork 0
Code Issues 14 Pull requests Releases Packages Activity Actions 6

Optimize caching and use of DB connections when Organisations are enabled

Browse source 
Closes #33353

Signed-off-by: Alexander Schwartz <aschwart@redhat.com>
This commit is contained in:
Alexander Schwartz 2024-09-27 18:05:01 +02:00 committed by Pedro Igor
parent d5d6390b1c
commit 5c503a55e9
3 changed files with 46 additions and 42 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

65
model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/organization/InfinispanOrganizationProvider.java vendored
View file

@ -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) {

13
model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/organization/OrganizationAdapter.java vendored
View file

@ -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

10
services/src/main/java/org/keycloak/organization/utils/Organizations.java
View file

@ -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()) {