mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-02 11:59:00 +08:00
287 lines
11 KiB
Go
287 lines
11 KiB
Go
// 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 meta
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/samber/lo"
|
|
"github.com/stretchr/testify/mock"
|
|
"github.com/stretchr/testify/suite"
|
|
|
|
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
|
|
"github.com/milvus-io/milvus/internal/proto/datapb"
|
|
"github.com/milvus-io/milvus/internal/proto/querypb"
|
|
. "github.com/milvus-io/milvus/internal/querycoordv2/params"
|
|
"github.com/milvus-io/milvus/internal/util/etcd"
|
|
"github.com/milvus-io/milvus/internal/util/typeutil"
|
|
)
|
|
|
|
type TargetManagerSuite struct {
|
|
suite.Suite
|
|
|
|
// Data
|
|
collections []int64
|
|
partitions map[int64][]int64
|
|
channels map[int64][]string
|
|
segments map[int64]map[int64][]int64 // CollectionID, PartitionID -> Segments
|
|
// Derived data
|
|
allChannels []string
|
|
allSegments []int64
|
|
|
|
kv *etcdkv.EtcdKV
|
|
meta *Meta
|
|
broker *MockBroker
|
|
// Test object
|
|
mgr *TargetManager
|
|
}
|
|
|
|
func (suite *TargetManagerSuite) SetupSuite() {
|
|
Params.Init()
|
|
suite.collections = []int64{1000, 1001}
|
|
suite.partitions = map[int64][]int64{
|
|
1000: {100, 101},
|
|
1001: {102, 103},
|
|
}
|
|
suite.channels = map[int64][]string{
|
|
1000: {"1000-dmc0", "1000-dmc1"},
|
|
1001: {"1001-dmc0", "1001-dmc1"},
|
|
}
|
|
suite.segments = map[int64]map[int64][]int64{
|
|
1000: {
|
|
100: {1, 2},
|
|
101: {3, 4},
|
|
},
|
|
1001: {
|
|
102: {5, 6},
|
|
103: {7, 8},
|
|
},
|
|
}
|
|
|
|
suite.allChannels = make([]string, 0)
|
|
suite.allSegments = make([]int64, 0)
|
|
for _, channels := range suite.channels {
|
|
suite.allChannels = append(suite.allChannels, channels...)
|
|
}
|
|
for _, partitions := range suite.segments {
|
|
for _, segments := range partitions {
|
|
suite.allSegments = append(suite.allSegments, segments...)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (suite *TargetManagerSuite) SetupTest() {
|
|
var err error
|
|
config := GenerateEtcdConfig()
|
|
cli, err := etcd.GetEtcdClient(
|
|
config.UseEmbedEtcd.GetAsBool(),
|
|
config.EtcdUseSSL.GetAsBool(),
|
|
config.Endpoints.GetAsStrings(),
|
|
config.EtcdTLSCert.GetValue(),
|
|
config.EtcdTLSKey.GetValue(),
|
|
config.EtcdTLSCACert.GetValue(),
|
|
config.EtcdTLSMinVersion.GetValue())
|
|
suite.Require().NoError(err)
|
|
suite.kv = etcdkv.NewEtcdKV(cli, config.MetaRootPath.GetValue())
|
|
|
|
// meta
|
|
store := NewMetaStore(suite.kv)
|
|
idAllocator := RandomIncrementIDAllocator()
|
|
suite.meta = NewMeta(idAllocator, store)
|
|
suite.broker = NewMockBroker(suite.T())
|
|
suite.mgr = NewTargetManager(suite.broker, suite.meta)
|
|
|
|
for _, collection := range suite.collections {
|
|
dmChannels := make([]*datapb.VchannelInfo, 0)
|
|
for _, channel := range suite.channels[collection] {
|
|
dmChannels = append(dmChannels, &datapb.VchannelInfo{
|
|
CollectionID: collection,
|
|
ChannelName: channel,
|
|
})
|
|
}
|
|
|
|
for partition, segments := range suite.segments[collection] {
|
|
allSegments := make([]*datapb.SegmentBinlogs, 0)
|
|
for _, segment := range segments {
|
|
|
|
allSegments = append(allSegments, &datapb.SegmentBinlogs{
|
|
SegmentID: segment,
|
|
InsertChannel: suite.channels[collection][0],
|
|
})
|
|
}
|
|
suite.broker.EXPECT().GetRecoveryInfo(mock.Anything, collection, partition).Return(dmChannels, allSegments, nil)
|
|
}
|
|
|
|
suite.mgr.UpdateCollectionNextTargetWithPartitions(collection, suite.partitions[collection]...)
|
|
}
|
|
}
|
|
|
|
func (suite *TargetManagerSuite) TearDownSuite() {
|
|
suite.kv.Close()
|
|
}
|
|
|
|
func (suite *TargetManagerSuite) TestUpdateCurrentTarget() {
|
|
collectionID := int64(1000)
|
|
suite.assertSegments(suite.getAllSegment(collectionID, suite.partitions[collectionID]),
|
|
suite.mgr.GetHistoricalSegmentsByCollection(collectionID, NextTarget))
|
|
suite.assertChannels(suite.channels[collectionID], suite.mgr.GetDmChannelsByCollection(collectionID, NextTarget))
|
|
suite.assertSegments([]int64{}, suite.mgr.GetHistoricalSegmentsByCollection(collectionID, CurrentTarget))
|
|
suite.assertChannels([]string{}, suite.mgr.GetDmChannelsByCollection(collectionID, CurrentTarget))
|
|
|
|
suite.mgr.UpdateCollectionCurrentTarget(collectionID)
|
|
suite.assertSegments([]int64{}, suite.mgr.GetHistoricalSegmentsByCollection(collectionID, NextTarget))
|
|
suite.assertChannels([]string{}, suite.mgr.GetDmChannelsByCollection(collectionID, NextTarget))
|
|
suite.assertSegments(suite.getAllSegment(collectionID, suite.partitions[collectionID]),
|
|
suite.mgr.GetHistoricalSegmentsByCollection(collectionID, CurrentTarget))
|
|
suite.assertChannels(suite.channels[collectionID], suite.mgr.GetDmChannelsByCollection(collectionID, CurrentTarget))
|
|
}
|
|
|
|
func (suite *TargetManagerSuite) TestUpdateNextTarget() {
|
|
collectionID := int64(1003)
|
|
suite.assertSegments([]int64{}, suite.mgr.GetHistoricalSegmentsByCollection(collectionID, NextTarget))
|
|
suite.assertChannels([]string{}, suite.mgr.GetDmChannelsByCollection(collectionID, NextTarget))
|
|
suite.assertSegments([]int64{}, suite.mgr.GetHistoricalSegmentsByCollection(collectionID, CurrentTarget))
|
|
suite.assertChannels([]string{}, suite.mgr.GetDmChannelsByCollection(collectionID, CurrentTarget))
|
|
|
|
suite.meta.PutCollection(&Collection{
|
|
CollectionLoadInfo: &querypb.CollectionLoadInfo{
|
|
CollectionID: collectionID,
|
|
ReplicaNumber: 1},
|
|
})
|
|
suite.meta.PutPartition(&Partition{
|
|
PartitionLoadInfo: &querypb.PartitionLoadInfo{
|
|
CollectionID: collectionID,
|
|
PartitionID: 1,
|
|
},
|
|
})
|
|
|
|
nextTargetChannels := []*datapb.VchannelInfo{
|
|
{
|
|
CollectionID: collectionID,
|
|
ChannelName: "channel-1",
|
|
},
|
|
{
|
|
CollectionID: collectionID,
|
|
ChannelName: "channel-2",
|
|
},
|
|
}
|
|
|
|
nextTargetSegments := []*datapb.SegmentBinlogs{
|
|
{
|
|
SegmentID: 11,
|
|
InsertChannel: "channel-1",
|
|
},
|
|
{
|
|
SegmentID: 12,
|
|
InsertChannel: "channel-2",
|
|
},
|
|
}
|
|
|
|
suite.broker.EXPECT().GetRecoveryInfo(mock.Anything, collectionID, int64(1)).Return(nextTargetChannels, nextTargetSegments, nil)
|
|
suite.mgr.UpdateCollectionNextTargetWithPartitions(collectionID, int64(1))
|
|
suite.assertSegments([]int64{11, 12}, suite.mgr.GetHistoricalSegmentsByCollection(collectionID, NextTarget))
|
|
suite.assertChannels([]string{"channel-1", "channel-2"}, suite.mgr.GetDmChannelsByCollection(collectionID, NextTarget))
|
|
suite.assertSegments([]int64{}, suite.mgr.GetHistoricalSegmentsByCollection(collectionID, CurrentTarget))
|
|
suite.assertChannels([]string{}, suite.mgr.GetDmChannelsByCollection(collectionID, CurrentTarget))
|
|
}
|
|
|
|
func (suite *TargetManagerSuite) TestRemovePartition() {
|
|
collectionID := int64(1000)
|
|
suite.assertSegments(suite.getAllSegment(collectionID, suite.partitions[collectionID]), suite.mgr.GetHistoricalSegmentsByCollection(collectionID, NextTarget))
|
|
suite.assertChannels(suite.channels[collectionID], suite.mgr.GetDmChannelsByCollection(collectionID, NextTarget))
|
|
suite.assertSegments([]int64{}, suite.mgr.GetHistoricalSegmentsByCollection(collectionID, CurrentTarget))
|
|
suite.assertChannels([]string{}, suite.mgr.GetDmChannelsByCollection(collectionID, CurrentTarget))
|
|
|
|
suite.mgr.RemovePartition(collectionID, 100)
|
|
suite.assertSegments([]int64{3, 4}, suite.mgr.GetHistoricalSegmentsByCollection(collectionID, NextTarget))
|
|
suite.assertChannels(suite.channels[collectionID], suite.mgr.GetDmChannelsByCollection(collectionID, NextTarget))
|
|
suite.assertSegments([]int64{}, suite.mgr.GetHistoricalSegmentsByCollection(collectionID, CurrentTarget))
|
|
suite.assertChannels([]string{}, suite.mgr.GetDmChannelsByCollection(collectionID, CurrentTarget))
|
|
}
|
|
|
|
func (suite *TargetManagerSuite) TestRemoveCollection() {
|
|
collectionID := int64(1000)
|
|
suite.assertSegments(suite.getAllSegment(collectionID, suite.partitions[collectionID]), suite.mgr.GetHistoricalSegmentsByCollection(collectionID, NextTarget))
|
|
suite.assertChannels(suite.channels[collectionID], suite.mgr.GetDmChannelsByCollection(collectionID, NextTarget))
|
|
suite.assertSegments([]int64{}, suite.mgr.GetHistoricalSegmentsByCollection(collectionID, CurrentTarget))
|
|
suite.assertChannels([]string{}, suite.mgr.GetDmChannelsByCollection(collectionID, CurrentTarget))
|
|
|
|
suite.mgr.RemoveCollection(collectionID)
|
|
suite.assertSegments([]int64{}, suite.mgr.GetHistoricalSegmentsByCollection(collectionID, NextTarget))
|
|
suite.assertChannels([]string{}, suite.mgr.GetDmChannelsByCollection(collectionID, NextTarget))
|
|
suite.assertSegments([]int64{}, suite.mgr.GetHistoricalSegmentsByCollection(collectionID, CurrentTarget))
|
|
suite.assertChannels([]string{}, suite.mgr.GetDmChannelsByCollection(collectionID, CurrentTarget))
|
|
|
|
collectionID = int64(1001)
|
|
suite.mgr.UpdateCollectionCurrentTarget(collectionID)
|
|
suite.assertSegments([]int64{}, suite.mgr.GetHistoricalSegmentsByCollection(collectionID, NextTarget))
|
|
suite.assertChannels([]string{}, suite.mgr.GetDmChannelsByCollection(collectionID, NextTarget))
|
|
suite.assertSegments(suite.getAllSegment(collectionID, suite.partitions[collectionID]), suite.mgr.GetHistoricalSegmentsByCollection(collectionID, CurrentTarget))
|
|
suite.assertChannels(suite.channels[collectionID], suite.mgr.GetDmChannelsByCollection(collectionID, CurrentTarget))
|
|
|
|
suite.mgr.RemoveCollection(collectionID)
|
|
suite.assertSegments([]int64{}, suite.mgr.GetHistoricalSegmentsByCollection(collectionID, NextTarget))
|
|
suite.assertChannels([]string{}, suite.mgr.GetDmChannelsByCollection(collectionID, NextTarget))
|
|
suite.assertSegments([]int64{}, suite.mgr.GetHistoricalSegmentsByCollection(collectionID, CurrentTarget))
|
|
suite.assertChannels([]string{}, suite.mgr.GetDmChannelsByCollection(collectionID, CurrentTarget))
|
|
}
|
|
|
|
func (suite *TargetManagerSuite) getAllSegment(collectionID int64, partitionIDs []int64) []int64 {
|
|
allSegments := make([]int64, 0)
|
|
for collection, partitions := range suite.segments {
|
|
if collectionID == collection {
|
|
for partition, segments := range partitions {
|
|
if lo.Contains(partitionIDs, partition) {
|
|
allSegments = append(allSegments, segments...)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return allSegments
|
|
}
|
|
|
|
func (suite *TargetManagerSuite) assertChannels(expected []string, actual map[string]*DmChannel) bool {
|
|
if !suite.Equal(len(expected), len(actual)) {
|
|
return false
|
|
}
|
|
|
|
set := typeutil.NewSet(expected...)
|
|
for _, channel := range actual {
|
|
set.Remove(channel.ChannelName)
|
|
}
|
|
|
|
return suite.Len(set, 0)
|
|
}
|
|
|
|
func (suite *TargetManagerSuite) assertSegments(expected []int64, actual map[int64]*datapb.SegmentInfo) bool {
|
|
if !suite.Equal(len(expected), len(actual)) {
|
|
return false
|
|
}
|
|
|
|
set := typeutil.NewUniqueSet(expected...)
|
|
for _, segment := range actual {
|
|
set.Remove(segment.ID)
|
|
}
|
|
|
|
return suite.Len(set, 0)
|
|
}
|
|
|
|
func TestTargetManager(t *testing.T) {
|
|
suite.Run(t, new(TargetManagerSuite))
|
|
}
|