2022-10-11 11:39:22 +08:00
|
|
|
// 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.
|
|
|
|
|
2022-09-15 18:48:32 +08:00
|
|
|
package job
|
|
|
|
|
|
|
|
import (
|
|
|
|
"context"
|
2023-05-18 09:17:23 +08:00
|
|
|
"fmt"
|
2022-09-15 18:48:32 +08:00
|
|
|
"testing"
|
|
|
|
|
2023-02-26 11:31:49 +08:00
|
|
|
"github.com/cockroachdb/errors"
|
2023-03-20 14:55:57 +08:00
|
|
|
"github.com/samber/lo"
|
2022-11-07 19:37:04 +08:00
|
|
|
"github.com/stretchr/testify/mock"
|
|
|
|
"github.com/stretchr/testify/suite"
|
|
|
|
|
2024-04-15 08:13:19 +08:00
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/rgpb"
|
2022-09-15 18:48:32 +08:00
|
|
|
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
|
2023-07-31 13:57:04 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/metastore"
|
|
|
|
"github.com/milvus-io/milvus/internal/metastore/kv/querycoord"
|
|
|
|
"github.com/milvus-io/milvus/internal/metastore/mocks"
|
2022-09-15 18:48:32 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/proto/datapb"
|
|
|
|
"github.com/milvus-io/milvus/internal/proto/querypb"
|
2023-05-30 17:41:28 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/querycoordv2/checkers"
|
2022-09-15 18:48:32 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/querycoordv2/meta"
|
2023-01-14 21:55:41 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/querycoordv2/observers"
|
2022-09-15 18:48:32 +08:00
|
|
|
. "github.com/milvus-io/milvus/internal/querycoordv2/params"
|
|
|
|
"github.com/milvus-io/milvus/internal/querycoordv2/session"
|
2024-08-20 16:49:02 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/util/proxyutil"
|
2024-06-25 21:18:15 +08:00
|
|
|
"github.com/milvus-io/milvus/pkg/kv"
|
2023-04-06 19:14:32 +08:00
|
|
|
"github.com/milvus-io/milvus/pkg/util/etcd"
|
2023-07-17 14:59:34 +08:00
|
|
|
"github.com/milvus-io/milvus/pkg/util/merr"
|
2023-09-05 10:31:48 +08:00
|
|
|
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
2022-09-15 18:48:32 +08:00
|
|
|
)
|
|
|
|
|
2022-10-20 11:29:27 +08:00
|
|
|
const (
|
|
|
|
defaultVecFieldID = 1
|
|
|
|
defaultIndexID = 1
|
|
|
|
)
|
|
|
|
|
2022-09-15 18:48:32 +08:00
|
|
|
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
|
2024-04-22 10:39:22 +08:00
|
|
|
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
|
2024-08-20 16:49:02 +08:00
|
|
|
proxyManager *proxyutil.MockProxyClientManager
|
2022-09-15 18:48:32 +08:00
|
|
|
|
|
|
|
// Test objects
|
|
|
|
scheduler *Scheduler
|
|
|
|
}
|
|
|
|
|
|
|
|
func (suite *JobSuite) SetupSuite() {
|
2023-09-05 10:31:48 +08:00
|
|
|
paramtable.Init()
|
2022-09-15 18:48:32 +08:00
|
|
|
|
|
|
|
suite.collections = []int64{1000, 1001}
|
|
|
|
suite.partitions = map[int64][]int64{
|
2023-03-20 14:55:57 +08:00
|
|
|
1000: {100, 101, 102},
|
|
|
|
1001: {103, 104, 105},
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
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},
|
2023-03-20 14:55:57 +08:00
|
|
|
102: {5, 6},
|
2022-09-15 18:48:32 +08:00
|
|
|
},
|
|
|
|
1001: {
|
|
|
|
103: {7, 8},
|
2023-03-20 14:55:57 +08:00
|
|
|
104: {9, 10},
|
|
|
|
105: {11, 12},
|
2022-09-15 18:48:32 +08:00
|
|
|
},
|
|
|
|
}
|
|
|
|
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,
|
|
|
|
})
|
|
|
|
}
|
2023-04-18 18:30:32 +08:00
|
|
|
|
|
|
|
segmentBinlogs := []*datapb.SegmentInfo{}
|
2022-09-15 18:48:32 +08:00
|
|
|
for partition, segments := range partitions {
|
|
|
|
for _, segment := range segments {
|
2023-04-18 18:30:32 +08:00
|
|
|
segmentBinlogs = append(segmentBinlogs, &datapb.SegmentInfo{
|
|
|
|
ID: segment,
|
|
|
|
PartitionID: partition,
|
2022-09-15 18:48:32 +08:00
|
|
|
InsertChannel: suite.channels[collection][segment%2],
|
|
|
|
})
|
|
|
|
}
|
|
|
|
}
|
2023-09-26 09:23:25 +08:00
|
|
|
suite.broker.EXPECT().GetRecoveryInfoV2(mock.Anything, collection).Return(vChannels, segmentBinlogs, nil).Maybe()
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
2023-03-20 14:55:57 +08:00
|
|
|
|
2023-11-02 23:52:16 +08:00
|
|
|
suite.broker.EXPECT().DescribeCollection(mock.Anything, mock.Anything).
|
2023-05-18 09:17:23 +08:00
|
|
|
Return(nil, nil)
|
2024-03-07 21:43:03 +08:00
|
|
|
suite.broker.EXPECT().ListIndexes(mock.Anything, mock.Anything).
|
2023-05-18 09:17:23 +08:00
|
|
|
Return(nil, nil)
|
|
|
|
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster = session.NewMockCluster(suite.T())
|
|
|
|
suite.cluster.EXPECT().
|
|
|
|
LoadPartitions(mock.Anything, mock.Anything, mock.Anything).
|
2023-10-11 21:01:35 +08:00
|
|
|
Return(merr.Success(), nil)
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster.EXPECT().
|
|
|
|
ReleasePartitions(mock.Anything, mock.Anything, mock.Anything).
|
2023-10-11 21:01:35 +08:00
|
|
|
Return(merr.Success(), nil).Maybe()
|
2024-08-20 16:49:02 +08:00
|
|
|
|
|
|
|
suite.proxyManager = proxyutil.NewMockProxyClientManager(suite.T())
|
|
|
|
suite.proxyManager.EXPECT().InvalidateCollectionMetaCache(mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func (suite *JobSuite) SetupTest() {
|
|
|
|
config := GenerateEtcdConfig()
|
2022-11-30 18:23:15 +08:00
|
|
|
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())
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.Require().NoError(err)
|
2022-11-17 18:59:09 +08:00
|
|
|
suite.kv = etcdkv.NewEtcdKV(cli, config.MetaRootPath.GetValue())
|
2022-09-15 18:48:32 +08:00
|
|
|
|
2023-07-31 13:57:04 +08:00
|
|
|
suite.store = querycoord.NewCatalog(suite.kv)
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.dist = meta.NewDistributionManager()
|
2023-01-30 10:19:48 +08:00
|
|
|
suite.nodeMgr = session.NewNodeManager()
|
|
|
|
suite.meta = meta.NewMeta(RandomIncrementIDAllocator(), suite.store, suite.nodeMgr)
|
2022-11-07 19:37:04 +08:00
|
|
|
suite.targetMgr = meta.NewTargetManager(suite.broker, suite.meta)
|
2023-01-14 21:55:41 +08:00
|
|
|
suite.targetObserver = observers.NewTargetObserver(suite.meta,
|
|
|
|
suite.targetMgr,
|
|
|
|
suite.dist,
|
|
|
|
suite.broker,
|
2023-11-08 11:36:22 +08:00
|
|
|
suite.cluster,
|
2023-01-14 21:55:41 +08:00
|
|
|
)
|
2023-09-27 16:27:27 +08:00
|
|
|
suite.targetObserver.Start()
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.scheduler = NewScheduler()
|
|
|
|
|
2023-09-27 16:27:27 +08:00
|
|
|
suite.scheduler.Start()
|
2023-01-10 20:35:39 +08:00
|
|
|
meta.GlobalFailedLoadCache = meta.NewFailedLoadCache()
|
2023-01-30 10:19:48 +08:00
|
|
|
|
2024-03-15 10:45:06 +08:00
|
|
|
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",
|
|
|
|
}))
|
2024-04-15 08:13:19 +08:00
|
|
|
|
|
|
|
suite.meta.HandleNodeUp(1000)
|
|
|
|
suite.meta.HandleNodeUp(2000)
|
|
|
|
suite.meta.HandleNodeUp(3000)
|
2023-05-30 17:41:28 +08:00
|
|
|
|
|
|
|
suite.checkerController = &checkers.CheckerController{}
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver = observers.NewCollectionObserver(
|
|
|
|
suite.dist,
|
|
|
|
suite.meta,
|
|
|
|
suite.targetMgr,
|
|
|
|
suite.targetObserver,
|
|
|
|
suite.checkerController,
|
2024-08-20 16:49:02 +08:00
|
|
|
suite.proxyManager,
|
2024-04-22 10:39:22 +08:00
|
|
|
)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func (suite *JobSuite) TearDownTest() {
|
|
|
|
suite.kv.Close()
|
|
|
|
suite.scheduler.Stop()
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver.Stop()
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func (suite *JobSuite) BeforeTest(suiteName, testName string) {
|
2023-03-20 14:55:57 +08:00
|
|
|
for collection, partitions := range suite.partitions {
|
|
|
|
suite.broker.EXPECT().
|
|
|
|
GetPartitions(mock.Anything, collection).
|
|
|
|
Return(partitions, nil)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
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,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.targetMgr,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.nodeMgr,
|
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err := job.Wait()
|
|
|
|
suite.NoError(err)
|
|
|
|
suite.EqualValues(1, suite.meta.GetReplicaNumber(collection))
|
2022-11-07 19:37:04 +08:00
|
|
|
suite.targetMgr.UpdateCollectionCurrentTarget(collection)
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.assertCollectionLoaded(collection)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
// 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,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.targetMgr,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.nodeMgr,
|
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err := job.Wait()
|
2023-07-17 14:59:34 +08:00
|
|
|
suite.NoError(err)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
// 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,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.targetMgr,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.nodeMgr,
|
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err := job.Wait()
|
2023-07-17 14:59:34 +08:00
|
|
|
suite.ErrorIs(err, merr.ErrParameterInvalid)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
|
2024-09-02 18:39:03 +08:00
|
|
|
// 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)
|
|
|
|
}
|
|
|
|
|
2022-09-15 18:48:32 +08:00
|
|
|
// 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,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.targetMgr,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.nodeMgr,
|
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err := job.Wait()
|
2023-07-17 14:59:34 +08:00
|
|
|
suite.NoError(err)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
2023-01-30 10:19:48 +08:00
|
|
|
|
2024-04-15 08:13:19 +08:00
|
|
|
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)
|
2023-01-30 10:19:48 +08:00
|
|
|
|
|
|
|
// 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,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2023-01-30 10:19:48 +08:00
|
|
|
suite.targetMgr,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2023-01-30 10:19:48 +08:00
|
|
|
suite.nodeMgr,
|
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err := job.Wait()
|
2024-08-20 10:30:55 +08:00
|
|
|
suite.ErrorIs(err, merr.ErrResourceGroupNodeNotEnough)
|
2023-01-30 10:19:48 +08:00
|
|
|
|
|
|
|
// Load with 3 replica on 3 rg
|
|
|
|
req = &querypb.LoadCollectionRequest{
|
2023-03-20 14:55:57 +08:00
|
|
|
CollectionID: 1001,
|
2023-01-30 10:19:48 +08:00
|
|
|
ReplicaNumber: 3,
|
|
|
|
ResourceGroups: []string{"rg1", "rg2", "rg3"},
|
|
|
|
}
|
|
|
|
job = NewLoadCollectionJob(
|
|
|
|
ctx,
|
|
|
|
req,
|
|
|
|
suite.dist,
|
|
|
|
suite.meta,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2023-01-30 10:19:48 +08:00
|
|
|
suite.targetMgr,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2023-01-30 10:19:48 +08:00
|
|
|
suite.nodeMgr,
|
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err = job.Wait()
|
2024-08-20 10:30:55 +08:00
|
|
|
suite.ErrorIs(err, merr.ErrResourceGroupNodeNotEnough)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
|
2022-09-21 17:54:50 +08:00
|
|
|
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,
|
2023-01-30 10:19:48 +08:00
|
|
|
ReplicaNumber: 5,
|
2022-09-21 17:54:50 +08:00
|
|
|
}
|
|
|
|
job := NewLoadCollectionJob(
|
|
|
|
ctx,
|
|
|
|
req,
|
|
|
|
suite.dist,
|
|
|
|
suite.meta,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2022-09-21 17:54:50 +08:00
|
|
|
suite.targetMgr,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2022-09-21 17:54:50 +08:00
|
|
|
suite.nodeMgr,
|
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err := job.Wait()
|
2024-08-20 10:30:55 +08:00
|
|
|
suite.ErrorIs(err, merr.ErrResourceGroupNodeNotEnough)
|
2022-09-21 17:54:50 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2022-09-15 18:48:32 +08:00
|
|
|
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{
|
2022-11-07 19:37:04 +08:00
|
|
|
CollectionID: collection,
|
|
|
|
PartitionIDs: suite.partitions[collection],
|
|
|
|
ReplicaNumber: 1,
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
job := NewLoadPartitionJob(
|
|
|
|
ctx,
|
|
|
|
req,
|
|
|
|
suite.dist,
|
|
|
|
suite.meta,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.targetMgr,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.nodeMgr,
|
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err := job.Wait()
|
|
|
|
suite.NoError(err)
|
|
|
|
suite.EqualValues(1, suite.meta.GetReplicaNumber(collection))
|
2023-08-04 10:31:06 +08:00
|
|
|
suite.targetMgr.UpdateCollectionCurrentTarget(collection)
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.assertCollectionLoaded(collection)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
// 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,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.targetMgr,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.nodeMgr,
|
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err := job.Wait()
|
2023-07-17 14:59:34 +08:00
|
|
|
suite.NoError(err)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
// 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(
|
2024-09-02 18:39:03 +08:00
|
|
|
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(
|
2022-09-15 18:48:32 +08:00
|
|
|
ctx,
|
|
|
|
req,
|
|
|
|
suite.dist,
|
|
|
|
suite.meta,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.targetMgr,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.nodeMgr,
|
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err := job.Wait()
|
2023-07-17 14:59:34 +08:00
|
|
|
suite.ErrorIs(err, merr.ErrParameterInvalid)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
// 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),
|
2023-03-20 14:55:57 +08:00
|
|
|
ReplicaNumber: 1,
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
job := NewLoadPartitionJob(
|
|
|
|
ctx,
|
|
|
|
req,
|
|
|
|
suite.dist,
|
|
|
|
suite.meta,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.targetMgr,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.nodeMgr,
|
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err := job.Wait()
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.NoError(err)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
// 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,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.targetMgr,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.nodeMgr,
|
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err := job.Wait()
|
2023-07-17 14:59:34 +08:00
|
|
|
suite.NoError(err)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
2023-01-30 10:19:48 +08:00
|
|
|
|
2024-04-15 08:13:19 +08:00
|
|
|
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)
|
2023-01-30 10:19:48 +08:00
|
|
|
|
|
|
|
// test load 3 replica in 1 rg, should pass rg check
|
|
|
|
req := &querypb.LoadPartitionsRequest{
|
2023-03-20 14:55:57 +08:00
|
|
|
CollectionID: 999,
|
|
|
|
PartitionIDs: []int64{888},
|
2023-01-30 10:19:48 +08:00
|
|
|
ReplicaNumber: 3,
|
|
|
|
ResourceGroups: []string{"rg1"},
|
|
|
|
}
|
|
|
|
job := NewLoadPartitionJob(
|
|
|
|
ctx,
|
|
|
|
req,
|
|
|
|
suite.dist,
|
|
|
|
suite.meta,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2023-01-30 10:19:48 +08:00
|
|
|
suite.targetMgr,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2023-01-30 10:19:48 +08:00
|
|
|
suite.nodeMgr,
|
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err := job.Wait()
|
2024-08-20 10:30:55 +08:00
|
|
|
suite.ErrorIs(err, merr.ErrResourceGroupNodeNotEnough)
|
2023-01-30 10:19:48 +08:00
|
|
|
|
|
|
|
// test load 3 replica in 3 rg, should pass rg check
|
|
|
|
req = &querypb.LoadPartitionsRequest{
|
2023-03-20 14:55:57 +08:00
|
|
|
CollectionID: 999,
|
|
|
|
PartitionIDs: []int64{888},
|
2023-01-30 10:19:48 +08:00
|
|
|
ReplicaNumber: 3,
|
|
|
|
ResourceGroups: []string{"rg1", "rg2", "rg3"},
|
|
|
|
}
|
|
|
|
job = NewLoadPartitionJob(
|
|
|
|
ctx,
|
|
|
|
req,
|
|
|
|
suite.dist,
|
|
|
|
suite.meta,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2023-01-30 10:19:48 +08:00
|
|
|
suite.targetMgr,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2023-01-30 10:19:48 +08:00
|
|
|
suite.nodeMgr,
|
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err = job.Wait()
|
2024-08-20 10:30:55 +08:00
|
|
|
suite.ErrorIs(err, merr.ErrResourceGroupNodeNotEnough)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
|
2023-03-20 14:55:57 +08:00
|
|
|
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,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
|
|
|
suite.targetMgr,
|
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.nodeMgr,
|
|
|
|
)
|
|
|
|
return job
|
|
|
|
}
|
|
|
|
newLoadColJob := func() *LoadCollectionJob {
|
|
|
|
req := &querypb.LoadCollectionRequest{
|
|
|
|
CollectionID: collection,
|
|
|
|
ReplicaNumber: 1,
|
|
|
|
}
|
|
|
|
job := NewLoadCollectionJob(
|
|
|
|
ctx,
|
|
|
|
req,
|
|
|
|
suite.dist,
|
|
|
|
suite.meta,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
|
|
|
suite.targetMgr,
|
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2023-03-20 14:55:57 +08:00
|
|
|
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()
|
2023-07-17 14:59:34 +08:00
|
|
|
suite.NoError(err)
|
2023-03-20 14:55:57 +08:00
|
|
|
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)
|
|
|
|
}
|
|
|
|
|
2022-09-21 17:54:50 +08:00
|
|
|
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],
|
2023-01-30 10:19:48 +08:00
|
|
|
ReplicaNumber: 5,
|
2022-09-21 17:54:50 +08:00
|
|
|
}
|
|
|
|
job := NewLoadPartitionJob(
|
|
|
|
ctx,
|
|
|
|
req,
|
|
|
|
suite.dist,
|
|
|
|
suite.meta,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2022-09-21 17:54:50 +08:00
|
|
|
suite.targetMgr,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2022-09-21 17:54:50 +08:00
|
|
|
suite.nodeMgr,
|
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err := job.Wait()
|
2024-08-20 10:30:55 +08:00
|
|
|
suite.ErrorIs(err, merr.ErrResourceGroupNodeNotEnough)
|
2022-09-21 17:54:50 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2022-09-15 18:48:32 +08:00
|
|
|
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,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
|
|
|
suite.cluster,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.targetMgr,
|
2023-01-14 21:55:41 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
|
2023-05-30 17:41:28 +08:00
|
|
|
suite.checkerController,
|
2022-09-15 18:48:32 +08:00
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err := job.Wait()
|
|
|
|
suite.NoError(err)
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.assertCollectionReleased(collection)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
// Test release again
|
|
|
|
for _, collection := range suite.collections {
|
|
|
|
req := &querypb.ReleaseCollectionRequest{
|
|
|
|
CollectionID: collection,
|
|
|
|
}
|
|
|
|
job := NewReleaseCollectionJob(
|
|
|
|
ctx,
|
|
|
|
req,
|
|
|
|
suite.dist,
|
|
|
|
suite.meta,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
|
|
|
suite.cluster,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.targetMgr,
|
2023-01-14 21:55:41 +08:00
|
|
|
suite.targetObserver,
|
2023-05-30 17:41:28 +08:00
|
|
|
suite.checkerController,
|
2022-09-15 18:48:32 +08:00
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err := job.Wait()
|
|
|
|
suite.NoError(err)
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.assertCollectionReleased(collection)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
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,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.targetMgr,
|
2023-01-14 21:55:41 +08:00
|
|
|
suite.targetObserver,
|
2023-05-30 17:41:28 +08:00
|
|
|
suite.checkerController,
|
2022-09-15 18:48:32 +08:00
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err := job.Wait()
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.NoError(err)
|
|
|
|
suite.assertPartitionReleased(collection, suite.partitions[collection]...)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
// 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,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.targetMgr,
|
2023-01-14 21:55:41 +08:00
|
|
|
suite.targetObserver,
|
2023-05-30 17:41:28 +08:00
|
|
|
suite.checkerController,
|
2022-09-15 18:48:32 +08:00
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err := job.Wait()
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.NoError(err)
|
|
|
|
suite.assertPartitionReleased(collection, suite.partitions[collection]...)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
// 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,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.targetMgr,
|
2023-01-14 21:55:41 +08:00
|
|
|
suite.targetObserver,
|
2023-05-30 17:41:28 +08:00
|
|
|
suite.checkerController,
|
2022-09-15 18:48:32 +08:00
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err := job.Wait()
|
2023-03-20 14:55:57 +08:00
|
|
|
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,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
|
|
|
suite.targetMgr,
|
|
|
|
suite.targetObserver,
|
2023-05-30 17:41:28 +08:00
|
|
|
suite.checkerController,
|
2023-03-20 14:55:57 +08:00
|
|
|
)
|
|
|
|
return job
|
|
|
|
}
|
|
|
|
newReleaseColJob := func(col int64) *ReleaseCollectionJob {
|
|
|
|
req := &querypb.ReleaseCollectionRequest{
|
|
|
|
CollectionID: col,
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
2023-03-20 14:55:57 +08:00
|
|
|
job := NewReleaseCollectionJob(
|
|
|
|
ctx,
|
|
|
|
req,
|
|
|
|
suite.dist,
|
|
|
|
suite.meta,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
|
|
|
suite.cluster,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetMgr,
|
|
|
|
suite.targetObserver,
|
2023-05-30 17:41:28 +08:00
|
|
|
suite.checkerController,
|
2023-03-20 14:55:57 +08:00
|
|
|
)
|
|
|
|
return job
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
2023-03-20 14:55:57 +08:00
|
|
|
|
|
|
|
// 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)
|
2023-07-19 16:54:58 +08:00
|
|
|
suite.False(suite.meta.Exist(col0))
|
2023-03-20 14:55:57 +08:00
|
|
|
|
|
|
|
// 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)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func (suite *JobSuite) TestLoadCollectionStoreFailed() {
|
|
|
|
// Store collection failed
|
2023-07-31 13:57:04 +08:00
|
|
|
store := mocks.NewQueryCoordCatalog(suite.T())
|
2023-01-30 10:19:48 +08:00
|
|
|
suite.meta = meta.NewMeta(RandomIncrementIDAllocator(), store, suite.nodeMgr)
|
|
|
|
|
|
|
|
store.EXPECT().SaveResourceGroup(mock.Anything, mock.Anything).Return(nil)
|
2024-04-15 08:13:19 +08:00
|
|
|
suite.meta.HandleNodeUp(1000)
|
|
|
|
suite.meta.HandleNodeUp(2000)
|
|
|
|
suite.meta.HandleNodeUp(3000)
|
2023-01-30 10:19:48 +08:00
|
|
|
|
2022-09-15 18:48:32 +08:00
|
|
|
for _, collection := range suite.collections {
|
|
|
|
if suite.loadTypes[collection] != querypb.LoadType_LoadCollection {
|
|
|
|
continue
|
|
|
|
}
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.broker.EXPECT().GetPartitions(mock.Anything, collection).Return(suite.partitions[collection], nil)
|
2022-09-15 18:48:32 +08:00
|
|
|
err := errors.New("failed to store collection")
|
|
|
|
store.EXPECT().SaveReplica(mock.Anything).Return(nil)
|
2023-03-20 14:55:57 +08:00
|
|
|
store.EXPECT().SaveCollection(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(err)
|
2022-09-15 18:48:32 +08:00
|
|
|
store.EXPECT().ReleaseReplicas(collection).Return(nil)
|
|
|
|
|
|
|
|
req := &querypb.LoadCollectionRequest{
|
|
|
|
CollectionID: collection,
|
|
|
|
}
|
|
|
|
job := NewLoadCollectionJob(
|
|
|
|
context.Background(),
|
|
|
|
req,
|
|
|
|
suite.dist,
|
|
|
|
suite.meta,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.targetMgr,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.nodeMgr,
|
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
loadErr := job.Wait()
|
|
|
|
suite.ErrorIs(loadErr, err)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func (suite *JobSuite) TestLoadPartitionStoreFailed() {
|
|
|
|
// Store partition failed
|
2023-07-31 13:57:04 +08:00
|
|
|
store := mocks.NewQueryCoordCatalog(suite.T())
|
2023-01-30 10:19:48 +08:00
|
|
|
suite.meta = meta.NewMeta(RandomIncrementIDAllocator(), store, suite.nodeMgr)
|
|
|
|
|
|
|
|
store.EXPECT().SaveResourceGroup(mock.Anything, mock.Anything).Return(nil)
|
2024-04-15 08:13:19 +08:00
|
|
|
suite.meta.HandleNodeUp(1000)
|
|
|
|
suite.meta.HandleNodeUp(2000)
|
|
|
|
suite.meta.HandleNodeUp(3000)
|
2023-01-30 10:19:48 +08:00
|
|
|
|
2024-04-15 08:13:19 +08:00
|
|
|
err := errors.New("failed to store collection")
|
2022-09-15 18:48:32 +08:00
|
|
|
for _, collection := range suite.collections {
|
|
|
|
if suite.loadTypes[collection] != querypb.LoadType_LoadPartition {
|
|
|
|
continue
|
|
|
|
}
|
|
|
|
|
|
|
|
store.EXPECT().SaveReplica(mock.Anything).Return(nil)
|
2023-03-20 14:55:57 +08:00
|
|
|
store.EXPECT().SaveCollection(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(err)
|
2022-09-15 18:48:32 +08:00
|
|
|
store.EXPECT().ReleaseReplicas(collection).Return(nil)
|
|
|
|
|
|
|
|
req := &querypb.LoadPartitionsRequest{
|
|
|
|
CollectionID: collection,
|
|
|
|
PartitionIDs: suite.partitions[collection],
|
|
|
|
}
|
|
|
|
job := NewLoadPartitionJob(
|
|
|
|
context.Background(),
|
|
|
|
req,
|
|
|
|
suite.dist,
|
|
|
|
suite.meta,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.targetMgr,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.nodeMgr,
|
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
loadErr := job.Wait()
|
|
|
|
suite.ErrorIs(loadErr, err)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func (suite *JobSuite) TestLoadCreateReplicaFailed() {
|
|
|
|
// Store replica failed
|
2023-01-30 10:19:48 +08:00
|
|
|
suite.meta = meta.NewMeta(ErrorIDAllocator(), suite.store, session.NewNodeManager())
|
2022-09-15 18:48:32 +08:00
|
|
|
for _, collection := range suite.collections {
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.broker.EXPECT().
|
|
|
|
GetPartitions(mock.Anything, collection).
|
|
|
|
Return(suite.partitions[collection], nil)
|
2022-09-15 18:48:32 +08:00
|
|
|
req := &querypb.LoadCollectionRequest{
|
|
|
|
CollectionID: collection,
|
|
|
|
}
|
|
|
|
job := NewLoadCollectionJob(
|
|
|
|
context.Background(),
|
|
|
|
req,
|
|
|
|
suite.dist,
|
|
|
|
suite.meta,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.targetMgr,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.nodeMgr,
|
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err := job.Wait()
|
2024-08-20 10:30:55 +08:00
|
|
|
suite.ErrorIs(err, merr.ErrResourceGroupNodeNotEnough)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2023-05-18 09:17:23 +08:00
|
|
|
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 {
|
2024-03-07 21:43:03 +08:00
|
|
|
return call.Method != "ListIndexes"
|
2023-05-18 09:17:23 +08:00
|
|
|
})
|
|
|
|
for _, collection := range suite.collections {
|
2024-03-07 21:43:03 +08:00
|
|
|
suite.broker.EXPECT().ListIndexes(mock.Anything, collection).Return(nil, getIndexErr)
|
2023-05-18 09:17:23 +08:00
|
|
|
loadCollectionReq := &querypb.LoadCollectionRequest{
|
|
|
|
CollectionID: collection,
|
|
|
|
}
|
|
|
|
loadCollectionJob := NewLoadCollectionJob(
|
|
|
|
context.Background(),
|
|
|
|
loadCollectionReq,
|
|
|
|
suite.dist,
|
|
|
|
suite.meta,
|
|
|
|
suite.broker,
|
|
|
|
suite.cluster,
|
|
|
|
suite.targetMgr,
|
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2023-05-18 09:17:23 +08:00
|
|
|
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,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2023-05-18 09:17:23 +08:00
|
|
|
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 {
|
2023-11-02 23:52:16 +08:00
|
|
|
return call.Method != "DescribeCollection"
|
2023-05-18 09:17:23 +08:00
|
|
|
})
|
|
|
|
for _, collection := range suite.collections {
|
2023-11-02 23:52:16 +08:00
|
|
|
suite.broker.EXPECT().DescribeCollection(mock.Anything, collection).Return(nil, getSchemaErr)
|
2023-05-18 09:17:23 +08:00
|
|
|
loadCollectionReq := &querypb.LoadCollectionRequest{
|
|
|
|
CollectionID: collection,
|
|
|
|
}
|
|
|
|
loadCollectionJob := NewLoadCollectionJob(
|
|
|
|
context.Background(),
|
|
|
|
loadCollectionReq,
|
|
|
|
suite.dist,
|
|
|
|
suite.meta,
|
|
|
|
suite.broker,
|
|
|
|
suite.cluster,
|
|
|
|
suite.targetMgr,
|
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2023-05-18 09:17:23 +08:00
|
|
|
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,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2023-05-18 09:17:23 +08:00
|
|
|
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 {
|
2024-03-07 21:43:03 +08:00
|
|
|
return call.Method != "ListIndexes" && call.Method != "DescribeCollection"
|
2023-05-18 09:17:23 +08:00
|
|
|
})
|
2023-11-02 23:52:16 +08:00
|
|
|
suite.broker.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(nil, nil)
|
2024-03-07 21:43:03 +08:00
|
|
|
suite.broker.EXPECT().ListIndexes(mock.Anything, mock.Anything).Return(nil, nil)
|
2023-05-18 09:17:23 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
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,
|
2023-05-30 17:41:28 +08:00
|
|
|
suite.checkerController,
|
2023-05-18 09:17:23 +08:00
|
|
|
)
|
|
|
|
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,
|
2023-05-30 17:41:28 +08:00
|
|
|
suite.checkerController,
|
2023-05-18 09:17:23 +08:00
|
|
|
)
|
|
|
|
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).
|
2023-10-11 21:01:35 +08:00
|
|
|
Return(merr.Success(), nil)
|
2023-05-18 09:17:23 +08:00
|
|
|
}
|
|
|
|
|
2023-03-20 14:55:57 +08:00
|
|
|
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,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
)
|
|
|
|
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,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err = job.Wait()
|
2023-07-17 14:59:34 +08:00
|
|
|
suite.NoError(err)
|
2023-03-20 14:55:57 +08:00
|
|
|
|
|
|
|
// 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,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err = job.Wait()
|
2023-07-17 14:59:34 +08:00
|
|
|
suite.NoError(err)
|
2023-03-20 14:55:57 +08:00
|
|
|
}
|
|
|
|
|
2022-09-15 18:48:32 +08:00
|
|
|
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,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.targetMgr,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2022-09-15 18:48:32 +08:00
|
|
|
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))
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.NotNil(suite.meta.GetPartitionsByCollection(collection))
|
2022-11-07 19:37:04 +08:00
|
|
|
suite.targetMgr.UpdateCollectionCurrentTarget(collection)
|
2022-09-15 18:48:32 +08:00
|
|
|
} else {
|
|
|
|
req := &querypb.LoadPartitionsRequest{
|
|
|
|
CollectionID: collection,
|
|
|
|
PartitionIDs: suite.partitions[collection],
|
|
|
|
}
|
|
|
|
job := NewLoadPartitionJob(
|
|
|
|
ctx,
|
|
|
|
req,
|
|
|
|
suite.dist,
|
|
|
|
suite.meta,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.cluster,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.targetMgr,
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.targetObserver,
|
2024-04-22 10:39:22 +08:00
|
|
|
suite.collectionObserver,
|
2022-09-15 18:48:32 +08:00
|
|
|
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))
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.NotNil(suite.meta.GetCollection(collection))
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.NotNil(suite.meta.GetPartitionsByCollection(collection))
|
2022-11-07 19:37:04 +08:00
|
|
|
suite.targetMgr.UpdateCollectionCurrentTarget(collection)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
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,
|
2023-05-18 09:17:23 +08:00
|
|
|
suite.broker,
|
|
|
|
suite.cluster,
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.targetMgr,
|
2023-01-14 21:55:41 +08:00
|
|
|
suite.targetObserver,
|
2023-05-30 17:41:28 +08:00
|
|
|
suite.checkerController,
|
2022-09-15 18:48:32 +08:00
|
|
|
)
|
|
|
|
suite.scheduler.Add(job)
|
|
|
|
err := job.Wait()
|
|
|
|
suite.NoError(err)
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.assertCollectionReleased(collection)
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2023-03-20 14:55:57 +08:00
|
|
|
func (suite *JobSuite) assertCollectionLoaded(collection int64) {
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.True(suite.meta.Exist(collection))
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.NotEqual(0, len(suite.meta.ReplicaManager.GetByCollection(collection)))
|
2022-09-15 18:48:32 +08:00
|
|
|
for _, channel := range suite.channels[collection] {
|
2022-11-07 19:37:04 +08:00
|
|
|
suite.NotNil(suite.targetMgr.GetDmChannel(collection, channel, meta.CurrentTarget))
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
2023-03-20 14:55:57 +08:00
|
|
|
for _, segments := range suite.segments[collection] {
|
|
|
|
for _, segment := range segments {
|
2023-10-25 00:44:12 +08:00
|
|
|
suite.NotNil(suite.targetMgr.GetSealedSegment(collection, segment, meta.CurrentTarget))
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2023-03-20 14:55:57 +08:00
|
|
|
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 {
|
2023-10-25 00:44:12 +08:00
|
|
|
suite.NotNil(suite.targetMgr.GetSealedSegment(collection, segment, meta.CurrentTarget))
|
2023-03-20 14:55:57 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func (suite *JobSuite) assertCollectionReleased(collection int64) {
|
2022-09-15 18:48:32 +08:00
|
|
|
suite.False(suite.meta.Exist(collection))
|
2023-03-20 14:55:57 +08:00
|
|
|
suite.Equal(0, len(suite.meta.ReplicaManager.GetByCollection(collection)))
|
2022-09-15 18:48:32 +08:00
|
|
|
for _, channel := range suite.channels[collection] {
|
2022-11-07 19:37:04 +08:00
|
|
|
suite.Nil(suite.targetMgr.GetDmChannel(collection, channel, meta.CurrentTarget))
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
for _, partitions := range suite.segments[collection] {
|
|
|
|
for _, segment := range partitions {
|
2023-10-25 00:44:12 +08:00
|
|
|
suite.Nil(suite.targetMgr.GetSealedSegment(collection, segment, meta.CurrentTarget))
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2023-03-20 14:55:57 +08:00
|
|
|
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 {
|
2023-10-25 00:44:12 +08:00
|
|
|
suite.Nil(suite.targetMgr.GetSealedSegment(collection, segment, meta.CurrentTarget))
|
2023-03-20 14:55:57 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2022-09-15 18:48:32 +08:00
|
|
|
func TestJob(t *testing.T) {
|
|
|
|
suite.Run(t, new(JobSuite))
|
|
|
|
}
|