cache count

Signed-off-by: vramik <vramik@redhat.com>

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;

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;
+    }
+}

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

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<String> 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());
         });
     }
 }