From 375923ba0f0d5c12e997ca0479ec64e7fab071c9 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 29 Jun 2026 01:06:04 +0000 Subject: [PATCH 01/10] Initial plan From 5ce811366a66d96587969a799b771ab2f94b6fe6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=BCseyin=20Akdo=C4=9Fan?= Date: Wed, 24 Jun 2026 17:03:06 +0300 Subject: [PATCH 02/10] Drop redundant isComposite() pre-query in composite-role expansion MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit RoleUtils.expandCompositeRolesStream gated each node on isComposite() before calling getCompositesStream(). On the JPA (cache-miss) path both run the getChildRoles named query -- and a FlushMode.AUTO flush before each -- so every composite node issued two getChildRoles round-trips while expanding the tree. getCompositesStream() already yields an empty stream for non-composite roles, making the pre-check redundant. Fetch the composites directly instead. Behaviour is unchanged (verified by RoleUtilsCompositeExpansionTest: identical expansion, isComposite() no longer called, cycles still terminate). On a 10-node composite chain the deterministic query-count model test (CompositeRoleExpansionQueryCountTest, jpa profile) shows getChildRoles executions drop from 19 to 10 -- the saving equals the number of composite nodes (C) per expansion. Signed-off-by: Hüseyin Akdoğan --- .../org/keycloak/models/utils/RoleUtils.java | 20 +- .../RoleUtilsCompositeExpansionTest.java | 182 ++++++++++++++++++ .../CompositeRoleExpansionQueryCountTest.java | 122 ++++++++++++ 3 files changed, 316 insertions(+), 8 deletions(-) create mode 100644 server-spi/src/test/java/org/keycloak/models/utils/RoleUtilsCompositeExpansionTest.java create mode 100644 testsuite/model/src/test/java/org/keycloak/testsuite/model/role/CompositeRoleExpansionQueryCountTest.java diff --git a/server-spi/src/main/java/org/keycloak/models/utils/RoleUtils.java b/server-spi/src/main/java/org/keycloak/models/utils/RoleUtils.java index 0322bac39e1c..f4ce086485ae 100644 --- a/server-spi/src/main/java/org/keycloak/models/utils/RoleUtils.java +++ b/server-spi/src/main/java/org/keycloak/models/utils/RoleUtils.java @@ -146,14 +146,18 @@ private static Stream expandCompositeRolesStream(RoleModel role, Set< RoleModel current = stack.pop(); sb.add(current); - if (current.isComposite()) { - current.getCompositesStream() - .filter(r -> !visited.contains(r)) - .forEach(r -> { - visited.add(r); - stack.add(r); - }); - } + // Fetch the composites directly rather than gating on isComposite() first. + // On the JPA (cache-miss) path isComposite() runs its own getChildRoles query -- + // and the FlushMode.AUTO flush before it -- immediately before getCompositesStream() + // runs another, doubling the per-node round-trips while expanding the tree. + // getCompositesStream() already yields an empty stream for non-composite roles, + // so the pre-check is redundant. + current.getCompositesStream() + .filter(r -> !visited.contains(r)) + .forEach(r -> { + visited.add(r); + stack.add(r); + }); } } diff --git a/server-spi/src/test/java/org/keycloak/models/utils/RoleUtilsCompositeExpansionTest.java b/server-spi/src/test/java/org/keycloak/models/utils/RoleUtilsCompositeExpansionTest.java new file mode 100644 index 000000000000..2373b977fd25 --- /dev/null +++ b/server-spi/src/test/java/org/keycloak/models/utils/RoleUtilsCompositeExpansionTest.java @@ -0,0 +1,182 @@ +package org.keycloak.models.utils; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Stream; + +import org.junit.Assert; +import org.junit.Test; +import org.keycloak.models.RoleContainerModel; +import org.keycloak.models.RoleModel; + +public class RoleUtilsCompositeExpansionTest { + + @Test + public void expandsFullTreeWithoutCallingIsComposite() { + // root -> a, b ; a -> c ; b -> c (diamond, c shared) ; c is a leaf + FakeRole c = new FakeRole("c"); + FakeRole a = new FakeRole("a", c); + FakeRole b = new FakeRole("b", c); + FakeRole root = new FakeRole("root", a, b); + + Set expanded = RoleUtils.expandCompositeRoles(Set.of(root)); + + Assert.assertEquals(Set.of(root, a, b, c), expanded); + + // The optimization: isComposite() is never used to drive the expansion. + for (FakeRole r : Arrays.asList(root, a, b, c)) { + Assert.assertEquals("isComposite() must not be called for " + r.id, 0, r.isCompositeCalls); + } + // Each reachable role is resolved exactly once thanks to the visited set. + Assert.assertEquals(1, root.getCompositesCalls); + Assert.assertEquals(1, a.getCompositesCalls); + Assert.assertEquals(1, b.getCompositesCalls); + Assert.assertEquals(1, c.getCompositesCalls); + } + + @Test + public void leafRoleExpandsToItself() { + FakeRole leaf = new FakeRole("leaf"); + + Assert.assertEquals(Set.of(leaf), RoleUtils.expandCompositeRoles(Set.of(leaf))); + Assert.assertEquals(0, leaf.isCompositeCalls); + } + + @Test + public void cyclicCompositesTerminate() { + // root -> a -> root (cycle) + FakeRole a = new FakeRole("a"); + FakeRole root = new FakeRole("root", a); + a.composites.add(root); + + Set expanded = RoleUtils.expandCompositeRoles(Set.of(root)); + + Assert.assertEquals(Set.of(root, a), expanded); + } + + // RoleModel stub recording isComposite()/getCompositesStream() calls; identity by id. + private static final class FakeRole implements RoleModel { + + final String id; + final List composites = new ArrayList<>(); + int isCompositeCalls = 0; + int getCompositesCalls = 0; + + FakeRole(String id, RoleModel... children) { + this.id = id; + this.composites.addAll(Arrays.asList(children)); + } + + @Override + public boolean isComposite() { + isCompositeCalls++; + return !composites.isEmpty(); + } + + @Override + public Stream getCompositesStream() { + getCompositesCalls++; + return new ArrayList<>(composites).stream(); + } + + @Override + public String getId() { + return id; + } + + @Override + public boolean equals(Object o) { + return o instanceof RoleModel && ((RoleModel) o).getId().equals(id); + } + + @Override + public int hashCode() { + return id.hashCode(); + } + + // --- unused RoleModel surface --------------------------------------------------------- + + @Override + public String getName() { + return id; + } + + @Override + public String getDescription() { + throw new UnsupportedOperationException(); + } + + @Override + public void setDescription(String description) { + throw new UnsupportedOperationException(); + } + + @Override + public void setName(String name) { + throw new UnsupportedOperationException(); + } + + @Override + public void addCompositeRole(RoleModel role) { + throw new UnsupportedOperationException(); + } + + @Override + public void removeCompositeRole(RoleModel role) { + throw new UnsupportedOperationException(); + } + + @Override + public Stream getCompositesStream(String search, Integer first, Integer max) { + throw new UnsupportedOperationException(); + } + + @Override + public boolean isClientRole() { + throw new UnsupportedOperationException(); + } + + @Override + public String getContainerId() { + throw new UnsupportedOperationException(); + } + + @Override + public RoleContainerModel getContainer() { + throw new UnsupportedOperationException(); + } + + @Override + public boolean hasRole(RoleModel role) { + throw new UnsupportedOperationException(); + } + + @Override + public void setSingleAttribute(String name, String value) { + throw new UnsupportedOperationException(); + } + + @Override + public void setAttribute(String name, List values) { + throw new UnsupportedOperationException(); + } + + @Override + public void removeAttribute(String name) { + throw new UnsupportedOperationException(); + } + + @Override + public Stream getAttributeStream(String name) { + throw new UnsupportedOperationException(); + } + + @Override + public Map> getAttributes() { + throw new UnsupportedOperationException(); + } + } +} diff --git a/testsuite/model/src/test/java/org/keycloak/testsuite/model/role/CompositeRoleExpansionQueryCountTest.java b/testsuite/model/src/test/java/org/keycloak/testsuite/model/role/CompositeRoleExpansionQueryCountTest.java new file mode 100644 index 000000000000..79568a161e77 --- /dev/null +++ b/testsuite/model/src/test/java/org/keycloak/testsuite/model/role/CompositeRoleExpansionQueryCountTest.java @@ -0,0 +1,122 @@ +package org.keycloak.testsuite.model.role; + +import java.util.ArrayDeque; +import java.util.Deque; +import java.util.HashSet; +import java.util.Set; + +import org.keycloak.connections.jpa.JpaConnectionProvider; +import org.keycloak.models.Constants; +import org.keycloak.models.KeycloakSession; +import org.keycloak.models.RealmModel; +import org.keycloak.models.RealmProvider; +import org.keycloak.models.RoleModel; +import org.keycloak.models.RoleProvider; +import org.keycloak.models.utils.RoleUtils; +import org.keycloak.testsuite.model.KeycloakModelTest; +import org.keycloak.testsuite.model.RequireProvider; + +import org.hibernate.SessionFactory; +import org.hibernate.stat.Statistics; +import org.junit.Test; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.is; + +@RequireProvider(RealmProvider.class) +@RequireProvider(RoleProvider.class) +public class CompositeRoleExpansionQueryCountTest extends KeycloakModelTest { + + // Deep composite chain r0 -> r1 -> ... -> r9 (r9 is a leaf): V = 10 nodes, C = 9 composite nodes. + private static final int CHAIN_LENGTH = 10; + + private String realmId; + private String rootRoleId; + + @Override + public void createEnvironment(KeycloakSession s) { + RealmModel realm = createRealm(s, "composite-expansion-realm"); + s.getContext().setRealm(realm); + realm.setDefaultRole(s.roles().addRealmRole(realm, Constants.DEFAULT_ROLES_ROLE_PREFIX + "-" + realm.getName())); + realmId = realm.getId(); + + RoleModel previous = null; + for (int i = 0; i < CHAIN_LENGTH; i++) { + RoleModel role = s.roles().addRealmRole(realm, "chain-" + i); + if (previous == null) { + rootRoleId = role.getId(); + } else { + previous.addCompositeRole(role); + } + previous = role; + } + } + + @Override + public void cleanEnvironment(KeycloakSession s) { + RealmModel realm = s.realms().getRealm(realmId); + s.getContext().setRealm(realm); + s.realms().removeRealm(realmId); + } + + @Test + public void droppingIsCompositePreQueryHalvesGetChildRolesOnCompositeNodes() { + long currentCount = withRealm(realmId, (s, realm) -> { + RoleModel root = s.roles().getRoleById(realm, rootRoleId); + return countQueries(s, () -> RoleUtils.expandCompositeRoles(Set.of(root))); + }); + + long prePatchCount = withRealm(realmId, (s, realm) -> { + RoleModel root = s.roles().getRoleById(realm, rootRoleId); + return countQueries(s, () -> expandWithIsCompositePreQuery(Set.of(root))); + }); + + // Current code: one getChildRoles per visited node (getCompositesStream on every node). + assertThat(currentCount, is((long) CHAIN_LENGTH)); // V = 10 + // Pre-patch code: isComposite() on every node + getCompositesStream() on composite nodes. + assertThat(prePatchCount, is((long) (2 * CHAIN_LENGTH - 1))); // V + C = 19 + // Saved exactly C = number of composite nodes. + assertThat(prePatchCount - currentCount, is((long) (CHAIN_LENGTH - 1))); // 9 + } + + private long countQueries(KeycloakSession session, Runnable walk) { + Statistics stats = session.getProvider(JpaConnectionProvider.class) + .getEntityManager() + .getEntityManagerFactory() + .unwrap(SessionFactory.class) + .getStatistics(); + stats.setStatisticsEnabled(true); + stats.clear(); + walk.run(); + // The only JDBC statements during the walk are the getChildRoles named-query executions: + // getCompositesStream() hydrates the child RoleEntity rows, so the subsequent getRoleById + // em.find calls are served from the persistence context without extra statements. + return stats.getPrepareStatementCount(); + } + + // Pre-patch walk: keeps the isComposite() guard so both variants run against the same database. + private static Set expandWithIsCompositePreQuery(Set roles) { + Set visited = new HashSet<>(); + Set result = new HashSet<>(); + for (RoleModel role : roles) { + if (visited.contains(role)) { + continue; + } + Deque stack = new ArrayDeque<>(); + stack.add(role); + while (!stack.isEmpty()) { + RoleModel current = stack.pop(); + result.add(current); + if (current.isComposite()) { + current.getCompositesStream() + .filter(r -> !visited.contains(r)) + .forEach(r -> { + visited.add(r); + stack.add(r); + }); + } + } + } + return result; + } +} From 7f9684be13fc37fa1267d20d4e0d6296777094c2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=BCseyin=20Akdo=C4=9Fan?= Date: Wed, 24 Jun 2026 17:46:27 +0300 Subject: [PATCH 03/10] Batch composite-role expansion with a level-wise IN query MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The composite-role walk in RoleUtils issues one getChildRoles query per role node. On a cache miss this is an N+1 over the whole role tree, each query preceded by a FlushMode.AUTO flush. Add a bulk primitive RoleProvider.getCompositeRolesStream(realm, parentIds) -- default resolves each parent individually, JpaRealmProvider overrides it with a single getChildRolesFromParentIds named query (IN clause) that hydrates the child RoleEntity rows so the follow-up getRoleById calls are served from the persistence context. RoleStorageManager delegates to local storage. RoleUtils.expandCompositeRoles(session, realm, roles) uses it to expand the tree one breadth-first level at a time: one query per level instead of one per node, while preserving the visited-set cycle protection and dedup. CompositeRoleExpansionQueryCountTest (jpa profile) proves it on a 31-node wide tree (root -> 5 -> 25): the per-node walk issues 31 getChildRoles queries, level-wise batching issues 3 (one per depth), and both produce the same expanded role set. Signed-off-by: Hüseyin Akdoğan --- .../keycloak/models/jpa/JpaRealmProvider.java | 16 +++++++ .../models/jpa/entities/RoleEntity.java | 1 + .../keycloak/storage/RoleStorageManager.java | 6 +++ .../org/keycloak/models/RoleProvider.java | 25 +++++++++++ .../org/keycloak/models/utils/RoleUtils.java | 38 ++++++++++++++++ .../CompositeRoleExpansionQueryCountTest.java | 43 +++++++++++++++++++ 6 files changed, 129 insertions(+) diff --git a/model/jpa/src/main/java/org/keycloak/models/jpa/JpaRealmProvider.java b/model/jpa/src/main/java/org/keycloak/models/jpa/JpaRealmProvider.java index 4b41f00446cb..8105dbe51e72 100644 --- a/model/jpa/src/main/java/org/keycloak/models/jpa/JpaRealmProvider.java +++ b/model/jpa/src/main/java/org/keycloak/models/jpa/JpaRealmProvider.java @@ -505,6 +505,22 @@ public RoleModel getRoleById(RealmModel realm, String id) { return adapter; } + @Override + public Stream getCompositeRolesStream(RealmModel realm, Set parentRoleIds) { + if (parentRoleIds == null || parentRoleIds.isEmpty()) { + return Stream.empty(); + } + // One query fetches the child roles for every parent in the frontier, so a composite-role + // tree expands with a single getChildRolesFromParentIds query per breadth-first level. The + // query hydrates the child RoleEntity rows, so the getRoleById calls below are served from + // the persistence context without extra round-trips. + TypedQuery query = em.createNamedQuery("getChildRolesFromParentIds", RoleEntity.class) + .setParameter("parentRoleIds", parentRoleIds); + return closing(query.getResultStream()) + .map(roleEntity -> session.roles().getRoleById(realm, roleEntity.getId())) + .filter(Objects::nonNull); + } + @Override public GroupModel getGroupById(RealmModel realm, String id) { GroupEntity groupEntity = em.find(GroupEntity.class, id); diff --git a/model/jpa/src/main/java/org/keycloak/models/jpa/entities/RoleEntity.java b/model/jpa/src/main/java/org/keycloak/models/jpa/entities/RoleEntity.java index 2b2921242514..b3c6f8196e37 100755 --- a/model/jpa/src/main/java/org/keycloak/models/jpa/entities/RoleEntity.java +++ b/model/jpa/src/main/java/org/keycloak/models/jpa/entities/RoleEntity.java @@ -61,6 +61,7 @@ @NamedQuery(name="getRoleIdsFromIdList", query="select role.id from RoleEntity role where role.realmId = :realm and role.id in :ids order by role.name ASC"), @NamedQuery(name="getRoleIdsByNameContainingFromIdList", query="select role.id from RoleEntity role where role.realmId = :realm and lower(role.name) like lower(concat('%',:search,'%')) and role.id in :ids order by role.name ASC"), @NamedQuery(name="getChildRoles", query="select r from RoleEntity r join CompositeRoleEntity c on r.id = c.childRole.id where c.parentRole.id = :parentRoleId"), + @NamedQuery(name="getChildRolesFromParentIds", query="select distinct r from RoleEntity r join CompositeRoleEntity c on r.id = c.childRole.id where c.parentRole.id in :parentRoleIds"), }) public class RoleEntity { diff --git a/model/storage-private/src/main/java/org/keycloak/storage/RoleStorageManager.java b/model/storage-private/src/main/java/org/keycloak/storage/RoleStorageManager.java index 612189effb8b..06dd40474735 100644 --- a/model/storage-private/src/main/java/org/keycloak/storage/RoleStorageManager.java +++ b/model/storage-private/src/main/java/org/keycloak/storage/RoleStorageManager.java @@ -17,6 +17,7 @@ package org.keycloak.storage; import java.util.Objects; +import java.util.Set; import java.util.stream.Stream; import org.keycloak.common.util.reflections.Types; @@ -161,6 +162,11 @@ public Stream getRolesStream(RealmModel realm, Stream ids, St return localStorage().getRolesStream(realm, ids, search, first, max); } + @Override + public Stream getCompositeRolesStream(RealmModel realm, Set parentRoleIds) { + return localStorage().getCompositeRolesStream(realm, parentRoleIds); + } + /** * Obtaining roles from an external role storage is time-bounded. In case the external role storage * isn't available at least roles from a local storage are returned. For this purpose diff --git a/server-spi/src/main/java/org/keycloak/models/RoleProvider.java b/server-spi/src/main/java/org/keycloak/models/RoleProvider.java index 39bd8fb3cabc..ac6b0bbce5f0 100644 --- a/server-spi/src/main/java/org/keycloak/models/RoleProvider.java +++ b/server-spi/src/main/java/org/keycloak/models/RoleProvider.java @@ -16,6 +16,8 @@ */ package org.keycloak.models; +import java.util.Objects; +import java.util.Set; import java.util.stream.Stream; import org.keycloak.provider.Provider; @@ -80,6 +82,29 @@ default Stream getRealmRolesStream(RealmModel realm) { */ Stream getRolesStream(RealmModel realm, Stream ids, String search, Integer first, Integer max); + /** + * Returns the direct composite child roles of all the given parent roles, de-duplicated. + *

+ * This is the bulk counterpart of {@link RoleModel#getCompositesStream()}: it lets callers expand + * a composite-role tree one breadth-first level at a time with a single lookup per level instead + * of one per role. The default implementation resolves each parent individually; storage-backed + * providers should override it with a batched lookup. + * + * @param realm Realm. Cannot be {@code null}. + * @param parentRoleIds Ids of the roles whose direct composite children should be returned. + * @return Stream of the child roles. Never returns {@code null}. + */ + default Stream getCompositeRolesStream(RealmModel realm, Set parentRoleIds) { + if (parentRoleIds == null || parentRoleIds.isEmpty()) { + return Stream.empty(); + } + return parentRoleIds.stream() + .map(id -> getRoleById(realm, id)) + .filter(Objects::nonNull) + .flatMap(RoleModel::getCompositesStream) + .distinct(); + } + /** * Removes given realm role from the given realm. * @param role Role to be removed. diff --git a/server-spi/src/main/java/org/keycloak/models/utils/RoleUtils.java b/server-spi/src/main/java/org/keycloak/models/utils/RoleUtils.java index f4ce086485ae..17587e76f255 100644 --- a/server-spi/src/main/java/org/keycloak/models/utils/RoleUtils.java +++ b/server-spi/src/main/java/org/keycloak/models/utils/RoleUtils.java @@ -27,6 +27,7 @@ import org.keycloak.models.ClientModel; import org.keycloak.models.GroupModel; +import org.keycloak.models.KeycloakSession; import org.keycloak.models.RealmModel; import org.keycloak.models.RoleContainerModel; import org.keycloak.models.RoleMapperModel; @@ -177,6 +178,43 @@ public static Set expandCompositeRoles(Set roles) { .collect(Collectors.toSet()); } + /** + * Bulk variant of {@link #expandCompositeRoles(Set)}. Expands the composite-role tree one + * breadth-first level at a time using {@code session.roles().getCompositeRolesStream}, issuing a + * single child lookup per level instead of one per role. Storage-backed providers batch that + * lookup, so this avoids the per-node N+1 on a cache miss. + * + * @param session Keycloak session. + * @param realm Realm the roles belong to. + * @param roles Roles to expand. + * @return new set with composite roles expanded + */ + public static Set expandCompositeRoles(KeycloakSession session, RealmModel realm, Set roles) { + Set result = new HashSet<>(); + Set enqueued = new HashSet<>(); + Set frontier = roles; + + while (!frontier.isEmpty()) { + result.addAll(frontier); + + Set parentIds = new HashSet<>(); + frontier.forEach(role -> { + if (enqueued.add(role.getId())) { + parentIds.add(role.getId()); + } + }); + if (parentIds.isEmpty()) { + break; + } + + frontier = session.roles().getCompositeRolesStream(realm, parentIds) + .filter(role -> !enqueued.contains(role.getId())) + .collect(Collectors.toSet()); + } + + return result; + } + /** * @param roles * @return stream with composite roles expanded diff --git a/testsuite/model/src/test/java/org/keycloak/testsuite/model/role/CompositeRoleExpansionQueryCountTest.java b/testsuite/model/src/test/java/org/keycloak/testsuite/model/role/CompositeRoleExpansionQueryCountTest.java index 79568a161e77..26e6191f574e 100644 --- a/testsuite/model/src/test/java/org/keycloak/testsuite/model/role/CompositeRoleExpansionQueryCountTest.java +++ b/testsuite/model/src/test/java/org/keycloak/testsuite/model/role/CompositeRoleExpansionQueryCountTest.java @@ -30,8 +30,14 @@ public class CompositeRoleExpansionQueryCountTest extends KeycloakModelTest { // Deep composite chain r0 -> r1 -> ... -> r9 (r9 is a leaf): V = 10 nodes, C = 9 composite nodes. private static final int CHAIN_LENGTH = 10; + // Wide tree: root -> BRANCHING mids -> BRANCHING leaves each. V = 1 + 5 + 25 = 31, depth = 3 levels. + private static final int BRANCHING = 5; + private static final int BUSHY_SIZE = 1 + BRANCHING + BRANCHING * BRANCHING; + private static final int BUSHY_LEVELS = 3; + private String realmId; private String rootRoleId; + private String bushyRootId; @Override public void createEnvironment(KeycloakSession s) { @@ -50,6 +56,16 @@ public void createEnvironment(KeycloakSession s) { } previous = role; } + + RoleModel bushyRoot = s.roles().addRealmRole(realm, "bushy-root"); + bushyRootId = bushyRoot.getId(); + for (int i = 0; i < BRANCHING; i++) { + RoleModel mid = s.roles().addRealmRole(realm, "bushy-mid-" + i); + bushyRoot.addCompositeRole(mid); + for (int j = 0; j < BRANCHING; j++) { + mid.addCompositeRole(s.roles().addRealmRole(realm, "bushy-leaf-" + i + "-" + j)); + } + } } @Override @@ -79,6 +95,33 @@ public void droppingIsCompositePreQueryHalvesGetChildRolesOnCompositeNodes() { assertThat(prePatchCount - currentCount, is((long) (CHAIN_LENGTH - 1))); // 9 } + @Test + public void levelWiseBatchingCollapsesGetChildRolesToOnePerLevel() { + Set perNodeIds = new HashSet<>(); + long perNodeCount = withRealm(realmId, (s, realm) -> { + RoleModel root = s.roles().getRoleById(realm, bushyRootId); + return countQueries(s, () -> collectIds(RoleUtils.expandCompositeRoles(Set.of(root)), perNodeIds)); + }); + + Set batchedIds = new HashSet<>(); + long batchedCount = withRealm(realmId, (s, realm) -> { + RoleModel root = s.roles().getRoleById(realm, bushyRootId); + return countQueries(s, () -> collectIds(RoleUtils.expandCompositeRoles(s, realm, Set.of(root)), batchedIds)); + }); + + // Both expansions yield the same role set ... + assertThat(perNodeIds.size(), is(BUSHY_SIZE)); // 31 + assertThat(batchedIds, is(perNodeIds)); + // ... but the per-node walk issues one getChildRoles query per node ... + assertThat(perNodeCount, is((long) BUSHY_SIZE)); // 31 + // ... whereas level-wise batching issues a single query per breadth-first level. + assertThat(batchedCount, is((long) BUSHY_LEVELS)); // 3 + } + + private static void collectIds(Set roles, Set into) { + roles.forEach(role -> into.add(role.getId())); + } + private long countQueries(KeycloakSession session, Runnable walk) { Statistics stats = session.getProvider(JpaConnectionProvider.class) .getEntityManager() From 9cc00e10743400b68d76e27591efd786ae460258 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=BCseyin=20Akdo=C4=9Fan?= Date: Wed, 24 Jun 2026 17:57:40 +0300 Subject: [PATCH 04/10] Use batched composite-role expansion on token/permission hot paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Route the composite-role expansion on the token and permission paths through the level-wise batched RoleUtils variants so the N+1 reduction reaches the flows the load test exercises. Add session-aware RoleUtils.getDeepUserRoleMappings(session, realm, user) and expandCompositeRolesStream(session, realm, roles), and switch: - TokenManager.getAccess: add a session-aware overload; the existing (UserModel, ClientModel, Stream) signature is kept and delegates to the per-node expansion, so the public API is unchanged, - DefaultClientSessionContext client-scope role expansion and user roles, - RolePolicyProvider (admin fine-grained permissions) deep user role mappings. The non-session methods are kept for callers without a session. Behaviour is unchanged: the batched expansion returns the same role set as the per-node walk (CompositeRoleExpansionQueryCountTest); a custom RoleProvider that does not override getCompositeRolesStream still resolves composites via the interface default (RoleProviderCompositeDefaultTest); RoleModelTest and CompositeRolesModelTest pass. Signed-off-by: Hüseyin Akdoğan --- .../provider/role/RolePolicyProvider.java | 2 +- .../org/keycloak/models/utils/RoleUtils.java | 18 +++ .../RoleProviderCompositeDefaultTest.java | 110 ++++++++++++++++++ .../keycloak/protocol/oidc/TokenManager.java | 14 ++- .../resources/account/AccountConsole.java | 2 +- .../util/DefaultClientSessionContext.java | 6 +- 6 files changed, 145 insertions(+), 7 deletions(-) create mode 100644 server-spi/src/test/java/org/keycloak/models/RoleProviderCompositeDefaultTest.java diff --git a/authz/policy/common/src/main/java/org/keycloak/authorization/policy/provider/role/RolePolicyProvider.java b/authz/policy/common/src/main/java/org/keycloak/authorization/policy/provider/role/RolePolicyProvider.java index 50d63f91f28f..3e9cd062b799 100644 --- a/authz/policy/common/src/main/java/org/keycloak/authorization/policy/provider/role/RolePolicyProvider.java +++ b/authz/policy/common/src/main/java/org/keycloak/authorization/policy/provider/role/RolePolicyProvider.java @@ -141,7 +141,7 @@ public Stream getPermissions(KeycloakSession session, ResourceType resou StoreFactory storeFactory = provider.getStoreFactory(); ResourceServer resourceServer = storeFactory.getResourceServerStore().findByClient(adminPermissionsClient); PolicyStore policyStore = storeFactory.getPolicyStore(); - List roleIds = getDeepUserRoleMappings(subject).stream().map(RoleModel::getId).toList(); + List roleIds = getDeepUserRoleMappings(session, realm, subject).stream().map(RoleModel::getId).toList(); Stream policies = Stream.of(); return Stream.concat(policies, policyStore.findDependentPolicies(resourceServer, resourceType.getType(), groupResourceType == null ? null : groupResourceType.getType(), RolePolicyProviderFactory.ID, "roles", roleIds)); diff --git a/server-spi/src/main/java/org/keycloak/models/utils/RoleUtils.java b/server-spi/src/main/java/org/keycloak/models/utils/RoleUtils.java index 17587e76f255..d56b20750192 100644 --- a/server-spi/src/main/java/org/keycloak/models/utils/RoleUtils.java +++ b/server-spi/src/main/java/org/keycloak/models/utils/RoleUtils.java @@ -225,6 +225,14 @@ public static Stream expandCompositeRolesStream(Stream rol return roles.flatMap(roleModel -> RoleUtils.expandCompositeRolesStream(roleModel, visited)); } + /** + * Bulk variant of {@link #expandCompositeRolesStream(Stream)} backed by + * {@link #expandCompositeRoles(KeycloakSession, RealmModel, Set)}. + */ + public static Stream expandCompositeRolesStream(KeycloakSession session, RealmModel realm, Stream roles) { + return expandCompositeRoles(session, realm, roles.collect(Collectors.toSet())).stream(); + } + /** * @param roleMapper @@ -254,6 +262,16 @@ public static Set getDeepUserRoleMappings(UserModel user) { return expandCompositeRoles(roleMappings); } + /** + * Bulk variant of {@link #getDeepUserRoleMappings(UserModel)} that expands composite roles with + * a single child lookup per breadth-first level instead of one per role on a cache miss. + */ + public static Set getDeepUserRoleMappings(KeycloakSession session, RealmModel realm, UserModel user) { + Set roleMappings = user.getRoleMappingsStream().collect(Collectors.toSet()); + user.getGroupsStream().forEach(group -> addGroupRoles(group, roleMappings)); + return expandCompositeRoles(session, realm, roleMappings); + } + private static void addGroupRoles(GroupModel group, Set roleMappings) { roleMappings.addAll(group.getRoleMappingsStream().collect(Collectors.toSet())); diff --git a/server-spi/src/test/java/org/keycloak/models/RoleProviderCompositeDefaultTest.java b/server-spi/src/test/java/org/keycloak/models/RoleProviderCompositeDefaultTest.java new file mode 100644 index 000000000000..a82507333012 --- /dev/null +++ b/server-spi/src/test/java/org/keycloak/models/RoleProviderCompositeDefaultTest.java @@ -0,0 +1,110 @@ +package org.keycloak.models; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +import org.junit.Assert; +import org.junit.Test; + +// A custom RoleProvider that does NOT override getCompositeRolesStream must still resolve composite +// children correctly through the interface default implementation (de-duplicated, per parent). +public class RoleProviderCompositeDefaultTest { + + @Test + public void defaultGetCompositeRolesStreamResolvesAndDeduplicates() { + // root -> a, b ; a -> c ; b -> c (diamond, c shared) ; c is a leaf + FakeRole c = new FakeRole("c"); + FakeRole a = new FakeRole("a", c); + FakeRole b = new FakeRole("b", c); + FakeRole root = new FakeRole("root", a, b); + FakeRoleProvider provider = new FakeRoleProvider(root, a, b, c); + + Assert.assertEquals(Set.of("a", "b"), childIds(provider, "root")); + // c is shared by a and b but must appear once + Assert.assertEquals(Set.of("c"), childIds(provider, "a", "b")); + Assert.assertEquals(Set.of(), childIds(provider, "c")); + Assert.assertEquals(Set.of(), provider.getCompositeRolesStream(null, Set.of()).collect(Collectors.toSet())); + } + + private static Set childIds(FakeRoleProvider provider, String... parentIds) { + return provider.getCompositeRolesStream(null, Set.of(parentIds)) + .map(RoleModel::getId).collect(Collectors.toSet()); + } + + // RoleProvider stub backed by an in-memory role map; getCompositeRolesStream is left to the default. + private static final class FakeRoleProvider implements RoleProvider { + + private final Map byId = new HashMap<>(); + + FakeRoleProvider(RoleModel... roles) { + for (RoleModel r : roles) byId.put(r.getId(), r); + } + + @Override + public RoleModel getRoleById(RealmModel realm, String id) { + return byId.get(id); + } + + // --- unused RoleProvider / RoleLookupProvider surface -------------------------------- + + @Override public RoleModel addRealmRole(RealmModel realm, String id, String name) { throw new UnsupportedOperationException(); } + @Override public Stream getRealmRolesStream(RealmModel realm, Integer first, Integer max) { throw new UnsupportedOperationException(); } + @Override public Stream getRolesStream(RealmModel realm, Stream ids, String search, Integer first, Integer max) { throw new UnsupportedOperationException(); } + @Override public boolean removeRole(RoleModel role) { throw new UnsupportedOperationException(); } + @Override public void removeRoles(RealmModel realm) { throw new UnsupportedOperationException(); } + @Override public RoleModel addClientRole(ClientModel client, String id, String name) { throw new UnsupportedOperationException(); } + @Override public Stream getClientRolesStream(ClientModel client, Integer first, Integer max) { throw new UnsupportedOperationException(); } + @Override public void removeRoles(ClientModel client) { throw new UnsupportedOperationException(); } + @Override public RoleModel getRealmRole(RealmModel realm, String name) { throw new UnsupportedOperationException(); } + @Override public Stream searchForRolesStream(RealmModel realm, String search, Integer first, Integer max) { throw new UnsupportedOperationException(); } + @Override public RoleModel getClientRole(ClientModel client, String name) { throw new UnsupportedOperationException(); } + @Override public Stream searchForClientRolesStream(ClientModel client, String search, Integer first, Integer max) { throw new UnsupportedOperationException(); } + @Override public Stream searchForClientRolesStream(RealmModel realm, Stream ids, String search, Integer first, Integer max) { throw new UnsupportedOperationException(); } + @Override public Stream searchForClientRolesStream(RealmModel realm, String search, Stream excludedIds, Integer first, Integer max) { throw new UnsupportedOperationException(); } + @Override public void close() { } + } + + // RoleModel stub; identity by id, composites held in memory. + private static final class FakeRole implements RoleModel { + + private final String id; + private final List composites = new ArrayList<>(); + + FakeRole(String id, RoleModel... children) { + this.id = id; + this.composites.addAll(Arrays.asList(children)); + } + + @Override public String getId() { return id; } + @Override public boolean isComposite() { return !composites.isEmpty(); } + @Override public Stream getCompositesStream() { return new ArrayList<>(composites).stream(); } + + @Override public boolean equals(Object o) { return o instanceof RoleModel && ((RoleModel) o).getId().equals(id); } + @Override public int hashCode() { return id.hashCode(); } + + // --- unused RoleModel surface -------------------------------------------------------- + + @Override public String getName() { return id; } + @Override public String getDescription() { throw new UnsupportedOperationException(); } + @Override public void setDescription(String description) { throw new UnsupportedOperationException(); } + @Override public void setName(String name) { throw new UnsupportedOperationException(); } + @Override public void addCompositeRole(RoleModel role) { throw new UnsupportedOperationException(); } + @Override public void removeCompositeRole(RoleModel role) { throw new UnsupportedOperationException(); } + @Override public Stream getCompositesStream(String search, Integer first, Integer max) { throw new UnsupportedOperationException(); } + @Override public boolean isClientRole() { throw new UnsupportedOperationException(); } + @Override public String getContainerId() { throw new UnsupportedOperationException(); } + @Override public RoleContainerModel getContainer() { throw new UnsupportedOperationException(); } + @Override public boolean hasRole(RoleModel role) { throw new UnsupportedOperationException(); } + @Override public void setSingleAttribute(String name, String value) { throw new UnsupportedOperationException(); } + @Override public void setAttribute(String name, List values) { throw new UnsupportedOperationException(); } + @Override public void removeAttribute(String name) { throw new UnsupportedOperationException(); } + @Override public Stream getAttributeStream(String name) { throw new UnsupportedOperationException(); } + @Override public Map> getAttributes() { throw new UnsupportedOperationException(); } + } +} diff --git a/services/src/main/java/org/keycloak/protocol/oidc/TokenManager.java b/services/src/main/java/org/keycloak/protocol/oidc/TokenManager.java index 8d51044062c9..c547a27d3756 100755 --- a/services/src/main/java/org/keycloak/protocol/oidc/TokenManager.java +++ b/services/src/main/java/org/keycloak/protocol/oidc/TokenManager.java @@ -473,8 +473,18 @@ public static void dettachClientSession(AuthenticatedClientSessionModel clientSe public static Set getAccess(UserModel user, ClientModel client, Stream clientScopes) { - Set roleMappings = RoleUtils.getDeepUserRoleMappings(user); + return getAccess(RoleUtils.getDeepUserRoleMappings(user), client, clientScopes, + RoleUtils::expandCompositeRolesStream); + } + + public static Set getAccess(KeycloakSession session, UserModel user, ClientModel client, Stream clientScopes) { + RealmModel realm = client.getRealm(); + return getAccess(RoleUtils.getDeepUserRoleMappings(session, realm, user), client, clientScopes, + scopeMappings -> RoleUtils.expandCompositeRolesStream(session, realm, scopeMappings)); + } + private static Set getAccess(Set roleMappings, ClientModel client, Stream clientScopes, + Function, Stream> expandScopeMappings) { if (client.isFullScopeAllowed()) { if (logger.isTraceEnabled()) { logger.tracef("Using full scope for client %s", client.getClientId()); @@ -499,7 +509,7 @@ public static Set getAccess(UserModel user, ClientModel client, Strea scopeMappings = Stream.concat(scopeMappings, clientScopesMappings); // 3 - Expand scope mappings - scopeMappings = RoleUtils.expandCompositeRolesStream(scopeMappings); + scopeMappings = expandScopeMappings.apply(scopeMappings); // Intersection of expanded user roles and expanded scopeMappings roleMappings.retainAll(scopeMappings.collect(Collectors.toSet())); diff --git a/services/src/main/java/org/keycloak/services/resources/account/AccountConsole.java b/services/src/main/java/org/keycloak/services/resources/account/AccountConsole.java index 5959223aa749..e66309aad455 100644 --- a/services/src/main/java/org/keycloak/services/resources/account/AccountConsole.java +++ b/services/src/main/java/org/keycloak/services/resources/account/AccountConsole.java @@ -347,7 +347,7 @@ static class AccountRoleChecker { AccountRoleChecker(KeycloakSession session, RealmModel realm, UserModel user) { this.accountClient = realm.getClientByClientId(Constants.ACCOUNT_MANAGEMENT_CLIENT_ID); ClientModel accountConsoleClient = realm.getClientByClientId(Constants.ACCOUNT_CONSOLE_CLIENT_ID); - this.scopeResolvedRoles = TokenManager.getAccess(user, accountConsoleClient, TokenManager.getRequestedClientScopes(session, null, accountConsoleClient, user)); + this.scopeResolvedRoles = TokenManager.getAccess(session, user, accountConsoleClient, TokenManager.getRequestedClientScopes(session, null, accountConsoleClient, user)); } boolean hasRole(String roleName) { diff --git a/services/src/main/java/org/keycloak/services/util/DefaultClientSessionContext.java b/services/src/main/java/org/keycloak/services/util/DefaultClientSessionContext.java index d2cb48c957ed..8b50300c028b 100644 --- a/services/src/main/java/org/keycloak/services/util/DefaultClientSessionContext.java +++ b/services/src/main/java/org/keycloak/services/util/DefaultClientSessionContext.java @@ -285,7 +285,7 @@ private boolean isClientScopePermittedForUser(ClientScopeModel clientScope) { } // Expand (resolve composite roles) - clientScopeRoles = RoleUtils.expandCompositeRoles(clientScopeRoles); + clientScopeRoles = RoleUtils.expandCompositeRoles(session, clientSession.getClient().getRealm(), clientScopeRoles); //remove roles that are not contained in requested audience if (attributes.get(Constants.REQUESTED_AUDIENCE_CLIENTS) != null) { @@ -304,7 +304,7 @@ private boolean isClientScopePermittedForUser(ClientScopeModel clientScope) { private Set loadRoles() { UserModel user = clientSession.getUserSession().getUser(); ClientModel client = clientSession.getClient(); - return TokenManager.getAccess(user, client, getClientScopesStream()); + return TokenManager.getAccess(session, user, client, getClientScopesStream()); } @@ -329,7 +329,7 @@ private Set loadProtocolMappers() { private Set loadUserRoles() { UserModel user = clientSession.getUserSession().getUser(); - return RoleUtils.getDeepUserRoleMappings(user); + return RoleUtils.getDeepUserRoleMappings(session, clientSession.getClient().getRealm(), user); } } From 59141b14b33a8fc56389964f7ca168ba5eedc10a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=BCseyin=20Akdo=C4=9Fan?= Date: Fri, 26 Jun 2026 17:57:50 +0300 Subject: [PATCH 05/10] Address review feedback on composite-role expansion tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - CompositeRoleExpansionQueryCountTest: require JpaConnectionProvider so the test is skipped rather than failing with an NPE on non-JPA backends (it reads Hibernate statistics via JpaConnectionProvider). - RoleProviderCompositeDefaultTest: pass a non-null RealmModel (honoring the getCompositeRolesStream contract) instead of null. Signed-off-by: Hüseyin Akdoğan --- .../models/RoleProviderCompositeDefaultTest.java | 12 ++++++++++-- .../role/CompositeRoleExpansionQueryCountTest.java | 1 + 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/server-spi/src/test/java/org/keycloak/models/RoleProviderCompositeDefaultTest.java b/server-spi/src/test/java/org/keycloak/models/RoleProviderCompositeDefaultTest.java index a82507333012..284b972a03f3 100644 --- a/server-spi/src/test/java/org/keycloak/models/RoleProviderCompositeDefaultTest.java +++ b/server-spi/src/test/java/org/keycloak/models/RoleProviderCompositeDefaultTest.java @@ -1,5 +1,6 @@ package org.keycloak.models; +import java.lang.reflect.Proxy; import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; @@ -16,6 +17,13 @@ // children correctly through the interface default implementation (de-duplicated, per parent). public class RoleProviderCompositeDefaultTest { + // Non-null realm satisfying the getCompositeRolesStream contract; the default impl only forwards + // it to getRoleById (ignored by the stub), so it is never actually invoked. + private static final RealmModel REALM = (RealmModel) Proxy.newProxyInstance( + RoleProviderCompositeDefaultTest.class.getClassLoader(), + new Class[]{RealmModel.class}, + (proxy, method, args) -> { throw new UnsupportedOperationException(); }); + @Test public void defaultGetCompositeRolesStreamResolvesAndDeduplicates() { // root -> a, b ; a -> c ; b -> c (diamond, c shared) ; c is a leaf @@ -29,11 +37,11 @@ public void defaultGetCompositeRolesStreamResolvesAndDeduplicates() { // c is shared by a and b but must appear once Assert.assertEquals(Set.of("c"), childIds(provider, "a", "b")); Assert.assertEquals(Set.of(), childIds(provider, "c")); - Assert.assertEquals(Set.of(), provider.getCompositeRolesStream(null, Set.of()).collect(Collectors.toSet())); + Assert.assertEquals(Set.of(), provider.getCompositeRolesStream(REALM, Set.of()).collect(Collectors.toSet())); } private static Set childIds(FakeRoleProvider provider, String... parentIds) { - return provider.getCompositeRolesStream(null, Set.of(parentIds)) + return provider.getCompositeRolesStream(REALM, Set.of(parentIds)) .map(RoleModel::getId).collect(Collectors.toSet()); } diff --git a/testsuite/model/src/test/java/org/keycloak/testsuite/model/role/CompositeRoleExpansionQueryCountTest.java b/testsuite/model/src/test/java/org/keycloak/testsuite/model/role/CompositeRoleExpansionQueryCountTest.java index 26e6191f574e..1c585fc7c3f1 100644 --- a/testsuite/model/src/test/java/org/keycloak/testsuite/model/role/CompositeRoleExpansionQueryCountTest.java +++ b/testsuite/model/src/test/java/org/keycloak/testsuite/model/role/CompositeRoleExpansionQueryCountTest.java @@ -25,6 +25,7 @@ @RequireProvider(RealmProvider.class) @RequireProvider(RoleProvider.class) +@RequireProvider(JpaConnectionProvider.class) public class CompositeRoleExpansionQueryCountTest extends KeycloakModelTest { // Deep composite chain r0 -> r1 -> ... -> r9 (r9 is a leaf): V = 10 nodes, C = 9 composite nodes. From 0a551df48e69c6fbaa283cfcc6de7c18aa84669e Mon Sep 17 00:00:00 2001 From: Alexander Schwartz Date: Sun, 28 Jun 2026 21:03:51 +0200 Subject: [PATCH 06/10] Use existing APIs Signed-off-by: Alexander Schwartz --- .../provider/role/RolePolicyProvider.java | 2 +- .../org/keycloak/models/utils/RoleUtils.java | 95 +++---------------- .../keycloak/protocol/oidc/TokenManager.java | 14 +-- .../resources/account/AccountConsole.java | 2 +- .../util/DefaultClientSessionContext.java | 6 +- 5 files changed, 22 insertions(+), 97 deletions(-) diff --git a/authz/policy/common/src/main/java/org/keycloak/authorization/policy/provider/role/RolePolicyProvider.java b/authz/policy/common/src/main/java/org/keycloak/authorization/policy/provider/role/RolePolicyProvider.java index 3e9cd062b799..50d63f91f28f 100644 --- a/authz/policy/common/src/main/java/org/keycloak/authorization/policy/provider/role/RolePolicyProvider.java +++ b/authz/policy/common/src/main/java/org/keycloak/authorization/policy/provider/role/RolePolicyProvider.java @@ -141,7 +141,7 @@ public Stream getPermissions(KeycloakSession session, ResourceType resou StoreFactory storeFactory = provider.getStoreFactory(); ResourceServer resourceServer = storeFactory.getResourceServerStore().findByClient(adminPermissionsClient); PolicyStore policyStore = storeFactory.getPolicyStore(); - List roleIds = getDeepUserRoleMappings(session, realm, subject).stream().map(RoleModel::getId).toList(); + List roleIds = getDeepUserRoleMappings(subject).stream().map(RoleModel::getId).toList(); Stream policies = Stream.of(); return Stream.concat(policies, policyStore.findDependentPolicies(resourceServer, resourceType.getType(), groupResourceType == null ? null : groupResourceType.getType(), RolePolicyProviderFactory.ID, "roles", roleIds)); diff --git a/server-spi/src/main/java/org/keycloak/models/utils/RoleUtils.java b/server-spi/src/main/java/org/keycloak/models/utils/RoleUtils.java index d56b20750192..05af33cc6e4a 100644 --- a/server-spi/src/main/java/org/keycloak/models/utils/RoleUtils.java +++ b/server-spi/src/main/java/org/keycloak/models/utils/RoleUtils.java @@ -17,8 +17,6 @@ package org.keycloak.models.utils; -import java.util.ArrayDeque; -import java.util.Deque; import java.util.HashSet; import java.util.Objects; import java.util.Set; @@ -33,6 +31,7 @@ import org.keycloak.models.RoleMapperModel; import org.keycloak.models.RoleModel; import org.keycloak.models.UserModel; +import org.keycloak.utils.KeycloakSessionUtil; /** * @author Stian Thorgersen @@ -129,71 +128,21 @@ public static boolean hasRoleFromGroup(Stream groups, RoleModel targ return groups.anyMatch(group -> hasRoleFromGroup(group, targetRole, checkParentGroup)); } - /** - * Recursively expands composite roles into their composite. - * @param role - * @param visited Track roles, which were already visited. Those will be ignored and won't be added to the stream. Besides that, - * the "visited" set itself will be updated as a result of this method call and all the tracked roles will be added to it - * @return Stream of containing all of the composite roles and their components. Never returns {@code null}. - */ - private static Stream expandCompositeRolesStream(RoleModel role, Set visited) { - Stream.Builder sb = Stream.builder(); - - if (!visited.contains(role)) { - Deque stack = new ArrayDeque<>(); - stack.add(role); - - while (!stack.isEmpty()) { - RoleModel current = stack.pop(); - sb.add(current); - - // Fetch the composites directly rather than gating on isComposite() first. - // On the JPA (cache-miss) path isComposite() runs its own getChildRoles query -- - // and the FlushMode.AUTO flush before it -- immediately before getCompositesStream() - // runs another, doubling the per-node round-trips while expanding the tree. - // getCompositesStream() already yields an empty stream for non-composite roles, - // so the pre-check is redundant. - current.getCompositesStream() - .filter(r -> !visited.contains(r)) - .forEach(r -> { - visited.add(r); - stack.add(r); - }); - } - } - - return sb.build(); - } - - /** * @param roles * @return new set with composite roles expanded */ public static Set expandCompositeRoles(Set roles) { - Set visited = new HashSet<>(); - - return roles.stream() - .flatMap(roleModel -> RoleUtils.expandCompositeRolesStream(roleModel, visited)) - .collect(Collectors.toSet()); - } + if (roles.isEmpty()) { + return roles; + } - /** - * Bulk variant of {@link #expandCompositeRoles(Set)}. Expands the composite-role tree one - * breadth-first level at a time using {@code session.roles().getCompositeRolesStream}, issuing a - * single child lookup per level instead of one per role. Storage-backed providers batch that - * lookup, so this avoids the per-node N+1 on a cache miss. - * - * @param session Keycloak session. - * @param realm Realm the roles belong to. - * @param roles Roles to expand. - * @return new set with composite roles expanded - */ - public static Set expandCompositeRoles(KeycloakSession session, RealmModel realm, Set roles) { Set result = new HashSet<>(); Set enqueued = new HashSet<>(); Set frontier = roles; + KeycloakSession session = KeycloakSessionUtil.getKeycloakSession(); + RealmModel realm = realmOf(roles.iterator().next()); while (!frontier.isEmpty()) { result.addAll(frontier); @@ -220,20 +169,9 @@ public static Set expandCompositeRoles(KeycloakSession session, Realm * @return stream with composite roles expanded */ public static Stream expandCompositeRolesStream(Stream roles) { - Set visited = new HashSet<>(); - - return roles.flatMap(roleModel -> RoleUtils.expandCompositeRolesStream(roleModel, visited)); + return RoleUtils.expandCompositeRoles(roles.collect(Collectors.toSet())).stream(); } - /** - * Bulk variant of {@link #expandCompositeRolesStream(Stream)} backed by - * {@link #expandCompositeRoles(KeycloakSession, RealmModel, Set)}. - */ - public static Stream expandCompositeRolesStream(KeycloakSession session, RealmModel realm, Stream roles) { - return expandCompositeRoles(session, realm, roles.collect(Collectors.toSet())).stream(); - } - - /** * @param roleMapper * @return all role mappings for the given mapper with composite roles expanded. @@ -262,23 +200,20 @@ public static Set getDeepUserRoleMappings(UserModel user) { return expandCompositeRoles(roleMappings); } - /** - * Bulk variant of {@link #getDeepUserRoleMappings(UserModel)} that expands composite roles with - * a single child lookup per breadth-first level instead of one per role on a cache miss. - */ - public static Set getDeepUserRoleMappings(KeycloakSession session, RealmModel realm, UserModel user) { - Set roleMappings = user.getRoleMappingsStream().collect(Collectors.toSet()); - user.getGroupsStream().forEach(group -> addGroupRoles(group, roleMappings)); - return expandCompositeRoles(session, realm, roleMappings); - } - - private static void addGroupRoles(GroupModel group, Set roleMappings) { roleMappings.addAll(group.getRoleMappingsStream().collect(Collectors.toSet())); if (group.getParentId() == null) return; addGroupRoles(group.getParent(), roleMappings); } + private static RealmModel realmOf(RoleModel role) { + RoleContainerModel container = role.getContainer(); + if (container instanceof RealmModel realmModel) { + return realmModel; + } + return ((ClientModel) container).getRealm(); + } + public static boolean isRealmRole(RoleModel r) { return r.getContainer() instanceof RealmModel; } diff --git a/services/src/main/java/org/keycloak/protocol/oidc/TokenManager.java b/services/src/main/java/org/keycloak/protocol/oidc/TokenManager.java index c547a27d3756..8d51044062c9 100755 --- a/services/src/main/java/org/keycloak/protocol/oidc/TokenManager.java +++ b/services/src/main/java/org/keycloak/protocol/oidc/TokenManager.java @@ -473,18 +473,8 @@ public static void dettachClientSession(AuthenticatedClientSessionModel clientSe public static Set getAccess(UserModel user, ClientModel client, Stream clientScopes) { - return getAccess(RoleUtils.getDeepUserRoleMappings(user), client, clientScopes, - RoleUtils::expandCompositeRolesStream); - } - - public static Set getAccess(KeycloakSession session, UserModel user, ClientModel client, Stream clientScopes) { - RealmModel realm = client.getRealm(); - return getAccess(RoleUtils.getDeepUserRoleMappings(session, realm, user), client, clientScopes, - scopeMappings -> RoleUtils.expandCompositeRolesStream(session, realm, scopeMappings)); - } + Set roleMappings = RoleUtils.getDeepUserRoleMappings(user); - private static Set getAccess(Set roleMappings, ClientModel client, Stream clientScopes, - Function, Stream> expandScopeMappings) { if (client.isFullScopeAllowed()) { if (logger.isTraceEnabled()) { logger.tracef("Using full scope for client %s", client.getClientId()); @@ -509,7 +499,7 @@ private static Set getAccess(Set roleMappings, ClientModel scopeMappings = Stream.concat(scopeMappings, clientScopesMappings); // 3 - Expand scope mappings - scopeMappings = expandScopeMappings.apply(scopeMappings); + scopeMappings = RoleUtils.expandCompositeRolesStream(scopeMappings); // Intersection of expanded user roles and expanded scopeMappings roleMappings.retainAll(scopeMappings.collect(Collectors.toSet())); diff --git a/services/src/main/java/org/keycloak/services/resources/account/AccountConsole.java b/services/src/main/java/org/keycloak/services/resources/account/AccountConsole.java index e66309aad455..5959223aa749 100644 --- a/services/src/main/java/org/keycloak/services/resources/account/AccountConsole.java +++ b/services/src/main/java/org/keycloak/services/resources/account/AccountConsole.java @@ -347,7 +347,7 @@ static class AccountRoleChecker { AccountRoleChecker(KeycloakSession session, RealmModel realm, UserModel user) { this.accountClient = realm.getClientByClientId(Constants.ACCOUNT_MANAGEMENT_CLIENT_ID); ClientModel accountConsoleClient = realm.getClientByClientId(Constants.ACCOUNT_CONSOLE_CLIENT_ID); - this.scopeResolvedRoles = TokenManager.getAccess(session, user, accountConsoleClient, TokenManager.getRequestedClientScopes(session, null, accountConsoleClient, user)); + this.scopeResolvedRoles = TokenManager.getAccess(user, accountConsoleClient, TokenManager.getRequestedClientScopes(session, null, accountConsoleClient, user)); } boolean hasRole(String roleName) { diff --git a/services/src/main/java/org/keycloak/services/util/DefaultClientSessionContext.java b/services/src/main/java/org/keycloak/services/util/DefaultClientSessionContext.java index 8b50300c028b..d2cb48c957ed 100644 --- a/services/src/main/java/org/keycloak/services/util/DefaultClientSessionContext.java +++ b/services/src/main/java/org/keycloak/services/util/DefaultClientSessionContext.java @@ -285,7 +285,7 @@ private boolean isClientScopePermittedForUser(ClientScopeModel clientScope) { } // Expand (resolve composite roles) - clientScopeRoles = RoleUtils.expandCompositeRoles(session, clientSession.getClient().getRealm(), clientScopeRoles); + clientScopeRoles = RoleUtils.expandCompositeRoles(clientScopeRoles); //remove roles that are not contained in requested audience if (attributes.get(Constants.REQUESTED_AUDIENCE_CLIENTS) != null) { @@ -304,7 +304,7 @@ private boolean isClientScopePermittedForUser(ClientScopeModel clientScope) { private Set loadRoles() { UserModel user = clientSession.getUserSession().getUser(); ClientModel client = clientSession.getClient(); - return TokenManager.getAccess(session, user, client, getClientScopesStream()); + return TokenManager.getAccess(user, client, getClientScopesStream()); } @@ -329,7 +329,7 @@ private Set loadProtocolMappers() { private Set loadUserRoles() { UserModel user = clientSession.getUserSession().getUser(); - return RoleUtils.getDeepUserRoleMappings(session, clientSession.getClient().getRealm(), user); + return RoleUtils.getDeepUserRoleMappings(user); } } From ea3f4ff9ba6afd75494f1253a5882b0bfe9cd00e Mon Sep 17 00:00:00 2001 From: Alexander Schwartz Date: Sun, 28 Jun 2026 22:59:14 +0200 Subject: [PATCH 07/10] Adding caching Signed-off-by: Alexander Schwartz --- .../cache/infinispan/RealmCacheSession.java | 45 +++++++++++++++ .../entities/CachedCompositeRoles.java | 55 +++++++++++++++++++ .../infinispan/stream/HasRolePredicate.java | 2 + 3 files changed, 102 insertions(+) create mode 100755 model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/entities/CachedCompositeRoles.java diff --git a/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/RealmCacheSession.java b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/RealmCacheSession.java index 935f2d388092..79feba52a0ff 100755 --- a/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/RealmCacheSession.java +++ b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/RealmCacheSession.java @@ -22,6 +22,7 @@ import java.util.LinkedList; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Set; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -50,6 +51,7 @@ import org.keycloak.models.cache.infinispan.entities.CachedClient; import org.keycloak.models.cache.infinispan.entities.CachedClientRole; import org.keycloak.models.cache.infinispan.entities.CachedClientScope; +import org.keycloak.models.cache.infinispan.entities.CachedCompositeRoles; import org.keycloak.models.cache.infinispan.entities.CachedGroup; import org.keycloak.models.cache.infinispan.entities.CachedRealm; import org.keycloak.models.cache.infinispan.entities.CachedRealmRole; @@ -1507,6 +1509,49 @@ public void addClientScopeToAllClients(RealmModel realm, ClientScopeModel client invalidationEvents.add(ClientScopeRemovedEvent.create(clientScope.getId(), realm.getId())); } + public Stream getCompositeRolesStream(RealmModel realm, Set parentRoleIds) { + String id = parentRoleIds.stream().sorted().toList().toString(); + CachedCompositeRoles cached = cache.get(id, CachedCompositeRoles.class); + if (cached != null && !cached.getRealm().equals(realm.getId())) { + cached = null; + } + if (cached != null) { + for (String compositeId : cached.getCompositeIds()) { + if (invalidations.contains(compositeId)) { + // If this composite role has been invalidated, it might be faster to fetch it and others from the DB as a batch + cached = null; + break; + } + if (managedRoles.containsKey(compositeId)) { + // If this composite role is already managed, continue + continue; + } + if (getCachedRole(realm, compositeId) == null) { + // If this composite role is not cached, it might be faster to fetch it and others from the DB as a batch + cached = null; + break; + } + } + } + if (cached != null) { + for (String parentId : cached.getParentIds()) { + if (cache.get(parentId, CachedRole.class).getRevision() > cached.getRevision()) { + // The parent has been modified, cache is invalid + cached = null; + break; + } + } + } + if (cached != null) { + return cached.getCompositeIds().stream().map(s -> getRoleById(realm, s)).filter(Objects::nonNull); + } + List compositeRoles = getRoleDelegate().getCompositeRolesStream(realm, parentRoleIds).toList(); + long loaded = cache.getCurrentRevision(id); + cached = new CachedCompositeRoles(loaded, id, new HashSet<>(parentRoleIds), compositeRoles.stream().map(RoleModel::getId).toList(), realm); + cache.addRevisioned(cached, startupRevision); + return compositeRoles.stream(); + } + // Don't cache ClientInitialAccessModel for now @Override public ClientInitialAccessModel createClientInitialAccessModel(RealmModel realm, int expiration, int count) { diff --git a/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/entities/CachedCompositeRoles.java b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/entities/CachedCompositeRoles.java new file mode 100755 index 000000000000..a20800b03044 --- /dev/null +++ b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/entities/CachedCompositeRoles.java @@ -0,0 +1,55 @@ +/* + * Copyright 2016 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.entities; + +import java.util.List; +import java.util.Set; + +import org.keycloak.models.RealmModel; + +/** + * @author Bill Burke + * @version $Revision: 1 $ + */ +public class CachedCompositeRoles extends AbstractRevisioned implements InRealm { + + private final Set parentIds; + final protected String realm; + private final List compositeIds; + + public CachedCompositeRoles(long revision, String id, Set parentIds, List compositeIds, RealmModel realm) { + super(revision, id); + this.parentIds = parentIds; + this.realm = realm.getId(); + this.compositeIds = compositeIds; + } + + @Override + public String getRealm() { + return realm; + } + + public Set getParentIds() { + return parentIds; + } + + public List getCompositeIds() { + return compositeIds; + } + +} diff --git a/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/stream/HasRolePredicate.java b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/stream/HasRolePredicate.java index 59ce4c36d652..c9efdc7225c3 100755 --- a/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/stream/HasRolePredicate.java +++ b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/stream/HasRolePredicate.java @@ -6,6 +6,7 @@ import org.keycloak.marshalling.Marshalling; import org.keycloak.models.cache.infinispan.entities.CachedClient; import org.keycloak.models.cache.infinispan.entities.CachedClientScope; +import org.keycloak.models.cache.infinispan.entities.CachedCompositeRoles; import org.keycloak.models.cache.infinispan.entities.CachedGroup; import org.keycloak.models.cache.infinispan.entities.CachedRole; import org.keycloak.models.cache.infinispan.entities.Revisioned; @@ -44,6 +45,7 @@ void setRole(String role) { public boolean test(Map.Entry entry) { Object value = entry.getValue(); return (value instanceof CachedRole cachedRole && cachedRole.getCachedComposites().ids().contains(role)) || + (value instanceof CachedCompositeRoles cachedCompositeRoles && (cachedCompositeRoles.getCompositeIds().contains(role) || cachedCompositeRoles.getParentIds().contains(role))) || (value instanceof CachedGroup cachedGroup && cachedGroup.getCachedRoleMappings().contains(role)) || (value instanceof RoleQuery roleQuery && roleQuery.getRoles().contains(role)) || (value instanceof CachedClient cachedClient && cachedClient.getScope().contains(role)) || From 84bc75e2cc9a530d0580392def25986175f00569 Mon Sep 17 00:00:00 2001 From: Alexander Schwartz Date: Sun, 28 Jun 2026 23:02:44 +0200 Subject: [PATCH 08/10] Remove unit test as this should be tested with integration tests as per Keycloak strategy Signed-off-by: Alexander Schwartz --- .../RoleUtilsCompositeExpansionTest.java | 182 ------------------ 1 file changed, 182 deletions(-) delete mode 100644 server-spi/src/test/java/org/keycloak/models/utils/RoleUtilsCompositeExpansionTest.java diff --git a/server-spi/src/test/java/org/keycloak/models/utils/RoleUtilsCompositeExpansionTest.java b/server-spi/src/test/java/org/keycloak/models/utils/RoleUtilsCompositeExpansionTest.java deleted file mode 100644 index 2373b977fd25..000000000000 --- a/server-spi/src/test/java/org/keycloak/models/utils/RoleUtilsCompositeExpansionTest.java +++ /dev/null @@ -1,182 +0,0 @@ -package org.keycloak.models.utils; - -import java.util.ArrayList; -import java.util.Arrays; -import java.util.List; -import java.util.Map; -import java.util.Set; -import java.util.stream.Stream; - -import org.junit.Assert; -import org.junit.Test; -import org.keycloak.models.RoleContainerModel; -import org.keycloak.models.RoleModel; - -public class RoleUtilsCompositeExpansionTest { - - @Test - public void expandsFullTreeWithoutCallingIsComposite() { - // root -> a, b ; a -> c ; b -> c (diamond, c shared) ; c is a leaf - FakeRole c = new FakeRole("c"); - FakeRole a = new FakeRole("a", c); - FakeRole b = new FakeRole("b", c); - FakeRole root = new FakeRole("root", a, b); - - Set expanded = RoleUtils.expandCompositeRoles(Set.of(root)); - - Assert.assertEquals(Set.of(root, a, b, c), expanded); - - // The optimization: isComposite() is never used to drive the expansion. - for (FakeRole r : Arrays.asList(root, a, b, c)) { - Assert.assertEquals("isComposite() must not be called for " + r.id, 0, r.isCompositeCalls); - } - // Each reachable role is resolved exactly once thanks to the visited set. - Assert.assertEquals(1, root.getCompositesCalls); - Assert.assertEquals(1, a.getCompositesCalls); - Assert.assertEquals(1, b.getCompositesCalls); - Assert.assertEquals(1, c.getCompositesCalls); - } - - @Test - public void leafRoleExpandsToItself() { - FakeRole leaf = new FakeRole("leaf"); - - Assert.assertEquals(Set.of(leaf), RoleUtils.expandCompositeRoles(Set.of(leaf))); - Assert.assertEquals(0, leaf.isCompositeCalls); - } - - @Test - public void cyclicCompositesTerminate() { - // root -> a -> root (cycle) - FakeRole a = new FakeRole("a"); - FakeRole root = new FakeRole("root", a); - a.composites.add(root); - - Set expanded = RoleUtils.expandCompositeRoles(Set.of(root)); - - Assert.assertEquals(Set.of(root, a), expanded); - } - - // RoleModel stub recording isComposite()/getCompositesStream() calls; identity by id. - private static final class FakeRole implements RoleModel { - - final String id; - final List composites = new ArrayList<>(); - int isCompositeCalls = 0; - int getCompositesCalls = 0; - - FakeRole(String id, RoleModel... children) { - this.id = id; - this.composites.addAll(Arrays.asList(children)); - } - - @Override - public boolean isComposite() { - isCompositeCalls++; - return !composites.isEmpty(); - } - - @Override - public Stream getCompositesStream() { - getCompositesCalls++; - return new ArrayList<>(composites).stream(); - } - - @Override - public String getId() { - return id; - } - - @Override - public boolean equals(Object o) { - return o instanceof RoleModel && ((RoleModel) o).getId().equals(id); - } - - @Override - public int hashCode() { - return id.hashCode(); - } - - // --- unused RoleModel surface --------------------------------------------------------- - - @Override - public String getName() { - return id; - } - - @Override - public String getDescription() { - throw new UnsupportedOperationException(); - } - - @Override - public void setDescription(String description) { - throw new UnsupportedOperationException(); - } - - @Override - public void setName(String name) { - throw new UnsupportedOperationException(); - } - - @Override - public void addCompositeRole(RoleModel role) { - throw new UnsupportedOperationException(); - } - - @Override - public void removeCompositeRole(RoleModel role) { - throw new UnsupportedOperationException(); - } - - @Override - public Stream getCompositesStream(String search, Integer first, Integer max) { - throw new UnsupportedOperationException(); - } - - @Override - public boolean isClientRole() { - throw new UnsupportedOperationException(); - } - - @Override - public String getContainerId() { - throw new UnsupportedOperationException(); - } - - @Override - public RoleContainerModel getContainer() { - throw new UnsupportedOperationException(); - } - - @Override - public boolean hasRole(RoleModel role) { - throw new UnsupportedOperationException(); - } - - @Override - public void setSingleAttribute(String name, String value) { - throw new UnsupportedOperationException(); - } - - @Override - public void setAttribute(String name, List values) { - throw new UnsupportedOperationException(); - } - - @Override - public void removeAttribute(String name) { - throw new UnsupportedOperationException(); - } - - @Override - public Stream getAttributeStream(String name) { - throw new UnsupportedOperationException(); - } - - @Override - public Map> getAttributes() { - throw new UnsupportedOperationException(); - } - } -} From 23fb815749619cc095900686624bd0636972c9fd Mon Sep 17 00:00:00 2001 From: Alexander Schwartz Date: Mon, 29 Jun 2026 01:18:55 +0200 Subject: [PATCH 09/10] Moving test Signed-off-by: Alexander Schwartz --- .../CompositeRoleExpansionQueryCountTest.java | 107 +++++++++++ .../CompositeRoleExpansionQueryCountTest.java | 166 ------------------ 2 files changed, 107 insertions(+), 166 deletions(-) create mode 100644 tests/base/src/test/java/org/keycloak/tests/model/CompositeRoleExpansionQueryCountTest.java delete mode 100644 testsuite/model/src/test/java/org/keycloak/testsuite/model/role/CompositeRoleExpansionQueryCountTest.java diff --git a/tests/base/src/test/java/org/keycloak/tests/model/CompositeRoleExpansionQueryCountTest.java b/tests/base/src/test/java/org/keycloak/tests/model/CompositeRoleExpansionQueryCountTest.java new file mode 100644 index 000000000000..8af93ba362fd --- /dev/null +++ b/tests/base/src/test/java/org/keycloak/tests/model/CompositeRoleExpansionQueryCountTest.java @@ -0,0 +1,107 @@ +package org.keycloak.tests.model; + +import java.util.HashSet; +import java.util.Set; + +import org.keycloak.connections.jpa.JpaConnectionProvider; +import org.keycloak.models.KeycloakSession; +import org.keycloak.models.RealmModel; +import org.keycloak.models.RoleModel; +import org.keycloak.models.utils.RoleUtils; +import org.keycloak.testframework.annotations.InjectRealm; +import org.keycloak.testframework.annotations.KeycloakIntegrationTest; +import org.keycloak.testframework.realm.ManagedRealm; +import org.keycloak.testframework.realm.RealmBuilder; +import org.keycloak.testframework.realm.RealmConfig; +import org.keycloak.testframework.realm.RoleBuilder; +import org.keycloak.testframework.remote.annotations.TestOnServer; + +import org.hibernate.SessionFactory; +import org.hibernate.stat.Statistics; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.is; + +@KeycloakIntegrationTest +public class CompositeRoleExpansionQueryCountTest { + + // Deep composite chain r0 -> r1 -> ... -> r9 (r9 is a leaf): V = 10 nodes, C = 9 composite nodes. + private static final int CHAIN_LENGTH = 10; + + // Wide tree: root -> BRANCHING mids -> BRANCHING leaves each. V = 1 + 5 + 25 = 31, depth = 3 levels. + private static final int BRANCHING = 5; + private static final int BUSHY_SIZE = 1 + BRANCHING + BRANCHING * BRANCHING; + private static final int BUSHY_LEVELS = 3; + + @InjectRealm(config = CompositeRoleExpansionRealmConfig.class) + ManagedRealm realm; + + @TestOnServer + public void droppingIsCompositePreQueryHalvesGetChildRolesOnCompositeNodes(KeycloakSession session) { + RealmModel realm = session.getContext().getRealm(); + RoleModel root = session.roles().getRealmRole(realm, "chain-0"); + long currentCount = countQueries(session, () -> RoleUtils.expandCompositeRoles(Set.of(root))); + // Current batched code: one query per breadth-first level (equals V for a chain since depth = V). + assertThat(currentCount, is((long) CHAIN_LENGTH)); // L = V = 10 + } + + @TestOnServer + public void levelWiseBatchingCollapsesGetChildRolesToOnePerLevel(KeycloakSession session) { + Set batchedIds = new HashSet<>(); + RealmModel realm = session.getContext().getRealm(); + RoleModel root = session.roles().getRealmRole(realm, "bushy-root"); + long batchedCount = countQueries(session, () -> collectIds(RoleUtils.expandCompositeRoles(Set.of(root)), batchedIds)); + assertThat(batchedIds.size(), is(BUSHY_SIZE)); // 31 + assertThat(batchedCount, is((long) BUSHY_LEVELS)); // 3 + } + + private static void collectIds(Set roles, Set into) { + roles.forEach(role -> into.add(role.getId())); + } + + private static long countQueries(KeycloakSession session, Runnable walk) { + Statistics stats = session.getProvider(JpaConnectionProvider.class) + .getEntityManager() + .getEntityManagerFactory() + .unwrap(SessionFactory.class) + .getStatistics(); + stats.setStatisticsEnabled(true); + stats.clear(); + walk.run(); + return stats.getPrepareStatementCount(); + } + + private static class CompositeRoleExpansionRealmConfig implements RealmConfig { + @Override + public RealmBuilder configure(RealmBuilder r) { + r.name("composite-role-expansion-realm"); + + RoleBuilder bushyRoot = RoleBuilder.create().name("bushy-root"); + r.realmRoles(bushyRoot); + for (int i = 0; i < BRANCHING; i++) { + String name = "bushy-mid-" + i; + RoleBuilder bushyMid = RoleBuilder.create().name(name); + r.realmRoles(bushyMid); + bushyRoot.realmComposite(name); + for (int j = 0; j < BRANCHING; j++) { + String leafName = "bushy-leaf-" + i + "-" + j; + r.realmRoles(leafName); + bushyMid.realmComposite(leafName); + } + } + + RoleBuilder previous = null; + for (int i = 0; i < CHAIN_LENGTH; i++) { + String name = "chain-" + i; + RoleBuilder role = RoleBuilder.create().name(name); + r.realmRoles(role); + if (previous != null) { + previous.realmComposite(name); + } + previous = role; + } + + return r; + } + } +} diff --git a/testsuite/model/src/test/java/org/keycloak/testsuite/model/role/CompositeRoleExpansionQueryCountTest.java b/testsuite/model/src/test/java/org/keycloak/testsuite/model/role/CompositeRoleExpansionQueryCountTest.java deleted file mode 100644 index 1c585fc7c3f1..000000000000 --- a/testsuite/model/src/test/java/org/keycloak/testsuite/model/role/CompositeRoleExpansionQueryCountTest.java +++ /dev/null @@ -1,166 +0,0 @@ -package org.keycloak.testsuite.model.role; - -import java.util.ArrayDeque; -import java.util.Deque; -import java.util.HashSet; -import java.util.Set; - -import org.keycloak.connections.jpa.JpaConnectionProvider; -import org.keycloak.models.Constants; -import org.keycloak.models.KeycloakSession; -import org.keycloak.models.RealmModel; -import org.keycloak.models.RealmProvider; -import org.keycloak.models.RoleModel; -import org.keycloak.models.RoleProvider; -import org.keycloak.models.utils.RoleUtils; -import org.keycloak.testsuite.model.KeycloakModelTest; -import org.keycloak.testsuite.model.RequireProvider; - -import org.hibernate.SessionFactory; -import org.hibernate.stat.Statistics; -import org.junit.Test; - -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.is; - -@RequireProvider(RealmProvider.class) -@RequireProvider(RoleProvider.class) -@RequireProvider(JpaConnectionProvider.class) -public class CompositeRoleExpansionQueryCountTest extends KeycloakModelTest { - - // Deep composite chain r0 -> r1 -> ... -> r9 (r9 is a leaf): V = 10 nodes, C = 9 composite nodes. - private static final int CHAIN_LENGTH = 10; - - // Wide tree: root -> BRANCHING mids -> BRANCHING leaves each. V = 1 + 5 + 25 = 31, depth = 3 levels. - private static final int BRANCHING = 5; - private static final int BUSHY_SIZE = 1 + BRANCHING + BRANCHING * BRANCHING; - private static final int BUSHY_LEVELS = 3; - - private String realmId; - private String rootRoleId; - private String bushyRootId; - - @Override - public void createEnvironment(KeycloakSession s) { - RealmModel realm = createRealm(s, "composite-expansion-realm"); - s.getContext().setRealm(realm); - realm.setDefaultRole(s.roles().addRealmRole(realm, Constants.DEFAULT_ROLES_ROLE_PREFIX + "-" + realm.getName())); - realmId = realm.getId(); - - RoleModel previous = null; - for (int i = 0; i < CHAIN_LENGTH; i++) { - RoleModel role = s.roles().addRealmRole(realm, "chain-" + i); - if (previous == null) { - rootRoleId = role.getId(); - } else { - previous.addCompositeRole(role); - } - previous = role; - } - - RoleModel bushyRoot = s.roles().addRealmRole(realm, "bushy-root"); - bushyRootId = bushyRoot.getId(); - for (int i = 0; i < BRANCHING; i++) { - RoleModel mid = s.roles().addRealmRole(realm, "bushy-mid-" + i); - bushyRoot.addCompositeRole(mid); - for (int j = 0; j < BRANCHING; j++) { - mid.addCompositeRole(s.roles().addRealmRole(realm, "bushy-leaf-" + i + "-" + j)); - } - } - } - - @Override - public void cleanEnvironment(KeycloakSession s) { - RealmModel realm = s.realms().getRealm(realmId); - s.getContext().setRealm(realm); - s.realms().removeRealm(realmId); - } - - @Test - public void droppingIsCompositePreQueryHalvesGetChildRolesOnCompositeNodes() { - long currentCount = withRealm(realmId, (s, realm) -> { - RoleModel root = s.roles().getRoleById(realm, rootRoleId); - return countQueries(s, () -> RoleUtils.expandCompositeRoles(Set.of(root))); - }); - - long prePatchCount = withRealm(realmId, (s, realm) -> { - RoleModel root = s.roles().getRoleById(realm, rootRoleId); - return countQueries(s, () -> expandWithIsCompositePreQuery(Set.of(root))); - }); - - // Current code: one getChildRoles per visited node (getCompositesStream on every node). - assertThat(currentCount, is((long) CHAIN_LENGTH)); // V = 10 - // Pre-patch code: isComposite() on every node + getCompositesStream() on composite nodes. - assertThat(prePatchCount, is((long) (2 * CHAIN_LENGTH - 1))); // V + C = 19 - // Saved exactly C = number of composite nodes. - assertThat(prePatchCount - currentCount, is((long) (CHAIN_LENGTH - 1))); // 9 - } - - @Test - public void levelWiseBatchingCollapsesGetChildRolesToOnePerLevel() { - Set perNodeIds = new HashSet<>(); - long perNodeCount = withRealm(realmId, (s, realm) -> { - RoleModel root = s.roles().getRoleById(realm, bushyRootId); - return countQueries(s, () -> collectIds(RoleUtils.expandCompositeRoles(Set.of(root)), perNodeIds)); - }); - - Set batchedIds = new HashSet<>(); - long batchedCount = withRealm(realmId, (s, realm) -> { - RoleModel root = s.roles().getRoleById(realm, bushyRootId); - return countQueries(s, () -> collectIds(RoleUtils.expandCompositeRoles(s, realm, Set.of(root)), batchedIds)); - }); - - // Both expansions yield the same role set ... - assertThat(perNodeIds.size(), is(BUSHY_SIZE)); // 31 - assertThat(batchedIds, is(perNodeIds)); - // ... but the per-node walk issues one getChildRoles query per node ... - assertThat(perNodeCount, is((long) BUSHY_SIZE)); // 31 - // ... whereas level-wise batching issues a single query per breadth-first level. - assertThat(batchedCount, is((long) BUSHY_LEVELS)); // 3 - } - - private static void collectIds(Set roles, Set into) { - roles.forEach(role -> into.add(role.getId())); - } - - private long countQueries(KeycloakSession session, Runnable walk) { - Statistics stats = session.getProvider(JpaConnectionProvider.class) - .getEntityManager() - .getEntityManagerFactory() - .unwrap(SessionFactory.class) - .getStatistics(); - stats.setStatisticsEnabled(true); - stats.clear(); - walk.run(); - // The only JDBC statements during the walk are the getChildRoles named-query executions: - // getCompositesStream() hydrates the child RoleEntity rows, so the subsequent getRoleById - // em.find calls are served from the persistence context without extra statements. - return stats.getPrepareStatementCount(); - } - - // Pre-patch walk: keeps the isComposite() guard so both variants run against the same database. - private static Set expandWithIsCompositePreQuery(Set roles) { - Set visited = new HashSet<>(); - Set result = new HashSet<>(); - for (RoleModel role : roles) { - if (visited.contains(role)) { - continue; - } - Deque stack = new ArrayDeque<>(); - stack.add(role); - while (!stack.isEmpty()) { - RoleModel current = stack.pop(); - result.add(current); - if (current.isComposite()) { - current.getCompositesStream() - .filter(r -> !visited.contains(r)) - .forEach(r -> { - visited.add(r); - stack.add(r); - }); - } - } - } - return result; - } -} From a4e4b8226c87992ddeca5855ab026e9668069dd5 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 29 Jun 2026 01:14:17 +0000 Subject: [PATCH 10/10] Fix review comments on batch composite-role expansion - RoleUtils.expandCompositeRoles: guard against null KeycloakSession by falling back to per-role stack expansion when no session context is available (fixes r3488823474) - RoleStorageManager.getCompositeRolesStream: split parent IDs into local and external; external (federated) IDs fall back to per-role getRoleById + getCompositesStream instead of being silently dropped (fixes r3488823485) - RealmCacheSession.getCompositeRolesStream: prefix cache key with realm ID to prevent cross-realm key collisions; use String.join for a more predictable key format (fixes r3488823491) - RealmCacheSession.getCompositeRolesStream: null-check the result of cache.get(parentId, CachedRole.class) before dereferencing it to prevent NPE on evicted entries (fixes r3488823493) - CompositeRoleExpansionQueryCountTest.countQueries: save and restore Hibernate statistics enabled state to avoid affecting other tests (fixes r3488823497) - CompositeRoleExpansionQueryCountTest: rename droppingIsCompositePreQueryHalvesGetChildRolesOnCompositeNodes to chainWalkIssuesOneQueryPerBreadthFirstLevel (fixes r3488823503) --- .../cache/infinispan/RealmCacheSession.java | 7 ++-- .../keycloak/storage/RoleStorageManager.java | 24 +++++++++++- .../org/keycloak/models/utils/RoleUtils.java | 38 ++++++++++++++++++- .../CompositeRoleExpansionQueryCountTest.java | 7 +++- 4 files changed, 69 insertions(+), 7 deletions(-) diff --git a/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/RealmCacheSession.java b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/RealmCacheSession.java index 79feba52a0ff..6787d60e810f 100755 --- a/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/RealmCacheSession.java +++ b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/RealmCacheSession.java @@ -1510,7 +1510,7 @@ public void addClientScopeToAllClients(RealmModel realm, ClientScopeModel client } public Stream getCompositeRolesStream(RealmModel realm, Set parentRoleIds) { - String id = parentRoleIds.stream().sorted().toList().toString(); + String id = realm.getId() + "::" + String.join(",", parentRoleIds.stream().sorted().collect(Collectors.toList())); CachedCompositeRoles cached = cache.get(id, CachedCompositeRoles.class); if (cached != null && !cached.getRealm().equals(realm.getId())) { cached = null; @@ -1535,8 +1535,9 @@ public Stream getCompositeRolesStream(RealmModel realm, Set p } if (cached != null) { for (String parentId : cached.getParentIds()) { - if (cache.get(parentId, CachedRole.class).getRevision() > cached.getRevision()) { - // The parent has been modified, cache is invalid + CachedRole cachedParent = cache.get(parentId, CachedRole.class); + if (cachedParent == null || cachedParent.getRevision() > cached.getRevision()) { + // The parent has been modified or evicted, cache is invalid cached = null; break; } diff --git a/model/storage-private/src/main/java/org/keycloak/storage/RoleStorageManager.java b/model/storage-private/src/main/java/org/keycloak/storage/RoleStorageManager.java index 06dd40474735..9dea97f2b701 100644 --- a/model/storage-private/src/main/java/org/keycloak/storage/RoleStorageManager.java +++ b/model/storage-private/src/main/java/org/keycloak/storage/RoleStorageManager.java @@ -18,6 +18,7 @@ import java.util.Objects; import java.util.Set; +import java.util.stream.Collectors; import java.util.stream.Stream; import org.keycloak.common.util.reflections.Types; @@ -164,7 +165,28 @@ public Stream getRolesStream(RealmModel realm, Stream ids, St @Override public Stream getCompositeRolesStream(RealmModel realm, Set parentRoleIds) { - return localStorage().getCompositeRolesStream(realm, parentRoleIds); + if (parentRoleIds == null || parentRoleIds.isEmpty()) { + return Stream.empty(); + } + // Split local and external IDs. Local IDs are handled by a single batched query; external + // (federated) IDs are resolved individually via getRoleById and then expanded per-role. + Set localIds = parentRoleIds.stream() + .filter(StorageId::isLocalStorage) + .collect(Collectors.toSet()); + Set externalIds = parentRoleIds.stream() + .filter(id -> !StorageId.isLocalStorage(id)) + .collect(Collectors.toSet()); + + Stream localComposites = localIds.isEmpty() + ? Stream.empty() + : localStorage().getCompositeRolesStream(realm, localIds); + + Stream externalComposites = externalIds.stream() + .map(id -> getRoleById(realm, id)) + .filter(Objects::nonNull) + .flatMap(RoleModel::getCompositesStream); + + return Stream.concat(localComposites, externalComposites).distinct(); } /** diff --git a/server-spi/src/main/java/org/keycloak/models/utils/RoleUtils.java b/server-spi/src/main/java/org/keycloak/models/utils/RoleUtils.java index 05af33cc6e4a..2472d5a619ef 100644 --- a/server-spi/src/main/java/org/keycloak/models/utils/RoleUtils.java +++ b/server-spi/src/main/java/org/keycloak/models/utils/RoleUtils.java @@ -17,6 +17,8 @@ package org.keycloak.models.utils; +import java.util.ArrayDeque; +import java.util.Deque; import java.util.HashSet; import java.util.Objects; import java.util.Set; @@ -137,11 +139,20 @@ public static Set expandCompositeRoles(Set roles) { return roles; } + KeycloakSession session = KeycloakSessionUtil.getKeycloakSession(); + if (session == null) { + // Outside a Resteasy/session-bound thread there is no KeycloakSession available. + // Fall back to per-role expansion, which works without a session context. + Set visited = new HashSet<>(); + return roles.stream() + .flatMap(role -> expandCompositeRolesWithoutSession(role, visited)) + .collect(Collectors.toSet()); + } + Set result = new HashSet<>(); Set enqueued = new HashSet<>(); Set frontier = roles; - KeycloakSession session = KeycloakSessionUtil.getKeycloakSession(); RealmModel realm = realmOf(roles.iterator().next()); while (!frontier.isEmpty()) { result.addAll(frontier); @@ -164,6 +175,31 @@ public static Set expandCompositeRoles(Set roles) { return result; } + private static Stream expandCompositeRolesWithoutSession(RoleModel role, Set visited) { + Stream.Builder sb = Stream.builder(); + + if (!visited.contains(role)) { + Deque stack = new ArrayDeque<>(); + stack.add(role); + + while (!stack.isEmpty()) { + RoleModel current = stack.pop(); + sb.add(current); + + if (current.isComposite()) { + current.getCompositesStream() + .filter(r -> !visited.contains(r)) + .forEach(r -> { + visited.add(r); + stack.add(r); + }); + } + } + } + + return sb.build(); + } + /** * @param roles * @return stream with composite roles expanded diff --git a/tests/base/src/test/java/org/keycloak/tests/model/CompositeRoleExpansionQueryCountTest.java b/tests/base/src/test/java/org/keycloak/tests/model/CompositeRoleExpansionQueryCountTest.java index 8af93ba362fd..445df9371ac1 100644 --- a/tests/base/src/test/java/org/keycloak/tests/model/CompositeRoleExpansionQueryCountTest.java +++ b/tests/base/src/test/java/org/keycloak/tests/model/CompositeRoleExpansionQueryCountTest.java @@ -37,7 +37,7 @@ public class CompositeRoleExpansionQueryCountTest { ManagedRealm realm; @TestOnServer - public void droppingIsCompositePreQueryHalvesGetChildRolesOnCompositeNodes(KeycloakSession session) { + public void chainWalkIssuesOneQueryPerBreadthFirstLevel(KeycloakSession session) { RealmModel realm = session.getContext().getRealm(); RoleModel root = session.roles().getRealmRole(realm, "chain-0"); long currentCount = countQueries(session, () -> RoleUtils.expandCompositeRoles(Set.of(root))); @@ -65,10 +65,13 @@ private static long countQueries(KeycloakSession session, Runnable walk) { .getEntityManagerFactory() .unwrap(SessionFactory.class) .getStatistics(); + boolean wasEnabled = stats.isStatisticsEnabled(); stats.setStatisticsEnabled(true); stats.clear(); walk.run(); - return stats.getPrepareStatementCount(); + long count = stats.getPrepareStatementCount(); + stats.setStatisticsEnabled(wasEnabled); + return count; } private static class CompositeRoleExpansionRealmConfig implements RealmConfig {