milvus/internal/proxy/meta_cache_test.go

883 lines
29 KiB
Go
Raw Normal View History

// 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 (
"context"
"fmt"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/cockroachdb/errors"
"github.com/milvus-io/milvus/internal/util/funcutil"
"github.com/milvus-io/milvus/internal/util/paramtable"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus-proto/go-api/commonpb"
"github.com/milvus-io/milvus-proto/go-api/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/schemapb"
"github.com/milvus-io/milvus/internal/log"
"github.com/milvus-io/milvus/internal/proto/internalpb"
"github.com/milvus-io/milvus/internal/proto/querypb"
"github.com/milvus-io/milvus/internal/proto/rootcoordpb"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/internal/util/crypto"
"github.com/milvus-io/milvus/internal/util/typeutil"
)
2021-06-21 17:28:03 +08:00
type MockRootCoordClientInterface struct {
types.RootCoord
Error bool
AccessCount int32
listPolicy func(ctx context.Context, in *internalpb.ListPolicyRequest) (*internalpb.ListPolicyResponse, error)
}
func (m *MockRootCoordClientInterface) IncAccessCount() {
atomic.AddInt32(&m.AccessCount, 1)
}
func (m *MockRootCoordClientInterface) GetAccessCount() int {
ret := atomic.LoadInt32(&m.AccessCount)
return int(ret)
}
2021-06-21 17:28:03 +08:00
func (m *MockRootCoordClientInterface) ShowPartitions(ctx context.Context, in *milvuspb.ShowPartitionsRequest) (*milvuspb.ShowPartitionsResponse, error) {
if m.Error {
return nil, errors.New("mocked error")
}
if in.CollectionName == "collection1" {
return &milvuspb.ShowPartitionsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
PartitionIDs: []typeutil.UniqueID{1, 2},
CreatedTimestamps: []uint64{100, 200},
CreatedUtcTimestamps: []uint64{100, 200},
PartitionNames: []string{"par1", "par2"},
}, nil
}
if in.CollectionName == "collection2" {
return &milvuspb.ShowPartitionsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
PartitionIDs: []typeutil.UniqueID{3, 4},
CreatedTimestamps: []uint64{201, 202},
CreatedUtcTimestamps: []uint64{201, 202},
PartitionNames: []string{"par1", "par2"},
}, nil
}
if in.CollectionName == "errorCollection" {
return &milvuspb.ShowPartitionsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
PartitionIDs: []typeutil.UniqueID{5, 6},
CreatedTimestamps: []uint64{201},
CreatedUtcTimestamps: []uint64{201},
PartitionNames: []string{"par1", "par2"},
}, nil
}
return &milvuspb.ShowPartitionsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
},
PartitionIDs: []typeutil.UniqueID{},
CreatedTimestamps: []uint64{},
CreatedUtcTimestamps: []uint64{},
PartitionNames: []string{},
}, nil
}
2021-06-21 17:28:03 +08:00
func (m *MockRootCoordClientInterface) DescribeCollection(ctx context.Context, in *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) {
if m.Error {
return nil, errors.New("mocked error")
}
m.IncAccessCount()
if in.CollectionName == "collection1" || in.CollectionID == 1 {
return &milvuspb.DescribeCollectionResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
CollectionID: typeutil.UniqueID(1),
Schema: &schemapb.CollectionSchema{
AutoID: true,
Name: "collection1",
},
}, nil
}
if in.CollectionName == "collection2" || in.CollectionID == 2 {
return &milvuspb.DescribeCollectionResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
CollectionID: typeutil.UniqueID(2),
Schema: &schemapb.CollectionSchema{
AutoID: true,
Name: "collection2",
},
}, nil
}
if in.CollectionName == "errorCollection" {
return &milvuspb.DescribeCollectionResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
CollectionID: typeutil.UniqueID(3),
Schema: &schemapb.CollectionSchema{
AutoID: true,
},
}, nil
}
err := fmt.Errorf("can't find collection: " + in.CollectionName)
return &milvuspb.DescribeCollectionResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_CollectionNotExists,
Reason: "describe collection failed: " + err.Error(),
},
Schema: nil,
}, nil
}
func (m *MockRootCoordClientInterface) GetCredential(ctx context.Context, req *rootcoordpb.GetCredentialRequest) (*rootcoordpb.GetCredentialResponse, error) {
if m.Error {
return nil, errors.New("mocked error")
}
m.IncAccessCount()
if req.Username == "mockUser" {
encryptedPassword, _ := crypto.PasswordEncrypt("mockPass")
return &rootcoordpb.GetCredentialResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
Username: "mockUser",
Password: encryptedPassword,
}, nil
}
err := fmt.Errorf("can't find credential: " + req.Username)
return nil, err
}
func (m *MockRootCoordClientInterface) ListCredUsers(ctx context.Context, req *milvuspb.ListCredUsersRequest) (*milvuspb.ListCredUsersResponse, error) {
if m.Error {
return nil, errors.New("mocked error")
}
return &milvuspb.ListCredUsersResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
Usernames: []string{"mockUser"},
}, nil
}
func (m *MockRootCoordClientInterface) ListPolicy(ctx context.Context, in *internalpb.ListPolicyRequest) (*internalpb.ListPolicyResponse, error) {
if m.listPolicy != nil {
return m.listPolicy(ctx, in)
}
return &internalpb.ListPolicyResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
}, nil
}
// Simulate the cache path and the
func TestMetaCache_GetCollection(t *testing.T) {
ctx := context.Background()
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
mgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, queryCoord, mgr)
assert.Nil(t, err)
id, err := globalMetaCache.GetCollectionID(ctx, "collection1")
assert.Nil(t, err)
assert.Equal(t, id, typeutil.UniqueID(1))
assert.Equal(t, rootCoord.GetAccessCount(), 1)
// should'nt be accessed to remote root coord.
schema, err := globalMetaCache.GetCollectionSchema(ctx, "collection1")
assert.Equal(t, rootCoord.GetAccessCount(), 1)
assert.Nil(t, err)
assert.Equal(t, schema, &schemapb.CollectionSchema{
AutoID: true,
Fields: []*schemapb.FieldSchema{},
Name: "collection1",
})
id, err = globalMetaCache.GetCollectionID(ctx, "collection2")
assert.Equal(t, rootCoord.GetAccessCount(), 2)
assert.Nil(t, err)
assert.Equal(t, id, typeutil.UniqueID(2))
schema, err = globalMetaCache.GetCollectionSchema(ctx, "collection2")
assert.Equal(t, rootCoord.GetAccessCount(), 2)
assert.Nil(t, err)
assert.Equal(t, schema, &schemapb.CollectionSchema{
AutoID: true,
Fields: []*schemapb.FieldSchema{},
Name: "collection2",
})
// test to get from cache, this should trigger root request
id, err = globalMetaCache.GetCollectionID(ctx, "collection1")
assert.Equal(t, rootCoord.GetAccessCount(), 2)
assert.Nil(t, err)
assert.Equal(t, id, typeutil.UniqueID(1))
schema, err = globalMetaCache.GetCollectionSchema(ctx, "collection1")
assert.Equal(t, rootCoord.GetAccessCount(), 2)
assert.Nil(t, err)
assert.Equal(t, schema, &schemapb.CollectionSchema{
AutoID: true,
Fields: []*schemapb.FieldSchema{},
Name: "collection1",
})
}
func TestMetaCache_GetCollectionName(t *testing.T) {
ctx := context.Background()
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
mgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, queryCoord, mgr)
assert.Nil(t, err)
collection, err := globalMetaCache.GetCollectionName(ctx, 1)
assert.Nil(t, err)
assert.Equal(t, collection, "collection1")
assert.Equal(t, rootCoord.GetAccessCount(), 1)
// should'nt be accessed to remote root coord.
schema, err := globalMetaCache.GetCollectionSchema(ctx, "collection1")
assert.Equal(t, rootCoord.GetAccessCount(), 1)
assert.Nil(t, err)
assert.Equal(t, schema, &schemapb.CollectionSchema{
AutoID: true,
Fields: []*schemapb.FieldSchema{},
Name: "collection1",
})
collection, err = globalMetaCache.GetCollectionName(ctx, 1)
assert.Equal(t, rootCoord.GetAccessCount(), 1)
assert.Nil(t, err)
assert.Equal(t, collection, "collection1")
schema, err = globalMetaCache.GetCollectionSchema(ctx, "collection2")
assert.Equal(t, rootCoord.GetAccessCount(), 2)
assert.Nil(t, err)
assert.Equal(t, schema, &schemapb.CollectionSchema{
AutoID: true,
Fields: []*schemapb.FieldSchema{},
Name: "collection2",
})
// test to get from cache, this should trigger root request
collection, err = globalMetaCache.GetCollectionName(ctx, 1)
assert.Equal(t, rootCoord.GetAccessCount(), 2)
assert.Nil(t, err)
assert.Equal(t, collection, "collection1")
schema, err = globalMetaCache.GetCollectionSchema(ctx, "collection1")
assert.Equal(t, rootCoord.GetAccessCount(), 2)
assert.Nil(t, err)
assert.Equal(t, schema, &schemapb.CollectionSchema{
AutoID: true,
Fields: []*schemapb.FieldSchema{},
Name: "collection1",
})
}
func TestMetaCache_GetCollectionFailure(t *testing.T) {
ctx := context.Background()
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
mgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, queryCoord, mgr)
assert.Nil(t, err)
rootCoord.Error = true
schema, err := globalMetaCache.GetCollectionSchema(ctx, "collection1")
assert.NotNil(t, err)
assert.Nil(t, schema)
rootCoord.Error = false
schema, err = globalMetaCache.GetCollectionSchema(ctx, "collection1")
assert.Nil(t, err)
assert.Equal(t, schema, &schemapb.CollectionSchema{
AutoID: true,
Fields: []*schemapb.FieldSchema{},
Name: "collection1",
})
rootCoord.Error = true
// should be cached with no error
assert.Nil(t, err)
assert.Equal(t, schema, &schemapb.CollectionSchema{
AutoID: true,
Fields: []*schemapb.FieldSchema{},
Name: "collection1",
})
}
func TestMetaCache_GetNonExistCollection(t *testing.T) {
ctx := context.Background()
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
mgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, queryCoord, mgr)
assert.Nil(t, err)
id, err := globalMetaCache.GetCollectionID(ctx, "collection3")
assert.NotNil(t, err)
assert.Equal(t, id, int64(0))
schema, err := globalMetaCache.GetCollectionSchema(ctx, "collection3")
assert.NotNil(t, err)
assert.Nil(t, schema)
}
func TestMetaCache_GetPartitionID(t *testing.T) {
ctx := context.Background()
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
mgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, queryCoord, mgr)
assert.Nil(t, err)
id, err := globalMetaCache.GetPartitionID(ctx, "collection1", "par1")
assert.Nil(t, err)
assert.Equal(t, id, typeutil.UniqueID(1))
id, err = globalMetaCache.GetPartitionID(ctx, "collection1", "par2")
assert.Nil(t, err)
assert.Equal(t, id, typeutil.UniqueID(2))
id, err = globalMetaCache.GetPartitionID(ctx, "collection2", "par1")
assert.Nil(t, err)
assert.Equal(t, id, typeutil.UniqueID(3))
id, err = globalMetaCache.GetPartitionID(ctx, "collection2", "par2")
assert.Nil(t, err)
assert.Equal(t, id, typeutil.UniqueID(4))
}
func TestMetaCache_ConcurrentTest1(t *testing.T) {
ctx := context.Background()
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
mgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, queryCoord, mgr)
assert.Nil(t, err)
var wg sync.WaitGroup
cnt := 100
getCollectionCacheFunc := func(wg *sync.WaitGroup) {
defer wg.Done()
for i := 0; i < cnt; i++ {
//GetCollectionSchema will never fail
schema, err := globalMetaCache.GetCollectionSchema(ctx, "collection1")
assert.Nil(t, err)
assert.Equal(t, schema, &schemapb.CollectionSchema{
AutoID: true,
Fields: []*schemapb.FieldSchema{},
Name: "collection1",
})
time.Sleep(10 * time.Millisecond)
}
}
getPartitionCacheFunc := func(wg *sync.WaitGroup) {
defer wg.Done()
for i := 0; i < cnt; i++ {
//GetPartitions may fail
globalMetaCache.GetPartitions(ctx, "collection1")
time.Sleep(10 * time.Millisecond)
}
}
invalidCacheFunc := func(wg *sync.WaitGroup) {
defer wg.Done()
for i := 0; i < cnt; i++ {
//periodically invalid collection cache
globalMetaCache.RemoveCollection(ctx, "collection1")
time.Sleep(10 * time.Millisecond)
}
}
wg.Add(1)
go getCollectionCacheFunc(&wg)
wg.Add(1)
go invalidCacheFunc(&wg)
wg.Add(1)
go getPartitionCacheFunc(&wg)
wg.Wait()
}
func TestMetaCache_GetPartitionError(t *testing.T) {
ctx := context.Background()
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
mgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, queryCoord, mgr)
assert.Nil(t, err)
// Test the case where ShowPartitionsResponse is not aligned
id, err := globalMetaCache.GetPartitionID(ctx, "errorCollection", "par1")
assert.NotNil(t, err)
log.Debug(err.Error())
assert.Equal(t, id, typeutil.UniqueID(0))
partitions, err2 := globalMetaCache.GetPartitions(ctx, "errorCollection")
assert.NotNil(t, err2)
log.Debug(err.Error())
assert.Equal(t, len(partitions), 0)
// Test non existed tables
id, err = globalMetaCache.GetPartitionID(ctx, "nonExisted", "par1")
assert.NotNil(t, err)
log.Debug(err.Error())
assert.Equal(t, id, typeutil.UniqueID(0))
// Test non existed partition
id, err = globalMetaCache.GetPartitionID(ctx, "collection1", "par3")
assert.NotNil(t, err)
log.Debug(err.Error())
assert.Equal(t, id, typeutil.UniqueID(0))
}
func TestMetaCache_GetShards(t *testing.T) {
var (
ctx = context.Background()
collectionName = "collection1"
)
rootCoord := &MockRootCoordClientInterface{}
qc := getQueryCoord()
qc.EXPECT().Init().Return(nil)
shardMgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, qc, shardMgr)
require.Nil(t, err)
qc.Init()
qc.Start()
defer qc.Stop()
t.Run("No collection in meta cache", func(t *testing.T) {
shards, err := globalMetaCache.GetShards(ctx, true, "non-exists")
assert.Error(t, err)
assert.Empty(t, shards)
})
t.Run("without shardLeaders in collection info invalid shardLeaders", func(t *testing.T) {
qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: "not implemented",
},
}, nil).Times(1)
qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
}, nil)
shards, err := globalMetaCache.GetShards(ctx, false, collectionName)
assert.Error(t, err)
assert.Empty(t, shards)
})
t.Run("without shardLeaders in collection info", func(t *testing.T) {
qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
Shards: []*querypb.ShardLeadersList{
{
ChannelName: "channel-1",
NodeIds: []int64{1, 2, 3},
NodeAddrs: []string{"localhost:9000", "localhost:9001", "localhost:9002"},
},
},
}, nil)
qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
}, nil)
shards, err := globalMetaCache.GetShards(ctx, true, collectionName)
assert.NoError(t, err)
assert.NotEmpty(t, shards)
assert.Equal(t, 1, len(shards))
assert.Equal(t, 3, len(shards["channel-1"]))
// get from cache
qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: "not implemented",
},
}, nil)
shards, err = globalMetaCache.GetShards(ctx, true, collectionName)
assert.NoError(t, err)
assert.NotEmpty(t, shards)
assert.Equal(t, 1, len(shards))
assert.Equal(t, 3, len(shards["channel-1"]))
})
}
func TestMetaCache_ClearShards(t *testing.T) {
var (
ctx = context.TODO()
collectionName = "collection1"
)
rootCoord := &MockRootCoordClientInterface{}
qc := getQueryCoord()
qc.EXPECT().Init().Return(nil)
mgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, qc, mgr)
require.Nil(t, err)
qc.Init()
qc.Start()
defer qc.Stop()
t.Run("Clear with no collection info", func(t *testing.T) {
globalMetaCache.ClearShards("collection_not_exist")
})
t.Run("Clear valid collection empty cache", func(t *testing.T) {
globalMetaCache.ClearShards(collectionName)
})
t.Run("Clear valid collection valid cache", func(t *testing.T) {
qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
Shards: []*querypb.ShardLeadersList{
{
ChannelName: "channel-1",
NodeIds: []int64{1, 2, 3},
NodeAddrs: []string{"localhost:9000", "localhost:9001", "localhost:9002"},
},
},
}, nil).Times(1)
qc.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
}, nil)
shards, err := globalMetaCache.GetShards(ctx, true, collectionName)
require.NoError(t, err)
require.NotEmpty(t, shards)
require.Equal(t, 1, len(shards))
require.Equal(t, 3, len(shards["channel-1"]))
globalMetaCache.ClearShards(collectionName)
qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: "not implemented",
},
}, nil)
shards, err = globalMetaCache.GetShards(ctx, true, collectionName)
assert.Error(t, err)
assert.Empty(t, shards)
})
}
func TestMetaCache_PolicyInfo(t *testing.T) {
client := &MockRootCoordClientInterface{}
qc := &types.MockQueryCoord{}
mgr := newShardClientMgr()
t.Run("InitMetaCache", func(t *testing.T) {
client.listPolicy = func(ctx context.Context, in *internalpb.ListPolicyRequest) (*internalpb.ListPolicyResponse, error) {
return nil, fmt.Errorf("mock error")
}
err := InitMetaCache(context.Background(), client, qc, mgr)
assert.NotNil(t, err)
client.listPolicy = func(ctx context.Context, in *internalpb.ListPolicyRequest) (*internalpb.ListPolicyResponse, error) {
return &internalpb.ListPolicyResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
PolicyInfos: []string{"policy1", "policy2", "policy3"},
}, nil
}
err = InitMetaCache(context.Background(), client, qc, mgr)
assert.Nil(t, err)
})
t.Run("GetPrivilegeInfo", func(t *testing.T) {
client.listPolicy = func(ctx context.Context, in *internalpb.ListPolicyRequest) (*internalpb.ListPolicyResponse, error) {
return &internalpb.ListPolicyResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
PolicyInfos: []string{"policy1", "policy2", "policy3"},
UserRoles: []string{funcutil.EncodeUserRoleCache("foo", "role1"), funcutil.EncodeUserRoleCache("foo", "role2"), funcutil.EncodeUserRoleCache("foo2", "role2")},
}, nil
}
err := InitMetaCache(context.Background(), client, qc, mgr)
assert.Nil(t, err)
policyInfos := globalMetaCache.GetPrivilegeInfo(context.Background())
assert.Equal(t, 3, len(policyInfos))
roles := globalMetaCache.GetUserRole("foo")
assert.Equal(t, 2, len(roles))
})
t.Run("GetPrivilegeInfo", func(t *testing.T) {
client.listPolicy = func(ctx context.Context, in *internalpb.ListPolicyRequest) (*internalpb.ListPolicyResponse, error) {
return &internalpb.ListPolicyResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
PolicyInfos: []string{"policy1", "policy2", "policy3"},
UserRoles: []string{funcutil.EncodeUserRoleCache("foo", "role1"), funcutil.EncodeUserRoleCache("foo", "role2"), funcutil.EncodeUserRoleCache("foo2", "role2")},
}, nil
}
err := InitMetaCache(context.Background(), client, qc, mgr)
assert.Nil(t, err)
err = globalMetaCache.RefreshPolicyInfo(typeutil.CacheOp{OpType: typeutil.CacheGrantPrivilege, OpKey: "policyX"})
assert.Nil(t, err)
policyInfos := globalMetaCache.GetPrivilegeInfo(context.Background())
assert.Equal(t, 4, len(policyInfos))
err = globalMetaCache.RefreshPolicyInfo(typeutil.CacheOp{OpType: typeutil.CacheRevokePrivilege, OpKey: "policyX"})
assert.Nil(t, err)
policyInfos = globalMetaCache.GetPrivilegeInfo(context.Background())
assert.Equal(t, 3, len(policyInfos))
err = globalMetaCache.RefreshPolicyInfo(typeutil.CacheOp{OpType: typeutil.CacheAddUserToRole, OpKey: funcutil.EncodeUserRoleCache("foo", "role3")})
assert.Nil(t, err)
roles := globalMetaCache.GetUserRole("foo")
assert.Equal(t, 3, len(roles))
err = globalMetaCache.RefreshPolicyInfo(typeutil.CacheOp{OpType: typeutil.CacheRemoveUserFromRole, OpKey: funcutil.EncodeUserRoleCache("foo", "role3")})
assert.Nil(t, err)
roles = globalMetaCache.GetUserRole("foo")
assert.Equal(t, 2, len(roles))
err = globalMetaCache.RefreshPolicyInfo(typeutil.CacheOp{OpType: typeutil.CacheGrantPrivilege, OpKey: ""})
assert.NotNil(t, err)
err = globalMetaCache.RefreshPolicyInfo(typeutil.CacheOp{OpType: 100, OpKey: "policyX"})
assert.NotNil(t, err)
})
}
func TestMetaCache_LoadCache(t *testing.T) {
ctx := context.Background()
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
mgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, queryCoord, mgr)
assert.Nil(t, err)
qcCounter := 0
queryCoord.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
CollectionIDs: []UniqueID{1, 2},
InMemoryPercentages: []int64{100, 50},
}, nil).Run(func(ctx context.Context, req *querypb.ShowCollectionsRequest) {
qcCounter++
})
t.Run("test IsCollectionLoaded", func(t *testing.T) {
info, err := globalMetaCache.GetCollectionInfo(ctx, "collection1")
assert.NoError(t, err)
assert.True(t, info.isLoaded)
// no collectionInfo of collection1, should access RootCoord
assert.Equal(t, rootCoord.GetAccessCount(), 1)
// not loaded, should access QueryCoord
assert.Equal(t, qcCounter, 1)
info, err = globalMetaCache.GetCollectionInfo(ctx, "collection1")
assert.NoError(t, err)
assert.True(t, info.isLoaded)
// shouldn't access QueryCoord or RootCoord again
assert.Equal(t, rootCoord.GetAccessCount(), 1)
assert.Equal(t, qcCounter, 1)
// test collection2 not fully loaded
info, err = globalMetaCache.GetCollectionInfo(ctx, "collection2")
assert.NoError(t, err)
assert.False(t, info.isLoaded)
// no collectionInfo of collection2, should access RootCoord
assert.Equal(t, rootCoord.GetAccessCount(), 2)
// not loaded, should access QueryCoord
assert.Equal(t, qcCounter, 2)
})
t.Run("test RemoveCollectionLoadCache", func(t *testing.T) {
globalMetaCache.RemoveCollection(ctx, "collection1")
info, err := globalMetaCache.GetCollectionInfo(ctx, "collection1")
assert.NoError(t, err)
assert.True(t, info.isLoaded)
// should access QueryCoord
assert.Equal(t, qcCounter, 3)
})
}
func TestMetaCache_RemoveCollection(t *testing.T) {
ctx := context.Background()
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
shardMgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, queryCoord, shardMgr)
assert.Nil(t, err)
queryCoord.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
CollectionIDs: []UniqueID{1, 2},
InMemoryPercentages: []int64{100, 50},
}, nil)
info, err := globalMetaCache.GetCollectionInfo(ctx, "collection1")
assert.NoError(t, err)
assert.True(t, info.isLoaded)
// no collectionInfo of collection1, should access RootCoord
assert.Equal(t, rootCoord.GetAccessCount(), 1)
info, err = globalMetaCache.GetCollectionInfo(ctx, "collection1")
assert.NoError(t, err)
assert.True(t, info.isLoaded)
// shouldn't access RootCoord again
assert.Equal(t, rootCoord.GetAccessCount(), 1)
globalMetaCache.RemoveCollection(ctx, "collection1")
// no collectionInfo of collection2, should access RootCoord
info, err = globalMetaCache.GetCollectionInfo(ctx, "collection1")
assert.NoError(t, err)
assert.True(t, info.isLoaded)
// shouldn't access RootCoord again
assert.Equal(t, rootCoord.GetAccessCount(), 2)
globalMetaCache.RemoveCollectionsByID(ctx, UniqueID(1))
// no collectionInfo of collection2, should access RootCoord
info, err = globalMetaCache.GetCollectionInfo(ctx, "collection1")
assert.NoError(t, err)
assert.True(t, info.isLoaded)
// shouldn't access RootCoord again
assert.Equal(t, rootCoord.GetAccessCount(), 3)
}
func TestMetaCache_ExpireShardLeaderCache(t *testing.T) {
ctx := context.Background()
rootCoord := &MockRootCoordClientInterface{}
queryCoord := &types.MockQueryCoord{}
shardMgr := newShardClientMgr()
err := InitMetaCache(ctx, rootCoord, queryCoord, shardMgr)
assert.Nil(t, err)
paramtable.Init()
paramtable.Get().Save(Params.ProxyCfg.ShardLeaderCacheInterval.Key, "1")
queryCoord.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
CollectionIDs: []UniqueID{1},
InMemoryPercentages: []int64{100},
}, nil)
queryCoord.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
Shards: []*querypb.ShardLeadersList{
{
ChannelName: "channel-1",
NodeIds: []int64{1, 2, 3},
NodeAddrs: []string{"localhost:9000", "localhost:9001", "localhost:9002"},
},
},
}, nil).Times(1)
nodeInfos, err := globalMetaCache.GetShards(ctx, true, "collection1")
assert.NoError(t, err)
assert.Len(t, nodeInfos["channel-1"], 3)
queryCoord.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
Shards: []*querypb.ShardLeadersList{
{
ChannelName: "channel-1",
NodeIds: []int64{1, 2},
NodeAddrs: []string{"localhost:9000", "localhost:9001"},
},
},
}, nil).Times(1)
assert.Eventually(t, func() bool {
nodeInfos, err := globalMetaCache.GetShards(ctx, true, "collection1")
assert.NoError(t, err)
return len(nodeInfos["channel-1"]) == 2
}, 3*time.Second, 1*time.Second)
queryCoord.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
Shards: []*querypb.ShardLeadersList{
{
ChannelName: "channel-1",
NodeIds: []int64{1, 2, 3},
NodeAddrs: []string{"localhost:9000", "localhost:9001", "localhost:9002"},
},
},
}, nil).Times(1)
assert.Eventually(t, func() bool {
nodeInfos, err := globalMetaCache.GetShards(ctx, true, "collection1")
assert.NoError(t, err)
return assert.Len(t, nodeInfos["channel-1"], 3)
}, 3*time.Second, 1*time.Second)
queryCoord.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
Shards: []*querypb.ShardLeadersList{
{
ChannelName: "channel-1",
NodeIds: []int64{1, 2, 3},
NodeAddrs: []string{"localhost:9000", "localhost:9001", "localhost:9002"},
},
{
ChannelName: "channel-2",
NodeIds: []int64{1, 2, 3},
NodeAddrs: []string{"localhost:9000", "localhost:9001", "localhost:9002"},
},
},
}, nil).Times(1)
assert.Eventually(t, func() bool {
nodeInfos, err := globalMetaCache.GetShards(ctx, true, "collection1")
assert.NoError(t, err)
return len(nodeInfos["channel-1"]) == 3 && len(nodeInfos["channel-2"]) == 3
}, 3*time.Second, 1*time.Second)
}