diff --git a/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/organization/CachedOrganization.java b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/organization/CachedOrganization.java index 000c6e9484..13a603aff6 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/organization/CachedOrganization.java +++ b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/organization/CachedOrganization.java @@ -25,7 +25,6 @@ import org.keycloak.models.IdentityProviderModel; import org.keycloak.models.OrganizationDomainModel; import org.keycloak.models.OrganizationModel; import org.keycloak.models.RealmModel; -import org.keycloak.models.UserModel; import org.keycloak.models.cache.infinispan.DefaultLazyLoader; import org.keycloak.models.cache.infinispan.LazyLoader; import org.keycloak.models.cache.infinispan.entities.AbstractRevisioned; diff --git a/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/organization/CachedOrganizationCount.java b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/organization/CachedOrganizationCount.java new file mode 100644 index 0000000000..8d029369aa --- /dev/null +++ b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/organization/CachedOrganizationCount.java @@ -0,0 +1,42 @@ +/* + * Copyright 2024 Red Hat, Inc. and/or its affiliates + * and other contributors as indicated by the @author tags. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.keycloak.models.cache.infinispan.organization; + +import org.keycloak.models.RealmModel; +import org.keycloak.models.cache.infinispan.entities.AbstractRevisioned; +import org.keycloak.models.cache.infinispan.entities.InRealm; + +public class CachedOrganizationCount extends AbstractRevisioned implements InRealm { + + private final RealmModel realm; + private final long count; + + public CachedOrganizationCount(Long revision, RealmModel realm, long count) { + super(revision, InfinispanOrganizationProvider.cacheKeyOrgCount(realm)); + this.realm = realm; + this.count = count; + } + + @Override + public String getRealm() { + return realm.getId(); + } + + public long getCount() { + return count; + } +} diff --git a/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/organization/InfinispanOrganizationProvider.java b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/organization/InfinispanOrganizationProvider.java index 74b7485897..4f08231b31 100644 --- a/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/organization/InfinispanOrganizationProvider.java +++ b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/organization/InfinispanOrganizationProvider.java @@ -41,11 +41,23 @@ public class InfinispanOrganizationProvider implements OrganizationProvider { this.realmCache = (RealmCacheSession) session.getProvider(CacheRealmProvider.class); } + static String cacheKeyOrgCount(RealmModel realm) { + return realm.getId() + ".org.count"; + } + @Override public OrganizationModel create(String name) { + registerCountInvalidation(); return orgDelegate.create(name); } + @Override + public boolean remove(OrganizationModel organization) { + registerOrganizationInvalidation(organization.getId()); + registerCountInvalidation(); + return orgDelegate.remove(organization); + } + @Override public OrganizationModel getById(String id) { CachedOrganization cached = realmCache.getCache().get(id, CachedOrganization.class); @@ -113,12 +125,6 @@ public class InfinispanOrganizationProvider implements OrganizationProvider { return getCacheDelegates(orgDelegate.getAllStream(attributes, first, max)); } - @Override - public boolean remove(OrganizationModel organization) { - registerOrganizationInvalidation(organization.getId()); - return orgDelegate.remove(organization); - } - @Override public void removeAll() { //TODO: won't scale, requires a better mechanism for bulk deleting organizations within a realm @@ -187,7 +193,20 @@ public class InfinispanOrganizationProvider implements OrganizationProvider { @Override public long count() { - return orgDelegate.count(); + String cacheKey = cacheKeyOrgCount(getRealm()); + CachedOrganizationCount cached = realmCache.getCache().get(cacheKey, CachedOrganizationCount.class); + + // cached and not invalidated + if (cached != null && !realmCache.getInvalidations().contains(cacheKey)) { + return cached.getCount(); + } + + Long loaded = realmCache.getCache().getCurrentRevision(cacheKey); + long count = orgDelegate.count(); + cached = new CachedOrganizationCount(loaded, getRealm(), count); + realmCache.getCache().addRevisioned(cached, realmCache.getStartupRevision()); + + return count; } @Override @@ -204,6 +223,10 @@ public class InfinispanOrganizationProvider implements OrganizationProvider { realmCache.registerInvalidation(orgId); } + private void registerCountInvalidation() { + realmCache.registerInvalidation(cacheKeyOrgCount(getRealm())); + } + private RealmModel getRealm() { RealmModel realm = session.getContext().getRealm(); if (realm == null) { diff --git a/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/organization/admin/OrganizationTest.java b/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/organization/admin/OrganizationTest.java index 6768d72db9..65cc4b37c5 100755 --- a/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/organization/admin/OrganizationTest.java +++ b/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/organization/admin/OrganizationTest.java @@ -42,10 +42,13 @@ import jakarta.ws.rs.NotFoundException; import jakarta.ws.rs.core.Response; import jakarta.ws.rs.core.Response.Status; import java.io.IOException; +import java.util.LinkedList; +import java.util.stream.IntStream; import org.junit.Test; import org.keycloak.admin.client.resource.OrganizationResource; import org.keycloak.admin.client.resource.RealmResource; import org.keycloak.common.Profile.Feature; +import org.keycloak.models.OrganizationModel; import org.keycloak.models.utils.KeycloakModelUtils; import org.keycloak.organization.OrganizationProvider; import org.keycloak.representations.idm.IdentityProviderRepresentation; @@ -452,13 +455,18 @@ public class OrganizationTest extends AbstractOrganizationTest { @Test public void testCount() { - for (int i = 0; i < 10; i++) { - createOrganization("kc.org." + i); - } + List orgIds = IntStream.range(0, 10) + .mapToObj(i -> createOrganization("kc.org." + i).getId()) + .collect(Collectors.toList()); getTestingClient().server(TEST_REALM_NAME).run((RunOnServer) session -> { OrganizationProvider orgProvider = session.getProvider(OrganizationProvider.class); assertEquals(10, orgProvider.count()); + + OrganizationModel org = orgProvider.getById(orgIds.get(0)); + orgProvider.remove(org); + + assertEquals(9, orgProvider.count()); }); } }