This is an automated email from the ASF dual-hosted git repository. yuqi4733 pushed a commit to branch fix-10545-batch-load-securable-objects in repository https://gitbox.apache.org/repos/asf/gravitino.git
commit ac25ca7f41a35bc6eadfdd8abd96628d27ba08ea Author: yuqi <[email protected]> AuthorDate: Wed Mar 25 20:12:01 2026 +0800 [#10545] improvement(auth): Batch-load securable objects to eliminate N+1 in loadRolePrivilege - Add listSecurableObjectsByRoleIds(List<Long>) to SecurableObjectMapper and SecurableObjectBaseSQLProvider: fetches securable objects for multiple roles in a single WHERE role_id IN (...) query. - Add RoleMetaService.batchListSecurableObjectsForRoles(List<Long>): issues the batch query and groups results by role ID; extracted buildSecurableObjectsFromPOs helper reused by both single and batch paths. - Rewrite JcasbinAuthorizer.loadRolePrivilege(): collect all unloaded role IDs, call the batch method once, load policies serially. Removes the per-role CompletableFuture + entityStore.get() pattern. Cold-path query count drops from 2+2N+T to 2+1+T (where N = roles, T = distinct object types). - Fix DBCP2 pool: minEvictableIdleTimeMillis 1000ms -> 30000ms, minIdle 0 -> 5, maxIdle 5 -> 10. Prevents connection churn that was adding 5-20ms per request. Co-Authored-By: Claude Sonnet 4.6 <[email protected]> --- .../relational/mapper/SecurableObjectMapper.java | 5 ++ .../mapper/SecurableObjectSQLProviderFactory.java | 4 ++ .../base/SecurableObjectBaseSQLProvider.java | 16 +++++ .../relational/service/RoleMetaService.java | 37 +++++++++++ .../session/SqlSessionFactoryHelper.java | 6 +- .../authorization/jcasbin/JcasbinAuthorizer.java | 77 ++++++++++++---------- .../jcasbin/TestJcasbinAuthorizer.java | 51 +++++++++----- 7 files changed, 145 insertions(+), 51 deletions(-) diff --git a/core/src/main/java/org/apache/gravitino/storage/relational/mapper/SecurableObjectMapper.java b/core/src/main/java/org/apache/gravitino/storage/relational/mapper/SecurableObjectMapper.java index d4e8fdf0a0..82b5b51a95 100644 --- a/core/src/main/java/org/apache/gravitino/storage/relational/mapper/SecurableObjectMapper.java +++ b/core/src/main/java/org/apache/gravitino/storage/relational/mapper/SecurableObjectMapper.java @@ -83,6 +83,11 @@ public interface SecurableObjectMapper { method = "listSecurableObjectsByRoleId") List<SecurableObjectPO> listSecurableObjectsByRoleId(@Param("roleId") Long roleId); + @SelectProvider( + type = SecurableObjectSQLProviderFactory.class, + method = "listSecurableObjectsByRoleIds") + List<SecurableObjectPO> listSecurableObjectsByRoleIds(@Param("roleIds") List<Long> roleIds); + @DeleteProvider( type = SecurableObjectSQLProviderFactory.class, method = "deleteSecurableObjectsByLegacyTimeline") diff --git a/core/src/main/java/org/apache/gravitino/storage/relational/mapper/SecurableObjectSQLProviderFactory.java b/core/src/main/java/org/apache/gravitino/storage/relational/mapper/SecurableObjectSQLProviderFactory.java index dab4bcf703..b404cd7061 100644 --- a/core/src/main/java/org/apache/gravitino/storage/relational/mapper/SecurableObjectSQLProviderFactory.java +++ b/core/src/main/java/org/apache/gravitino/storage/relational/mapper/SecurableObjectSQLProviderFactory.java @@ -89,6 +89,10 @@ public class SecurableObjectSQLProviderFactory { return getProvider().listSecurableObjectsByRoleId(roleId); } + public static String listSecurableObjectsByRoleIds(@Param("roleIds") List<Long> roleIds) { + return getProvider().listSecurableObjectsByRoleIds(roleIds); + } + public static String deleteSecurableObjectsByLegacyTimeline( @Param("legacyTimeline") Long legacyTimeline, @Param("limit") int limit) { return getProvider().deleteSecurableObjectsByLegacyTimeline(legacyTimeline, limit); diff --git a/core/src/main/java/org/apache/gravitino/storage/relational/mapper/provider/base/SecurableObjectBaseSQLProvider.java b/core/src/main/java/org/apache/gravitino/storage/relational/mapper/provider/base/SecurableObjectBaseSQLProvider.java index 69e9d9e2e3..4cce0fdf3d 100644 --- a/core/src/main/java/org/apache/gravitino/storage/relational/mapper/provider/base/SecurableObjectBaseSQLProvider.java +++ b/core/src/main/java/org/apache/gravitino/storage/relational/mapper/provider/base/SecurableObjectBaseSQLProvider.java @@ -191,6 +191,22 @@ public class SecurableObjectBaseSQLProvider { + " WHERE role_id = #{roleId} AND deleted_at = 0"; } + public String listSecurableObjectsByRoleIds(@Param("roleIds") List<Long> roleIds) { + return "<script>" + + "SELECT role_id as roleId, metadata_object_id as metadataObjectId," + + " type as type, privilege_names as privilegeNames," + + " privilege_conditions as privilegeConditions, current_version as currentVersion," + + " last_version as lastVersion, deleted_at as deletedAt" + + " FROM " + + SECURABLE_OBJECT_TABLE_NAME + + " WHERE role_id IN " + + "<foreach collection='roleIds' item='item' open='(' separator=',' close=')'>" + + "#{item}" + + "</foreach>" + + " AND deleted_at = 0" + + "</script>"; + } + public String deleteSecurableObjectsByLegacyTimeline( @Param("legacyTimeline") Long legacyTimeline, @Param("limit") int limit) { return "DELETE FROM " diff --git a/core/src/main/java/org/apache/gravitino/storage/relational/service/RoleMetaService.java b/core/src/main/java/org/apache/gravitino/storage/relational/service/RoleMetaService.java index d0fe1db7dd..9e97cf7af1 100644 --- a/core/src/main/java/org/apache/gravitino/storage/relational/service/RoleMetaService.java +++ b/core/src/main/java/org/apache/gravitino/storage/relational/service/RoleMetaService.java @@ -21,6 +21,7 @@ package org.apache.gravitino.storage.relational.service; import static org.apache.gravitino.metrics.source.MetricsSource.GRAVITINO_RELATIONAL_STORE_METRIC_NAME; import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.Lists; import com.google.common.collect.Sets; import java.io.IOException; @@ -338,6 +339,37 @@ public class RoleMetaService { SecurableObjectMapper.class, mapper -> mapper.listSecurableObjectsByRoleId(roleId)); } + /** + * Batch-loads securable objects for multiple roles in a single SQL query and returns a map from + * role ID to the resolved {@link SecurableObject} list. This eliminates the N+1 query pattern + * that occurs when loading securable objects for each role individually. + * + * @param roleIds the list of role IDs to load + * @return a map from role ID to its list of resolved securable objects + */ + @Monitored( + metricsSource = GRAVITINO_RELATIONAL_STORE_METRIC_NAME, + baseMetricName = "batchListSecurableObjectsForRoles") + public static Map<Long, List<SecurableObject>> batchListSecurableObjectsForRoles( + List<Long> roleIds) { + if (roleIds.isEmpty()) { + return ImmutableMap.of(); + } + List<SecurableObjectPO> allPOs = + SessionUtils.getWithoutCommit( + SecurableObjectMapper.class, mapper -> mapper.listSecurableObjectsByRoleIds(roleIds)); + + Map<Long, List<SecurableObjectPO>> byRoleId = + allPOs.stream().collect(Collectors.groupingBy(SecurableObjectPO::getRoleId)); + + ImmutableMap.Builder<Long, List<SecurableObject>> builder = ImmutableMap.builder(); + for (Long roleId : roleIds) { + List<SecurableObjectPO> pos = byRoleId.getOrDefault(roleId, Collections.emptyList()); + builder.put(roleId, buildSecurableObjectsFromPOs(pos)); + } + return builder.build(); + } + @Monitored( metricsSource = GRAVITINO_RELATIONAL_STORE_METRIC_NAME, baseMetricName = "listRolesByNamespace") @@ -398,6 +430,11 @@ public class RoleMetaService { private static List<SecurableObject> listSecurableObjects(RolePO po) { List<SecurableObjectPO> securableObjectPOs = listSecurableObjectsByRoleId(po.getRoleId()); + return buildSecurableObjectsFromPOs(securableObjectPOs); + } + + private static List<SecurableObject> buildSecurableObjectsFromPOs( + List<SecurableObjectPO> securableObjectPOs) { List<SecurableObject> securableObjects = Lists.newArrayList(); securableObjectPOs.stream() diff --git a/core/src/main/java/org/apache/gravitino/storage/relational/session/SqlSessionFactoryHelper.java b/core/src/main/java/org/apache/gravitino/storage/relational/session/SqlSessionFactoryHelper.java index bc406af3b5..db71b4ae41 100644 --- a/core/src/main/java/org/apache/gravitino/storage/relational/session/SqlSessionFactoryHelper.java +++ b/core/src/main/java/org/apache/gravitino/storage/relational/session/SqlSessionFactoryHelper.java @@ -88,15 +88,15 @@ public class SqlSessionFactoryHelper { dataSource.setMaxWaitMillis( config.get(Configs.ENTITY_RELATIONAL_JDBC_BACKEND_WAIT_MILLISECONDS)); dataSource.setMaxTotal(config.get(Configs.ENTITY_RELATIONAL_JDBC_BACKEND_MAX_CONNECTIONS)); - dataSource.setMaxIdle(5); - dataSource.setMinIdle(0); + dataSource.setMaxIdle(10); + dataSource.setMinIdle(5); dataSource.setLogAbandoned(true); dataSource.setRemoveAbandonedOnBorrow(true); dataSource.setRemoveAbandonedTimeout(60); dataSource.setTimeBetweenEvictionRunsMillis(Duration.ofMillis(10 * 60 * 1000L).toMillis()); dataSource.setTestOnBorrow(true); dataSource.setTestWhileIdle(true); - dataSource.setMinEvictableIdleTimeMillis(1000); + dataSource.setMinEvictableIdleTimeMillis(30_000); dataSource.setNumTestsPerEvictionRun(BaseObjectPoolConfig.DEFAULT_NUM_TESTS_PER_EVICTION_RUN); dataSource.setTestOnReturn(BaseObjectPoolConfig.DEFAULT_TEST_ON_RETURN); dataSource.setSoftMinEvictableIdleTimeMillis( diff --git a/server-common/src/main/java/org/apache/gravitino/server/authorization/jcasbin/JcasbinAuthorizer.java b/server-common/src/main/java/org/apache/gravitino/server/authorization/jcasbin/JcasbinAuthorizer.java index 7630f71ae3..62220b1205 100644 --- a/server-common/src/main/java/org/apache/gravitino/server/authorization/jcasbin/JcasbinAuthorizer.java +++ b/server-common/src/main/java/org/apache/gravitino/server/authorization/jcasbin/JcasbinAuthorizer.java @@ -27,14 +27,16 @@ import java.nio.charset.StandardCharsets; import java.security.Principal; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.Objects; import java.util.Optional; -import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executor; import java.util.concurrent.Executors; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; import org.apache.commons.io.IOUtils; import org.apache.commons.lang3.StringUtils; import org.apache.gravitino.Configs; @@ -55,6 +57,7 @@ import org.apache.gravitino.exceptions.NoSuchUserException; import org.apache.gravitino.meta.RoleEntity; import org.apache.gravitino.meta.UserEntity; import org.apache.gravitino.server.authorization.MetadataIdConverter; +import org.apache.gravitino.storage.relational.service.RoleMetaService; import org.apache.gravitino.utils.MetadataObjectUtil; import org.apache.gravitino.utils.NameIdentifierUtil; import org.apache.gravitino.utils.PrincipalUtils; @@ -490,48 +493,56 @@ public class JcasbinAuthorizer implements GravitinoAuthorizer { () -> { EntityStore entityStore = GravitinoEnv.getInstance().entityStore(); NameIdentifier userNameIdentifier = NameIdentifierUtil.ofUser(metalake, username); - List<RoleEntity> entities; + List<RoleEntity> roleStubs; try { - entities = + roleStubs = entityStore .relationOperations() .listEntitiesByRelation( SupportsRelationOperations.Type.ROLE_USER_REL, userNameIdentifier, Entity.EntityType.USER); - List<CompletableFuture<Void>> loadRoleFutures = new ArrayList<>(); - for (RoleEntity role : entities) { - Long roleId = role.id(); - allowEnforcer.addRoleForUser(String.valueOf(userId), String.valueOf(roleId)); - denyEnforcer.addRoleForUser(String.valueOf(userId), String.valueOf(roleId)); - if (loadedRoles.getIfPresent(roleId) != null) { - continue; - } - CompletableFuture<Void> loadRoleFuture = - CompletableFuture.supplyAsync( - () -> { - try { - return entityStore.get( - NameIdentifierUtil.ofRole(metalake, role.name()), - Entity.EntityType.ROLE, - RoleEntity.class); - } catch (Exception e) { - throw new RuntimeException("Failed to load role: " + role.name(), e); - } - }, - executor) - .thenAcceptAsync( - roleEntity -> { - loadPolicyByRoleEntity(roleEntity); - loadedRoles.put(roleId, true); - }, - executor); - loadRoleFutures.add(loadRoleFuture); - } - CompletableFuture.allOf(loadRoleFutures.toArray(new CompletableFuture[0])).join(); } catch (IOException e) { throw new RuntimeException(e); } + + // Register user-role associations in enforcers for all roles. + for (RoleEntity role : roleStubs) { + allowEnforcer.addRoleForUser(String.valueOf(userId), String.valueOf(role.id())); + denyEnforcer.addRoleForUser(String.valueOf(userId), String.valueOf(role.id())); + } + + // Collect stubs for roles whose policies have not yet been loaded into the enforcer. + List<RoleEntity> unloadedRoleStubs = + roleStubs.stream() + .filter(role -> loadedRoles.getIfPresent(role.id()) == null) + .collect(Collectors.toList()); + if (unloadedRoleStubs.isEmpty()) { + return; + } + + // Batch-fetch securable objects for all unloaded roles in a single query, + // eliminating the N+1 pattern of per-role entityStore.get() calls. + List<Long> unloadedRoleIds = + unloadedRoleStubs.stream().map(RoleEntity::id).collect(Collectors.toList()); + Map<Long, List<SecurableObject>> secObjsByRoleId = + RoleMetaService.batchListSecurableObjectsForRoles(unloadedRoleIds); + + for (RoleEntity stub : unloadedRoleStubs) { + List<SecurableObject> securableObjects = + secObjsByRoleId.getOrDefault(stub.id(), Collections.emptyList()); + RoleEntity fullRole = + RoleEntity.builder() + .withId(stub.id()) + .withName(stub.name()) + .withNamespace(stub.namespace()) + .withProperties(stub.properties()) + .withAuditInfo(stub.auditInfo()) + .withSecurableObjects(securableObjects) + .build(); + loadPolicyByRoleEntity(fullRole); + loadedRoles.put(stub.id(), true); + } }); } diff --git a/server-common/src/test/java/org/apache/gravitino/server/authorization/jcasbin/TestJcasbinAuthorizer.java b/server-common/src/test/java/org/apache/gravitino/server/authorization/jcasbin/TestJcasbinAuthorizer.java index 7bafa045db..dbca29ad0c 100644 --- a/server-common/src/test/java/org/apache/gravitino/server/authorization/jcasbin/TestJcasbinAuthorizer.java +++ b/server-common/src/test/java/org/apache/gravitino/server/authorization/jcasbin/TestJcasbinAuthorizer.java @@ -38,7 +38,9 @@ import java.io.IOException; import java.lang.reflect.Field; import java.security.Principal; import java.util.ArrayList; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.concurrent.Executor; import java.util.stream.Collectors; @@ -64,6 +66,7 @@ import org.apache.gravitino.server.ServerConfig; import org.apache.gravitino.server.authorization.MetadataIdConverter; import org.apache.gravitino.storage.relational.po.SecurableObjectPO; import org.apache.gravitino.storage.relational.service.OwnerMetaService; +import org.apache.gravitino.storage.relational.service.RoleMetaService; import org.apache.gravitino.storage.relational.utils.POConverters; import org.apache.gravitino.utils.NameIdentifierUtil; import org.apache.gravitino.utils.NamespaceUtil; @@ -98,6 +101,12 @@ public class TestJcasbinAuthorizer { private static SupportsRelationOperations supportsRelationOperations = mock(SupportsRelationOperations.class); + /** + * Shared map from role ID to securable objects, consulted by the {@link RoleMetaService} mock. + * Test methods must populate this map before exercising code paths that load new roles. + */ + private static final Map<Long, List<SecurableObject>> roleSecurableObjectsMap = new HashMap<>(); + private static MockedStatic<PrincipalUtils> principalUtilsMockedStatic; private static MockedStatic<GravitinoEnv> gravitinoEnvMockedStatic; @@ -106,12 +115,33 @@ public class TestJcasbinAuthorizer { private static MockedStatic<OwnerMetaService> ownerMetaServiceMockedStatic; + private static MockedStatic<RoleMetaService> roleMetaServiceMockedStatic; + private static JcasbinAuthorizer jcasbinAuthorizer; private static ObjectMapper objectMapper = new ObjectMapper(); @BeforeAll public static void setup() throws IOException { + // Pre-populate known constant roles so tests don't need to set them up individually. + roleSecurableObjectsMap.put(ALLOW_ROLE_ID, ImmutableList.of(getAllowSecurableObject())); + roleSecurableObjectsMap.put(DENY_ROLE_ID, ImmutableList.of(getDenySecurableObject())); + + // Mock RoleMetaService.batchListSecurableObjectsForRoles to avoid real DB access. + roleMetaServiceMockedStatic = mockStatic(RoleMetaService.class); + roleMetaServiceMockedStatic + .when(() -> RoleMetaService.batchListSecurableObjectsForRoles(any())) + .thenAnswer( + inv -> { + List<Long> ids = inv.getArgument(0); + com.google.common.collect.ImmutableMap.Builder<Long, List<SecurableObject>> result = + com.google.common.collect.ImmutableMap.builder(); + for (Long id : ids) { + result.put(id, roleSecurableObjectsMap.getOrDefault(id, ImmutableList.of())); + } + return result.build(); + }); + OwnerMetaService ownerMetaService = mock(OwnerMetaService.class); ownerMetaServiceMockedStatic = mockStatic(OwnerMetaService.class); ownerMetaServiceMockedStatic.when(OwnerMetaService::getInstance).thenReturn(ownerMetaService); @@ -164,6 +194,9 @@ public class TestJcasbinAuthorizer { if (gravitinoEnvMockedStatic != null) { gravitinoEnvMockedStatic.close(); } + if (roleMetaServiceMockedStatic != null) { + roleMetaServiceMockedStatic.close(); + } } @Test @@ -476,11 +509,7 @@ public class TestJcasbinAuthorizer { ImmutableList.of( buildManageGrantsSecurableObject( metalakeGrantRoleId, MetadataObject.Type.METALAKE, METALAKE))); - when(entityStore.get( - eq(NameIdentifierUtil.ofRole(METALAKE, metalakeGrantRole.name())), - eq(Entity.EntityType.ROLE), - eq(RoleEntity.class))) - .thenReturn(metalakeGrantRole); + roleSecurableObjectsMap.put(metalakeGrantRoleId, metalakeGrantRole.securableObjects()); when(supportsRelationOperations.listEntitiesByRelation( eq(SupportsRelationOperations.Type.ROLE_USER_REL), eq(userNameIdentifier), @@ -503,11 +532,7 @@ public class TestJcasbinAuthorizer { ImmutableList.of( buildManageGrantsSecurableObject( catalogGrantRoleId, MetadataObject.Type.CATALOG, "testCatalog"))); - when(entityStore.get( - eq(NameIdentifierUtil.ofRole(METALAKE, catalogGrantRole.name())), - eq(Entity.EntityType.ROLE), - eq(RoleEntity.class))) - .thenReturn(catalogGrantRole); + roleSecurableObjectsMap.put(catalogGrantRoleId, catalogGrantRole.securableObjects()); when(supportsRelationOperations.listEntitiesByRelation( eq(SupportsRelationOperations.Type.ROLE_USER_REL), eq(userNameIdentifier), @@ -536,11 +561,7 @@ public class TestJcasbinAuthorizer { tableGrantRoleId, MetadataObject.Type.TABLE, "testCatalog.testSchema.testTable"))); - when(entityStore.get( - eq(NameIdentifierUtil.ofRole(METALAKE, tableGrantRole.name())), - eq(Entity.EntityType.ROLE), - eq(RoleEntity.class))) - .thenReturn(tableGrantRole); + roleSecurableObjectsMap.put(tableGrantRoleId, tableGrantRole.securableObjects()); when(supportsRelationOperations.listEntitiesByRelation( eq(SupportsRelationOperations.Type.ROLE_USER_REL), eq(userNameIdentifier),
