milvus/internal/querycoordv2/job/job_test.go

1525 lines
39 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 job
import (
"context"
"fmt"
"testing"
"github.com/cockroachdb/errors"
"github.com/samber/lo"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/suite"
"github.com/milvus-io/milvus-proto/go-api/v2/rgpb"
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
"github.com/milvus-io/milvus/internal/metastore"
"github.com/milvus-io/milvus/internal/metastore/kv/querycoord"
"github.com/milvus-io/milvus/internal/metastore/mocks"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/proto/querypb"
"github.com/milvus-io/milvus/internal/querycoordv2/checkers"
"github.com/milvus-io/milvus/internal/querycoordv2/meta"
"github.com/milvus-io/milvus/internal/querycoordv2/observers"
. "github.com/milvus-io/milvus/internal/querycoordv2/params"
"github.com/milvus-io/milvus/internal/querycoordv2/session"
"github.com/milvus-io/milvus/internal/util/proxyutil"
"github.com/milvus-io/milvus/pkg/kv"
"github.com/milvus-io/milvus/pkg/util/etcd"
"github.com/milvus-io/milvus/pkg/util/merr"
"github.com/milvus-io/milvus/pkg/util/paramtable"
)
const (
defaultVecFieldID = 1
defaultIndexID = 1
)
type JobSuite struct {
suite.Suite
// Data
collections []int64
partitions map[int64][]int64
channels map[int64][]string
segments map[int64]map[int64][]int64 // CollectionID, PartitionID -> Segments
loadTypes map[int64]querypb.LoadType
// Dependencies
kv kv.MetaKv
store metastore.QueryCoordCatalog
dist *meta.DistributionManager
meta *meta.Meta
cluster *session.MockCluster
targetMgr *meta.TargetManager
targetObserver *observers.TargetObserver
collectionObserver *observers.CollectionObserver
broker *meta.MockBroker
nodeMgr *session.NodeManager
checkerController *checkers.CheckerController
proxyManager *proxyutil.MockProxyClientManager
// Test objects
scheduler *Scheduler
}
func (suite *JobSuite) SetupSuite() {
paramtable.Init()
suite.collections = []int64{1000, 1001}
suite.partitions = map[int64][]int64{
1000: {100, 101, 102},
1001: {103, 104, 105},
}
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},
102: {5, 6},
},
1001: {
103: {7, 8},
104: {9, 10},
105: {11, 12},
},
}
suite.loadTypes = map[int64]querypb.LoadType{
1000: querypb.LoadType_LoadCollection,
1001: querypb.LoadType_LoadPartition,
}
suite.broker = meta.NewMockBroker(suite.T())
for collection, partitions := range suite.segments {
vChannels := []*datapb.VchannelInfo{}
for _, channel := range suite.channels[collection] {
vChannels = append(vChannels, &datapb.VchannelInfo{
CollectionID: collection,
ChannelName: channel,
})
}
segmentBinlogs := []*datapb.SegmentInfo{}
for partition, segments := range partitions {
for _, segment := range segments {
segmentBinlogs = append(segmentBinlogs, &datapb.SegmentInfo{
ID: segment,
PartitionID: partition,
InsertChannel: suite.channels[collection][segment%2],
})
}
}
suite.broker.EXPECT().GetRecoveryInfoV2(mock.Anything, collection).Return(vChannels, segmentBinlogs, nil).Maybe()
}
suite.broker.EXPECT().DescribeCollection(mock.Anything, mock.Anything).
Return(nil, nil)
suite.broker.EXPECT().ListIndexes(mock.Anything, mock.Anything).
Return(nil, nil)
suite.cluster = session.NewMockCluster(suite.T())
suite.cluster.EXPECT().
LoadPartitions(mock.Anything, mock.Anything, mock.Anything).
Return(merr.Success(), nil)
suite.cluster.EXPECT().
ReleasePartitions(mock.Anything, mock.Anything, mock.Anything).
Return(merr.Success(), nil).Maybe()
suite.proxyManager = proxyutil.NewMockProxyClientManager(suite.T())
suite.proxyManager.EXPECT().InvalidateCollectionMetaCache(mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
}
func (suite *JobSuite) SetupTest() {
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())
suite.store = querycoord.NewCatalog(suite.kv)
suite.dist = meta.NewDistributionManager()
suite.nodeMgr = session.NewNodeManager()
suite.meta = meta.NewMeta(RandomIncrementIDAllocator(), suite.store, suite.nodeMgr)
suite.targetMgr = meta.NewTargetManager(suite.broker, suite.meta)
suite.targetObserver = observers.NewTargetObserver(suite.meta,
suite.targetMgr,
suite.dist,
suite.broker,
suite.cluster,
)
suite.targetObserver.Start()
suite.scheduler = NewScheduler()
suite.scheduler.Start()
meta.GlobalFailedLoadCache = meta.NewFailedLoadCache()
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
NodeID: 1000,
Address: "localhost",
Hostname: "localhost",
}))
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
NodeID: 2000,
Address: "localhost",
Hostname: "localhost",
}))
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
NodeID: 3000,
Address: "localhost",
Hostname: "localhost",
}))
suite.meta.HandleNodeUp(1000)
suite.meta.HandleNodeUp(2000)
suite.meta.HandleNodeUp(3000)
suite.checkerController = &checkers.CheckerController{}
suite.collectionObserver = observers.NewCollectionObserver(
suite.dist,
suite.meta,
suite.targetMgr,
suite.targetObserver,
suite.checkerController,
suite.proxyManager,
)
}
func (suite *JobSuite) TearDownTest() {
suite.kv.Close()
suite.scheduler.Stop()
suite.targetObserver.Stop()
}
func (suite *JobSuite) BeforeTest(suiteName, testName string) {
for collection, partitions := range suite.partitions {
suite.broker.EXPECT().
GetPartitions(mock.Anything, collection).
Return(partitions, nil)
}
}
func (suite *JobSuite) TestLoadCollection() {
ctx := context.Background()
// Test load collection
for _, collection := range suite.collections {
if suite.loadTypes[collection] != querypb.LoadType_LoadCollection {
continue
}
// Load with 1 replica
req := &querypb.LoadCollectionRequest{
CollectionID: collection,
// It will be set to 1
// ReplicaNumber: 1,
}
job := NewLoadCollectionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.NoError(err)
suite.EqualValues(1, suite.meta.GetReplicaNumber(collection))
suite.targetMgr.UpdateCollectionCurrentTarget(collection)
suite.assertCollectionLoaded(collection)
}
// Test load again
for _, collection := range suite.collections {
if suite.loadTypes[collection] != querypb.LoadType_LoadCollection {
continue
}
req := &querypb.LoadCollectionRequest{
CollectionID: collection,
}
job := NewLoadCollectionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.NoError(err)
}
// Test load existed collection with different replica number
for _, collection := range suite.collections {
if suite.loadTypes[collection] != querypb.LoadType_LoadCollection {
continue
}
req := &querypb.LoadCollectionRequest{
CollectionID: collection,
ReplicaNumber: 3,
}
job := NewLoadCollectionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.ErrorIs(err, merr.ErrParameterInvalid)
}
// Test load existed collection with different load fields
for _, collection := range suite.collections {
if suite.loadTypes[collection] != querypb.LoadType_LoadCollection {
continue
}
req := &querypb.LoadCollectionRequest{
CollectionID: collection,
LoadFields: []int64{100, 101},
}
job := NewLoadCollectionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.ErrorIs(err, merr.ErrParameterInvalid)
}
// Test load partition while collection exists
for _, collection := range suite.collections {
if suite.loadTypes[collection] != querypb.LoadType_LoadCollection {
continue
}
// Load with 1 replica
req := &querypb.LoadPartitionsRequest{
CollectionID: collection,
PartitionIDs: suite.partitions[collection],
ReplicaNumber: 1,
}
job := NewLoadPartitionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.NoError(err)
}
cfg := &rgpb.ResourceGroupConfig{
Requests: &rgpb.ResourceGroupLimit{
NodeNum: 0,
},
Limits: &rgpb.ResourceGroupLimit{
NodeNum: 0,
},
}
suite.meta.ResourceManager.AddResourceGroup("rg1", cfg)
suite.meta.ResourceManager.AddResourceGroup("rg2", cfg)
suite.meta.ResourceManager.AddResourceGroup("rg3", cfg)
// Load with 3 replica on 1 rg
req := &querypb.LoadCollectionRequest{
CollectionID: 1001,
ReplicaNumber: 3,
ResourceGroups: []string{"rg1"},
}
job := NewLoadCollectionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.ErrorIs(err, merr.ErrResourceGroupNodeNotEnough)
// Load with 3 replica on 3 rg
req = &querypb.LoadCollectionRequest{
CollectionID: 1001,
ReplicaNumber: 3,
ResourceGroups: []string{"rg1", "rg2", "rg3"},
}
job = NewLoadCollectionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
err = job.Wait()
suite.ErrorIs(err, merr.ErrResourceGroupNodeNotEnough)
}
func (suite *JobSuite) TestLoadCollectionWithReplicas() {
ctx := context.Background()
// Test load collection
for _, collection := range suite.collections {
if suite.loadTypes[collection] != querypb.LoadType_LoadCollection {
continue
}
// Load with 3 replica
req := &querypb.LoadCollectionRequest{
CollectionID: collection,
ReplicaNumber: 5,
}
job := NewLoadCollectionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.ErrorIs(err, merr.ErrResourceGroupNodeNotEnough)
}
}
func (suite *JobSuite) TestLoadPartition() {
ctx := context.Background()
// Test load partition
for _, collection := range suite.collections {
if suite.loadTypes[collection] != querypb.LoadType_LoadPartition {
continue
}
// Load with 1 replica
req := &querypb.LoadPartitionsRequest{
CollectionID: collection,
PartitionIDs: suite.partitions[collection],
ReplicaNumber: 1,
}
job := NewLoadPartitionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.NoError(err)
suite.EqualValues(1, suite.meta.GetReplicaNumber(collection))
suite.targetMgr.UpdateCollectionCurrentTarget(collection)
suite.assertCollectionLoaded(collection)
}
// Test load partition again
for _, collection := range suite.collections {
if suite.loadTypes[collection] != querypb.LoadType_LoadPartition {
continue
}
// Load with 1 replica
req := &querypb.LoadPartitionsRequest{
CollectionID: collection,
PartitionIDs: suite.partitions[collection],
// ReplicaNumber: 1,
}
job := NewLoadPartitionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.NoError(err)
}
// Test load partition with different replica number
for _, collection := range suite.collections {
if suite.loadTypes[collection] != querypb.LoadType_LoadPartition {
continue
}
req := &querypb.LoadPartitionsRequest{
CollectionID: collection,
PartitionIDs: suite.partitions[collection],
ReplicaNumber: 3,
}
job := NewLoadPartitionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.ErrorIs(err, merr.ErrParameterInvalid)
}
// Test load partition with different load fields
for _, collection := range suite.collections {
if suite.loadTypes[collection] != querypb.LoadType_LoadPartition {
continue
}
req := &querypb.LoadPartitionsRequest{
CollectionID: collection,
PartitionIDs: suite.partitions[collection],
LoadFields: []int64{100, 101},
}
job := NewLoadPartitionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.ErrorIs(err, merr.ErrParameterInvalid)
}
// Test load partition with more partition
for _, collection := range suite.collections {
if suite.loadTypes[collection] != querypb.LoadType_LoadPartition {
continue
}
req := &querypb.LoadPartitionsRequest{
CollectionID: collection,
PartitionIDs: append(suite.partitions[collection], 200),
ReplicaNumber: 1,
}
job := NewLoadPartitionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.NoError(err)
}
// Test load collection while partitions exists
for _, collection := range suite.collections {
if suite.loadTypes[collection] != querypb.LoadType_LoadPartition {
continue
}
req := &querypb.LoadCollectionRequest{
CollectionID: collection,
ReplicaNumber: 1,
}
job := NewLoadCollectionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.NoError(err)
}
cfg := &rgpb.ResourceGroupConfig{
Requests: &rgpb.ResourceGroupLimit{
NodeNum: 1,
},
Limits: &rgpb.ResourceGroupLimit{
NodeNum: 1,
},
}
suite.meta.ResourceManager.AddResourceGroup("rg1", cfg)
suite.meta.ResourceManager.AddResourceGroup("rg2", cfg)
suite.meta.ResourceManager.AddResourceGroup("rg3", cfg)
// test load 3 replica in 1 rg, should pass rg check
req := &querypb.LoadPartitionsRequest{
CollectionID: 999,
PartitionIDs: []int64{888},
ReplicaNumber: 3,
ResourceGroups: []string{"rg1"},
}
job := NewLoadPartitionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.ErrorIs(err, merr.ErrResourceGroupNodeNotEnough)
// test load 3 replica in 3 rg, should pass rg check
req = &querypb.LoadPartitionsRequest{
CollectionID: 999,
PartitionIDs: []int64{888},
ReplicaNumber: 3,
ResourceGroups: []string{"rg1", "rg2", "rg3"},
}
job = NewLoadPartitionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
err = job.Wait()
suite.ErrorIs(err, merr.ErrResourceGroupNodeNotEnough)
}
func (suite *JobSuite) TestDynamicLoad() {
ctx := context.Background()
collection := suite.collections[0]
p0, p1, p2 := suite.partitions[collection][0], suite.partitions[collection][1], suite.partitions[collection][2]
newLoadPartJob := func(partitions ...int64) *LoadPartitionJob {
req := &querypb.LoadPartitionsRequest{
CollectionID: collection,
PartitionIDs: partitions,
ReplicaNumber: 1,
}
job := NewLoadPartitionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
return job
}
newLoadColJob := func() *LoadCollectionJob {
req := &querypb.LoadCollectionRequest{
CollectionID: collection,
ReplicaNumber: 1,
}
job := NewLoadCollectionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
return job
}
// loaded: none
// action: load p0, p1, p2
// expect: p0, p1, p2 loaded
job := newLoadPartJob(p0, p1, p2)
suite.scheduler.Add(job)
err := job.Wait()
suite.NoError(err)
suite.targetMgr.UpdateCollectionCurrentTarget(collection)
suite.assertPartitionLoaded(collection, p0, p1, p2)
// loaded: p0, p1, p2
// action: load p0, p1, p2
// expect: do nothing, p0, p1, p2 loaded
job = newLoadPartJob(p0, p1, p2)
suite.scheduler.Add(job)
err = job.Wait()
suite.NoError(err)
suite.assertPartitionLoaded(collection)
// loaded: p0, p1
// action: load p2
// expect: p0, p1, p2 loaded
suite.releaseAll()
job = newLoadPartJob(p0, p1)
suite.scheduler.Add(job)
err = job.Wait()
suite.NoError(err)
suite.targetMgr.UpdateCollectionCurrentTarget(collection)
suite.assertPartitionLoaded(collection, p0, p1)
job = newLoadPartJob(p2)
suite.scheduler.Add(job)
err = job.Wait()
suite.NoError(err)
suite.targetMgr.UpdateCollectionCurrentTarget(collection)
suite.assertPartitionLoaded(collection, p2)
// loaded: p0, p1
// action: load p1, p2
// expect: p0, p1, p2 loaded
suite.releaseAll()
job = newLoadPartJob(p0, p1)
suite.scheduler.Add(job)
err = job.Wait()
suite.NoError(err)
suite.targetMgr.UpdateCollectionCurrentTarget(collection)
suite.assertPartitionLoaded(collection, p0, p1)
job = newLoadPartJob(p1, p2)
suite.scheduler.Add(job)
err = job.Wait()
suite.NoError(err)
suite.targetMgr.UpdateCollectionCurrentTarget(collection)
suite.assertPartitionLoaded(collection, p2)
// loaded: p0, p1
// action: load col
// expect: col loaded
suite.releaseAll()
job = newLoadPartJob(p0, p1)
suite.scheduler.Add(job)
err = job.Wait()
suite.NoError(err)
suite.targetMgr.UpdateCollectionCurrentTarget(collection)
suite.assertPartitionLoaded(collection, p0, p1)
colJob := newLoadColJob()
suite.scheduler.Add(colJob)
err = colJob.Wait()
suite.NoError(err)
suite.targetMgr.UpdateCollectionCurrentTarget(collection)
suite.assertPartitionLoaded(collection, p2)
}
func (suite *JobSuite) TestLoadPartitionWithReplicas() {
ctx := context.Background()
// Test load partitions
for _, collection := range suite.collections {
if suite.loadTypes[collection] != querypb.LoadType_LoadPartition {
continue
}
// Load with 3 replica
req := &querypb.LoadPartitionsRequest{
CollectionID: collection,
PartitionIDs: suite.partitions[collection],
ReplicaNumber: 5,
}
job := NewLoadPartitionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.ErrorIs(err, merr.ErrResourceGroupNodeNotEnough)
}
}
func (suite *JobSuite) TestReleaseCollection() {
ctx := context.Background()
suite.loadAll()
// Test release collection and partition
for _, collection := range suite.collections {
req := &querypb.ReleaseCollectionRequest{
CollectionID: collection,
}
job := NewReleaseCollectionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.checkerController,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.NoError(err)
suite.assertCollectionReleased(collection)
}
// Test release again
for _, collection := range suite.collections {
req := &querypb.ReleaseCollectionRequest{
CollectionID: collection,
}
job := NewReleaseCollectionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.checkerController,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.NoError(err)
suite.assertCollectionReleased(collection)
}
}
func (suite *JobSuite) TestReleasePartition() {
ctx := context.Background()
suite.loadAll()
// Test release partition
for _, collection := range suite.collections {
req := &querypb.ReleasePartitionsRequest{
CollectionID: collection,
PartitionIDs: suite.partitions[collection],
}
job := NewReleasePartitionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.checkerController,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.NoError(err)
suite.assertPartitionReleased(collection, suite.partitions[collection]...)
}
// Test release again
for _, collection := range suite.collections {
req := &querypb.ReleasePartitionsRequest{
CollectionID: collection,
PartitionIDs: suite.partitions[collection],
}
job := NewReleasePartitionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.checkerController,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.NoError(err)
suite.assertPartitionReleased(collection, suite.partitions[collection]...)
}
// Test release partial partitions
suite.releaseAll()
suite.loadAll()
for _, collection := range suite.collections {
req := &querypb.ReleasePartitionsRequest{
CollectionID: collection,
PartitionIDs: suite.partitions[collection][1:],
}
job := NewReleasePartitionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.checkerController,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.NoError(err)
suite.True(suite.meta.Exist(collection))
partitions := suite.meta.GetPartitionsByCollection(collection)
suite.Len(partitions, 1)
suite.Equal(suite.partitions[collection][0], partitions[0].GetPartitionID())
suite.assertPartitionReleased(collection, suite.partitions[collection][1:]...)
}
}
func (suite *JobSuite) TestDynamicRelease() {
ctx := context.Background()
col0, col1 := suite.collections[0], suite.collections[1]
p0, p1, p2 := suite.partitions[col0][0], suite.partitions[col0][1], suite.partitions[col0][2]
p3, p4, p5 := suite.partitions[col1][0], suite.partitions[col1][1], suite.partitions[col1][2]
newReleasePartJob := func(col int64, partitions ...int64) *ReleasePartitionJob {
req := &querypb.ReleasePartitionsRequest{
CollectionID: col,
PartitionIDs: partitions,
}
job := NewReleasePartitionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.checkerController,
)
return job
}
newReleaseColJob := func(col int64) *ReleaseCollectionJob {
req := &querypb.ReleaseCollectionRequest{
CollectionID: col,
}
job := NewReleaseCollectionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.checkerController,
)
return job
}
// loaded: p0, p1, p2
// action: release p0
// expect: p0 released, p1, p2 loaded
suite.loadAll()
job := newReleasePartJob(col0, p0)
suite.scheduler.Add(job)
err := job.Wait()
suite.NoError(err)
suite.assertPartitionReleased(col0, p0)
suite.assertPartitionLoaded(col0, p1, p2)
// loaded: p1, p2
// action: release p0, p1
// expect: p1 released, p2 loaded
job = newReleasePartJob(col0, p0, p1)
suite.scheduler.Add(job)
err = job.Wait()
suite.NoError(err)
suite.assertPartitionReleased(col0, p0, p1)
suite.assertPartitionLoaded(col0, p2)
// loaded: p2
// action: release p2
// expect: loadType=col: col loaded, p2 released
job = newReleasePartJob(col0, p2)
suite.scheduler.Add(job)
err = job.Wait()
suite.NoError(err)
suite.assertPartitionReleased(col0, p0, p1, p2)
suite.False(suite.meta.Exist(col0))
// loaded: p0, p1, p2
// action: release col
// expect: col released
suite.releaseAll()
suite.loadAll()
releaseColJob := newReleaseColJob(col0)
suite.scheduler.Add(releaseColJob)
err = releaseColJob.Wait()
suite.NoError(err)
suite.assertCollectionReleased(col0)
suite.assertPartitionReleased(col0, p0, p1, p2)
// loaded: p3, p4, p5
// action: release p3, p4, p5
// expect: loadType=partition: col released
suite.releaseAll()
suite.loadAll()
job = newReleasePartJob(col1, p3, p4, p5)
suite.scheduler.Add(job)
err = job.Wait()
suite.NoError(err)
suite.assertCollectionReleased(col1)
suite.assertPartitionReleased(col1, p3, p4, p5)
}
func (suite *JobSuite) TestLoadCollectionStoreFailed() {
// Store collection failed
store := mocks.NewQueryCoordCatalog(suite.T())
suite.meta = meta.NewMeta(RandomIncrementIDAllocator(), store, suite.nodeMgr)
store.EXPECT().SaveResourceGroup(mock.Anything, mock.Anything).Return(nil)
suite.meta.HandleNodeUp(1000)
suite.meta.HandleNodeUp(2000)
suite.meta.HandleNodeUp(3000)
for _, collection := range suite.collections {
if suite.loadTypes[collection] != querypb.LoadType_LoadCollection {
continue
}
suite.broker.EXPECT().GetPartitions(mock.Anything, collection).Return(suite.partitions[collection], nil)
err := errors.New("failed to store collection")
store.EXPECT().SaveReplica(mock.Anything).Return(nil)
store.EXPECT().SaveCollection(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(err)
store.EXPECT().ReleaseReplicas(collection).Return(nil)
req := &querypb.LoadCollectionRequest{
CollectionID: collection,
}
job := NewLoadCollectionJob(
context.Background(),
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
loadErr := job.Wait()
suite.ErrorIs(loadErr, err)
}
}
func (suite *JobSuite) TestLoadPartitionStoreFailed() {
// Store partition failed
store := mocks.NewQueryCoordCatalog(suite.T())
suite.meta = meta.NewMeta(RandomIncrementIDAllocator(), store, suite.nodeMgr)
store.EXPECT().SaveResourceGroup(mock.Anything, mock.Anything).Return(nil)
suite.meta.HandleNodeUp(1000)
suite.meta.HandleNodeUp(2000)
suite.meta.HandleNodeUp(3000)
err := errors.New("failed to store collection")
for _, collection := range suite.collections {
if suite.loadTypes[collection] != querypb.LoadType_LoadPartition {
continue
}
store.EXPECT().SaveReplica(mock.Anything).Return(nil)
store.EXPECT().SaveCollection(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(err)
store.EXPECT().ReleaseReplicas(collection).Return(nil)
req := &querypb.LoadPartitionsRequest{
CollectionID: collection,
PartitionIDs: suite.partitions[collection],
}
job := NewLoadPartitionJob(
context.Background(),
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
loadErr := job.Wait()
suite.ErrorIs(loadErr, err)
}
}
func (suite *JobSuite) TestLoadCreateReplicaFailed() {
// Store replica failed
suite.meta = meta.NewMeta(ErrorIDAllocator(), suite.store, session.NewNodeManager())
for _, collection := range suite.collections {
suite.broker.EXPECT().
GetPartitions(mock.Anything, collection).
Return(suite.partitions[collection], nil)
req := &querypb.LoadCollectionRequest{
CollectionID: collection,
}
job := NewLoadCollectionJob(
context.Background(),
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.ErrorIs(err, merr.ErrResourceGroupNodeNotEnough)
}
}
func (suite *JobSuite) TestCallLoadPartitionFailed() {
// call LoadPartitions failed at get index info
getIndexErr := fmt.Errorf("mock get index error")
suite.broker.ExpectedCalls = lo.Filter(suite.broker.ExpectedCalls, func(call *mock.Call, _ int) bool {
return call.Method != "ListIndexes"
})
for _, collection := range suite.collections {
suite.broker.EXPECT().ListIndexes(mock.Anything, collection).Return(nil, getIndexErr)
loadCollectionReq := &querypb.LoadCollectionRequest{
CollectionID: collection,
}
loadCollectionJob := NewLoadCollectionJob(
context.Background(),
loadCollectionReq,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(loadCollectionJob)
err := loadCollectionJob.Wait()
suite.T().Logf("%s", err)
suite.ErrorIs(err, getIndexErr)
loadPartitionReq := &querypb.LoadPartitionsRequest{
CollectionID: collection,
PartitionIDs: suite.partitions[collection],
}
loadPartitionJob := NewLoadPartitionJob(
context.Background(),
loadPartitionReq,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(loadPartitionJob)
err = loadPartitionJob.Wait()
suite.ErrorIs(err, getIndexErr)
}
// call LoadPartitions failed at get schema
getSchemaErr := fmt.Errorf("mock get schema error")
suite.broker.ExpectedCalls = lo.Filter(suite.broker.ExpectedCalls, func(call *mock.Call, _ int) bool {
return call.Method != "DescribeCollection"
})
for _, collection := range suite.collections {
suite.broker.EXPECT().DescribeCollection(mock.Anything, collection).Return(nil, getSchemaErr)
loadCollectionReq := &querypb.LoadCollectionRequest{
CollectionID: collection,
}
loadCollectionJob := NewLoadCollectionJob(
context.Background(),
loadCollectionReq,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(loadCollectionJob)
err := loadCollectionJob.Wait()
suite.ErrorIs(err, getSchemaErr)
loadPartitionReq := &querypb.LoadPartitionsRequest{
CollectionID: collection,
PartitionIDs: suite.partitions[collection],
}
loadPartitionJob := NewLoadPartitionJob(
context.Background(),
loadPartitionReq,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(loadPartitionJob)
err = loadPartitionJob.Wait()
suite.ErrorIs(err, getSchemaErr)
}
suite.broker.ExpectedCalls = lo.Filter(suite.broker.ExpectedCalls, func(call *mock.Call, _ int) bool {
return call.Method != "ListIndexes" && call.Method != "DescribeCollection"
})
suite.broker.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(nil, nil)
suite.broker.EXPECT().ListIndexes(mock.Anything, mock.Anything).Return(nil, nil)
}
func (suite *JobSuite) TestCallReleasePartitionFailed() {
ctx := context.Background()
suite.loadAll()
releasePartitionErr := fmt.Errorf("mock release partitions error")
suite.cluster.ExpectedCalls = lo.Filter(suite.cluster.ExpectedCalls, func(call *mock.Call, _ int) bool {
return call.Method != "ReleasePartitions"
})
suite.cluster.EXPECT().ReleasePartitions(mock.Anything, mock.Anything, mock.Anything).
Return(nil, releasePartitionErr)
for _, collection := range suite.collections {
releaseCollectionReq := &querypb.ReleaseCollectionRequest{
CollectionID: collection,
}
releaseCollectionJob := NewReleaseCollectionJob(
ctx,
releaseCollectionReq,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.checkerController,
)
suite.scheduler.Add(releaseCollectionJob)
err := releaseCollectionJob.Wait()
suite.NoError(err)
releasePartitionReq := &querypb.ReleasePartitionsRequest{
CollectionID: collection,
PartitionIDs: suite.partitions[collection],
}
releasePartitionJob := NewReleasePartitionJob(
ctx,
releasePartitionReq,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.checkerController,
)
suite.scheduler.Add(releasePartitionJob)
err = releasePartitionJob.Wait()
suite.NoError(err)
}
suite.cluster.ExpectedCalls = lo.Filter(suite.cluster.ExpectedCalls, func(call *mock.Call, _ int) bool {
return call.Method != "ReleasePartitions"
})
suite.cluster.EXPECT().ReleasePartitions(mock.Anything, mock.Anything, mock.Anything).
Return(merr.Success(), nil)
}
func (suite *JobSuite) TestSyncNewCreatedPartition() {
newPartition := int64(999)
// test sync new created partition
suite.loadAll()
req := &querypb.SyncNewCreatedPartitionRequest{
CollectionID: suite.collections[0],
PartitionID: newPartition,
}
job := NewSyncNewCreatedPartitionJob(
context.Background(),
req,
suite.meta,
suite.cluster,
suite.broker,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.NoError(err)
partition := suite.meta.CollectionManager.GetPartition(newPartition)
suite.NotNil(partition)
suite.Equal(querypb.LoadStatus_Loaded, partition.GetStatus())
// test collection not loaded
req = &querypb.SyncNewCreatedPartitionRequest{
CollectionID: int64(888),
PartitionID: newPartition,
}
job = NewSyncNewCreatedPartitionJob(
context.Background(),
req,
suite.meta,
suite.cluster,
suite.broker,
)
suite.scheduler.Add(job)
err = job.Wait()
suite.NoError(err)
// test collection loaded, but its loadType is loadPartition
req = &querypb.SyncNewCreatedPartitionRequest{
CollectionID: suite.collections[1],
PartitionID: newPartition,
}
job = NewSyncNewCreatedPartitionJob(
context.Background(),
req,
suite.meta,
suite.cluster,
suite.broker,
)
suite.scheduler.Add(job)
err = job.Wait()
suite.NoError(err)
}
func (suite *JobSuite) loadAll() {
ctx := context.Background()
for _, collection := range suite.collections {
if suite.loadTypes[collection] == querypb.LoadType_LoadCollection {
req := &querypb.LoadCollectionRequest{
CollectionID: collection,
}
job := NewLoadCollectionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.NoError(err)
suite.EqualValues(1, suite.meta.GetReplicaNumber(collection))
suite.True(suite.meta.Exist(collection))
suite.NotNil(suite.meta.GetCollection(collection))
suite.NotNil(suite.meta.GetPartitionsByCollection(collection))
suite.targetMgr.UpdateCollectionCurrentTarget(collection)
} else {
req := &querypb.LoadPartitionsRequest{
CollectionID: collection,
PartitionIDs: suite.partitions[collection],
}
job := NewLoadPartitionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.collectionObserver,
suite.nodeMgr,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.NoError(err)
suite.EqualValues(1, suite.meta.GetReplicaNumber(collection))
suite.True(suite.meta.Exist(collection))
suite.NotNil(suite.meta.GetCollection(collection))
suite.NotNil(suite.meta.GetPartitionsByCollection(collection))
suite.targetMgr.UpdateCollectionCurrentTarget(collection)
}
}
}
func (suite *JobSuite) releaseAll() {
ctx := context.Background()
for _, collection := range suite.collections {
req := &querypb.ReleaseCollectionRequest{
CollectionID: collection,
}
job := NewReleaseCollectionJob(
ctx,
req,
suite.dist,
suite.meta,
suite.broker,
suite.cluster,
suite.targetMgr,
suite.targetObserver,
suite.checkerController,
)
suite.scheduler.Add(job)
err := job.Wait()
suite.NoError(err)
suite.assertCollectionReleased(collection)
}
}
func (suite *JobSuite) assertCollectionLoaded(collection int64) {
suite.True(suite.meta.Exist(collection))
suite.NotEqual(0, len(suite.meta.ReplicaManager.GetByCollection(collection)))
for _, channel := range suite.channels[collection] {
suite.NotNil(suite.targetMgr.GetDmChannel(collection, channel, meta.CurrentTarget))
}
for _, segments := range suite.segments[collection] {
for _, segment := range segments {
suite.NotNil(suite.targetMgr.GetSealedSegment(collection, segment, meta.CurrentTarget))
}
}
}
func (suite *JobSuite) assertPartitionLoaded(collection int64, partitionIDs ...int64) {
suite.True(suite.meta.Exist(collection))
suite.NotEqual(0, len(suite.meta.ReplicaManager.GetByCollection(collection)))
for _, channel := range suite.channels[collection] {
suite.NotNil(suite.targetMgr.GetDmChannel(collection, channel, meta.CurrentTarget))
}
for partitionID, segments := range suite.segments[collection] {
if !lo.Contains(partitionIDs, partitionID) {
continue
}
suite.NotNil(suite.meta.GetPartition(partitionID))
for _, segment := range segments {
suite.NotNil(suite.targetMgr.GetSealedSegment(collection, segment, meta.CurrentTarget))
}
}
}
func (suite *JobSuite) assertCollectionReleased(collection int64) {
suite.False(suite.meta.Exist(collection))
suite.Equal(0, len(suite.meta.ReplicaManager.GetByCollection(collection)))
for _, channel := range suite.channels[collection] {
suite.Nil(suite.targetMgr.GetDmChannel(collection, channel, meta.CurrentTarget))
}
for _, partitions := range suite.segments[collection] {
for _, segment := range partitions {
suite.Nil(suite.targetMgr.GetSealedSegment(collection, segment, meta.CurrentTarget))
}
}
}
func (suite *JobSuite) assertPartitionReleased(collection int64, partitionIDs ...int64) {
for _, partition := range partitionIDs {
suite.Nil(suite.meta.GetPartition(partition))
segments := suite.segments[collection][partition]
for _, segment := range segments {
suite.Nil(suite.targetMgr.GetSealedSegment(collection, segment, meta.CurrentTarget))
}
}
}
func TestJob(t *testing.T) {
suite.Run(t, new(JobSuite))
}