fix: avoid repeated grantee scans in ListPolicy (#51404)

Related: #50236

- Reuse the initially loaded grantee keys/values when ListPolicy checks
legacy grantee ID ownership.
- Avoid reloading the full grantee-privileges tree once per legacy
grant.
- Add a regression test that verifies ListPolicy performs only one full
grantee-prefix load for legacy grants.

Signed-off-by: shaoting-huang <shaoting.huang@zilliz.com>
This commit is contained in:
sthuang
2026-07-16 17:40:37 +08:00
committed by GitHub
parent 0fe1a5ba6f
commit 6c7300f15c
2 changed files with 65 additions and 2 deletions
+15 -2
View File
@@ -1491,11 +1491,24 @@ func shouldRemoveGranteeIDSubtree(ctx context.Context, granteePrefix string, idS
}
func (kc *Catalog) loadGranteeIDPrefix(ctx context.Context, tenant string, granteeKey string, idStr string) ([]string, []string, string, error) {
return kc.loadGranteeIDPrefixWithLoadedGrantees(ctx, tenant, granteeKey, idStr, nil, nil)
}
func (kc *Catalog) loadGranteeIDPrefixWithLoadedGrantees(ctx context.Context, tenant string, granteeKey string, idStr string, granteeKeys []string, granteeValues []string) ([]string, []string, string, error) {
var firstPrefix string
newID := crypto.GranteeID(granteeKey)
if isLegacyGranteeID(idStr) && idStr != newID {
granteeIDKey := funcutil.HandleTenantForEtcdPrefix(GranteeIDPrefix, tenant, idStr)
otherGranteeKey, err := kc.findOtherGranteeWithID(ctx, tenant, granteeKey, idStr)
var (
otherGranteeKey string
err error
)
if granteeKeys != nil {
granteePrefix := funcutil.HandleTenantForEtcdPrefix(GranteePrefix, tenant)
otherGranteeKey, err = findOtherGranteeWithIDFromKeys(ctx, granteePrefix, granteeKey, idStr, granteeKeys, granteeValues)
} else {
otherGranteeKey, err = kc.findOtherGranteeWithID(ctx, tenant, granteeKey, idStr)
}
if err != nil {
return nil, nil, granteeIDKey, err
}
@@ -1981,7 +1994,7 @@ func (kc *Catalog) ListPolicy(ctx context.Context, tenant string) ([]*milvuspb.G
continue
}
logicalGranteeKey := fmt.Sprintf("%s/%s/%s/%s", GranteePrefix, grantInfos[0], grantInfos[1], grantInfos[2])
idKeys, _, granteeIDKey, err := kc.loadGranteeIDPrefix(ctx, tenant, logicalGranteeKey, values[i])
idKeys, _, granteeIDKey, err := kc.loadGranteeIDPrefixWithLoadedGrantees(ctx, tenant, logicalGranteeKey, values[i], keys, values)
if err != nil {
mlog.Error(ctx, "fail to load the grantee ids", mlog.String("key", granteeIDKey), mlog.Err(err))
return []*milvuspb.GrantEntity{}, err
@@ -3254,6 +3254,56 @@ func TestRBACGrantSharedLegacyGranteeIDMigrationFailsClosed(t *testing.T) {
assert.Equal(t, "donor-user", donorUser)
}
func TestRBACListPolicyLegacyGranteeIDReusesLoadedGrantKeys(t *testing.T) {
ctx := context.Background()
tenant := util.DefaultTenant
kvmock := mocks.NewTxnKV(t)
c := NewCatalog(kvmock)
granteePrefix := funcutil.HandleTenantForEtcdPrefix(GranteePrefix, tenant)
granteeIDPrefix := funcutil.HandleTenantForEtcdPrefix(GranteeIDPrefix, tenant)
objectType := commonpb.ObjectType_Collection.String()
firstLogicalKey := fmt.Sprintf("%s/%s/%s/%s", GranteePrefix, "role1", objectType, funcutil.CombineObjectName(util.DefaultDBName, "coll1"))
secondLogicalKey := fmt.Sprintf("%s/%s/%s/%s", GranteePrefix, "role2", objectType, funcutil.CombineObjectName(util.DefaultDBName, "coll2"))
firstEtcdKey := fmt.Sprintf("%s%s/%s/%s", granteePrefix, "role1", objectType, funcutil.CombineObjectName(util.DefaultDBName, "coll1"))
secondEtcdKey := fmt.Sprintf("%s%s/%s/%s", granteePrefix, "role2", objectType, funcutil.CombineObjectName(util.DefaultDBName, "coll2"))
firstLegacyID := crypto.MD5(firstLogicalKey)
secondLegacyID := crypto.MD5(secondLogicalKey)
var granteePrefixLoads atomic.Int32
kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return(
func(ctx context.Context, key string) []string {
switch {
case key == granteePrefix:
granteePrefixLoads.Inc()
return []string{firstEtcdKey, secondEtcdKey}
case strings.HasPrefix(key, granteeIDPrefix):
return []string{key + "PrivilegeLoad"}
default:
return nil
}
},
func(ctx context.Context, key string) []string {
switch {
case key == granteePrefix:
return []string{firstLegacyID, secondLegacyID}
case strings.HasPrefix(key, granteeIDPrefix):
return []string{"root"}
default:
return nil
}
},
func(ctx context.Context, key string) error {
return nil
},
)
policies, err := c.ListPolicy(ctx, tenant)
require.NoError(t, err)
require.Len(t, policies, 2)
assert.Equal(t, int32(1), granteePrefixLoads.Load(), "ListPolicy should reuse the initially loaded grantee keys for legacy ID collision checks")
}
func TestRBACGrantMigrationIgnoresUnreferencedComputedLegacyID(t *testing.T) {
ctx := context.Background()
etcdCli, _ := etcd.GetEtcdClient(