Skip to content

Commit

Permalink
enhance: Add in-memory cache for casbin enforcer result (milvus-io#35271
Browse files Browse the repository at this point in the history
)

See also milvus-io#35270

---------

Signed-off-by: Congqi Xia <[email protected]>
  • Loading branch information
congqixia authored Aug 5, 2024
1 parent a62118c commit a2b5175
Show file tree
Hide file tree
Showing 6 changed files with 190 additions and 15 deletions.
28 changes: 20 additions & 8 deletions internal/datacoord/compaction_policy_clustering_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -213,7 +213,8 @@ func (s *ClusteringCompactionPolicySuite) TestTimeIntervalLogic() {
size: *atomic.NewInt64(1024 * 1024 * 1024 * 10),
},
}, true},
{"very recent partition stats and enough new data",
{
"very recent partition stats and enough new data",
[]*datapb.PartitionStatsInfo{
{
CollectionID: collectionID,
Expand All @@ -228,8 +229,11 @@ func (s *ClusteringCompactionPolicySuite) TestTimeIntervalLogic() {
{
size: *atomic.NewInt64(1024 * 1024 * 1024 * 10),
},
}, false},
{"very old partition stats and not enough new data",
},
false,
},
{
"very old partition stats and not enough new data",
[]*datapb.PartitionStatsInfo{
{
CollectionID: collectionID,
Expand All @@ -244,8 +248,11 @@ func (s *ClusteringCompactionPolicySuite) TestTimeIntervalLogic() {
{
size: *atomic.NewInt64(1024),
},
}, true},
{"partition stats and enough new data",
},
true,
},
{
"partition stats and enough new data",
[]*datapb.PartitionStatsInfo{
{
CollectionID: collectionID,
Expand All @@ -262,8 +269,11 @@ func (s *ClusteringCompactionPolicySuite) TestTimeIntervalLogic() {
SegmentInfo: &datapb.SegmentInfo{ID: 9999},
size: *atomic.NewInt64(1024 * 1024 * 1024 * 10),
},
}, true},
{"partition stats and not enough new data",
},
true,
},
{
"partition stats and not enough new data",
[]*datapb.PartitionStatsInfo{
{
CollectionID: collectionID,
Expand All @@ -280,7 +290,9 @@ func (s *ClusteringCompactionPolicySuite) TestTimeIntervalLogic() {
SegmentInfo: &datapb.SegmentInfo{ID: 9999},
size: *atomic.NewInt64(1024),
},
}, false},
},
false,
},
}

for _, test := range tests {
Expand Down
11 changes: 4 additions & 7 deletions internal/datacoord/compaction_trigger_v2_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,21 +5,18 @@ import (
"strconv"
"testing"

"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"

"github.com/milvus-io/milvus/pkg/util/paramtable"

"github.com/milvus-io/milvus/internal/metastore/model"
"github.com/milvus-io/milvus/pkg/common"

"github.com/samber/lo"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/suite"
"go.uber.org/zap"

"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/metastore/model"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/pkg/common"
"github.com/milvus-io/milvus/pkg/log"
"github.com/milvus-io/milvus/pkg/util/paramtable"
)

func TestCompactionTriggerManagerSuite(t *testing.T) {
Expand Down
1 change: 1 addition & 0 deletions internal/proxy/meta_cache.go
Original file line number Diff line number Diff line change
Expand Up @@ -1074,6 +1074,7 @@ func (m *MetaCache) RefreshPolicyInfo(op typeutil.CacheOp) (err error) {
if le != nil {
log.Error("failed to load policy after RefreshPolicyInfo", zap.Error(le))
}
CleanPrivilegeCache()
}
}()
if op.OpType != typeutil.CacheRefresh {
Expand Down
84 changes: 84 additions & 0 deletions internal/proxy/privilege_cache.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,84 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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 proxy

import (
"fmt"
"sync"

"go.uber.org/atomic"

"github.com/milvus-io/milvus/pkg/util/typeutil"
)

var (
priCacheInitOnce sync.Once
priCacheMut sync.RWMutex
priCache *PrivilegeCache
ver atomic.Int64
)

func getPriCache() *PrivilegeCache {
priCacheMut.RLock()
c := priCache
priCacheMut.RUnlock()

if c == nil {
priCacheInitOnce.Do(func() {
priCacheMut.Lock()
defer priCacheMut.Unlock()
c = &PrivilegeCache{
version: ver.Inc(),
values: typeutil.ConcurrentMap[string, bool]{},
}
priCache = c
})
}

return c
}

func CleanPrivilegeCache() {
priCacheMut.Lock()
defer priCacheMut.Unlock()
priCache = &PrivilegeCache{
version: ver.Inc(),
values: typeutil.ConcurrentMap[string, bool]{},
}
}

func GetPrivilegeCache(roleName, object, objectPrivilege string) (isPermit, cached bool, version int64) {
key := fmt.Sprintf("%s_%s_%s", roleName, object, objectPrivilege)
c := getPriCache()
isPermit, cached = c.values.Get(key)
return isPermit, cached, c.version
}

func SetPrivilegeCache(roleName, object, objectPrivilege string, isPermit bool, version int64) {
key := fmt.Sprintf("%s_%s_%s", roleName, object, objectPrivilege)
c := getPriCache()
if c.version == version {
c.values.Insert(key, isPermit)
}
}

// PrivilegeCache is a cache for privilege enforce result
// version provides version control when any policy updates
type PrivilegeCache struct {
values typeutil.ConcurrentMap[string, bool]
version int64
}
76 changes: 76 additions & 0 deletions internal/proxy/privilege_cache_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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 proxy

import (
"testing"

"github.com/stretchr/testify/suite"
)

type PrivilegeCacheSuite struct {
suite.Suite
}

func (s *PrivilegeCacheSuite) TearDownTest() {
CleanPrivilegeCache()
}

func (s *PrivilegeCacheSuite) TestGetPrivilege() {
// get current version
_, _, version := GetPrivilegeCache("", "", "")
SetPrivilegeCache("test-role", "test-object", "read", true, version)
SetPrivilegeCache("test-role", "test-object", "delete", false, version)

type testCase struct {
tag string
input [3]string
expectIsPermit bool
expectExists bool
}

testCases := []testCase{
{tag: "exist_true", input: [3]string{"test-role", "test-object", "read"}, expectIsPermit: true, expectExists: true},
{tag: "exist_false", input: [3]string{"test-role", "test-object", "delete"}, expectIsPermit: false, expectExists: true},
{tag: "not_exist", input: [3]string{"guest", "test-object", "delete"}, expectIsPermit: false, expectExists: false},
}

for _, tc := range testCases {
s.Run(tc.tag, func() {
isPermit, exists, _ := GetPrivilegeCache(tc.input[0], tc.input[1], tc.input[2])
s.Equal(tc.expectIsPermit, isPermit)
s.Equal(tc.expectExists, exists)
})
}
}

func (s *PrivilegeCacheSuite) TestSetPrivilegeVersion() {
// get current version
_, _, version := GetPrivilegeCache("", "", "")
CleanPrivilegeCache()

SetPrivilegeCache("test-role", "test-object", "read", true, version)

isPermit, exists, nextVersion := GetPrivilegeCache("test-role", "test-object", "read")
s.False(isPermit)
s.False(exists)
s.NotEqual(version, nextVersion)
}

func TestPrivilegeSuite(t *testing.T) {
suite.Run(t, new(PrivilegeCacheSuite))
}
5 changes: 5 additions & 0 deletions internal/proxy/privilege_interceptor.go
Original file line number Diff line number Diff line change
Expand Up @@ -135,10 +135,15 @@ func PrivilegeInterceptor(ctx context.Context, req interface{}) (context.Context
for _, roleName := range roleNames {
permitFunc := func(resName string) (bool, error) {
object := funcutil.PolicyForResource(dbName, objectType, resName)
isPermit, cached, version := GetPrivilegeCache(roleName, object, objectPrivilege)
if cached {
return isPermit, nil
}
isPermit, err := e.Enforce(roleName, object, objectPrivilege)
if err != nil {
return false, err
}
SetPrivilegeCache(roleName, object, objectPrivilege, isPermit, version)
return isPermit, nil
}

Expand Down

0 comments on commit a2b5175

Please sign in to comment.