mirror of
https://gitee.com/milvus-io/milvus.git
synced 2025-12-08 01:58:34 +08:00
39754af727
issue: #43794 Fix the issue where L0 segments were not correctly selecting appropriate delegators during loading, which could cause load failures or incorrect delegator assignments. Changes include: - Add special handling for L0 segments in delegator selection logic - Find delegators that are missing the L0 segment for direct loading - Fallback to existing serviceable delegator selection when no suitable delegator is found for L0 segments - Add comprehensive test coverage for L0 segment loading scenarios - Test delegator selection when some delegators are missing segments - Test fallback behavior when all delegators already have the segment - Test error handling when no delegators are available Signed-off-by: Wei Liu <wei.liu@zilliz.com>
680 lines
27 KiB
Go
680 lines
27 KiB
Go
// Licensed to the LF AI & Data foundation under one
|
|
// or more contributor license agreements. See the NOTICE file
|
|
// distributed with this work for additional information
|
|
// regarding copyright ownership. The ASF licenses this file
|
|
// to you under the Apache License, Version 2.0 (the
|
|
// "License"); you may not use this file except in compliance
|
|
// with the License. You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package task
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/bytedance/mockey"
|
|
"github.com/cockroachdb/errors"
|
|
"github.com/stretchr/testify/mock"
|
|
"github.com/stretchr/testify/suite"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
|
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
|
|
"github.com/milvus-io/milvus/internal/metastore/kv/querycoord"
|
|
"github.com/milvus-io/milvus/internal/querycoordv2/meta"
|
|
"github.com/milvus-io/milvus/internal/querycoordv2/params"
|
|
"github.com/milvus-io/milvus/internal/querycoordv2/session"
|
|
"github.com/milvus-io/milvus/internal/querycoordv2/utils"
|
|
"github.com/milvus-io/milvus/pkg/v2/kv"
|
|
"github.com/milvus-io/milvus/pkg/v2/proto/datapb"
|
|
"github.com/milvus-io/milvus/pkg/v2/proto/indexpb"
|
|
"github.com/milvus-io/milvus/pkg/v2/proto/querypb"
|
|
"github.com/milvus-io/milvus/pkg/v2/util/etcd"
|
|
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
|
|
)
|
|
|
|
type ExecutorTestSuite struct {
|
|
suite.Suite
|
|
|
|
// Dependencies
|
|
kv kv.MetaKv
|
|
meta *meta.Meta
|
|
dist *meta.DistributionManager
|
|
target *meta.TargetManager
|
|
broker *meta.MockBroker
|
|
nodeMgr *session.NodeManager
|
|
cluster *session.MockCluster
|
|
|
|
// Test object
|
|
executor *Executor
|
|
ctx context.Context
|
|
}
|
|
|
|
func (suite *ExecutorTestSuite) SetupSuite() {
|
|
paramtable.Init()
|
|
}
|
|
|
|
func (suite *ExecutorTestSuite) SetupTest() {
|
|
var err error
|
|
config := params.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())
|
|
|
|
// Initialize components
|
|
store := querycoord.NewCatalog(suite.kv)
|
|
idAllocator := params.RandomIncrementIDAllocator()
|
|
suite.nodeMgr = session.NewNodeManager()
|
|
suite.meta = meta.NewMeta(idAllocator, store, suite.nodeMgr)
|
|
suite.dist = meta.NewDistributionManager()
|
|
suite.broker = meta.NewMockBroker(suite.T())
|
|
suite.target = meta.NewTargetManager(suite.broker, suite.meta)
|
|
suite.cluster = session.NewMockCluster(suite.T())
|
|
|
|
suite.executor = NewExecutor(suite.meta, suite.dist, suite.broker, suite.target, suite.cluster, suite.nodeMgr)
|
|
suite.ctx = context.Background()
|
|
}
|
|
|
|
func (suite *ExecutorTestSuite) TearDownTest() {
|
|
suite.kv.Close()
|
|
}
|
|
|
|
func (suite *ExecutorTestSuite) TestReleaseSegmentWithServiceableLeader() {
|
|
// Setup collection and replica
|
|
collection := utils.CreateTestCollection(1, 1)
|
|
suite.meta.CollectionManager.PutCollection(suite.ctx, collection)
|
|
suite.meta.CollectionManager.PutPartition(suite.ctx, utils.CreateTestPartition(1, 1))
|
|
|
|
replica := utils.CreateTestReplica(1, 1, []int64{1, 2})
|
|
suite.meta.ReplicaManager.Put(suite.ctx, replica)
|
|
|
|
// Setup nodes
|
|
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 1,
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 2,
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 1)
|
|
suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 2)
|
|
|
|
// Create segment task
|
|
action := NewSegmentAction(1, ActionTypeReduce, "test-channel", 100)
|
|
task, err := NewSegmentTask(suite.ctx, time.Second*10, WrapIDSource(0), 1, replica, action)
|
|
suite.NoError(err)
|
|
|
|
// Setup serviceable leader view (UnServiceableError = nil means serviceable)
|
|
serviceableLeaderView := utils.CreateTestLeaderView(1, 1, "test-channel", map[int64]int64{100: 1}, map[int64]*meta.Segment{})
|
|
serviceableLeaderView.UnServiceableError = nil // serviceable
|
|
suite.dist.LeaderViewManager.Update(1, serviceableLeaderView)
|
|
|
|
// Setup non-serviceable leader view
|
|
nonServiceableLeaderView := utils.CreateTestLeaderView(2, 1, "test-channel", map[int64]int64{100: 2}, map[int64]*meta.Segment{})
|
|
nonServiceableLeaderView.UnServiceableError = errors.New("not serviceable") // not serviceable
|
|
suite.dist.LeaderViewManager.Update(2, nonServiceableLeaderView)
|
|
|
|
// Mock cluster response
|
|
suite.cluster.EXPECT().ReleaseSegments(mock.Anything, int64(1), mock.Anything).Return(&commonpb.Status{
|
|
ErrorCode: commonpb.ErrorCode_Success,
|
|
}, nil).Once()
|
|
|
|
// Execute release segment
|
|
suite.executor.releaseSegment(task, 0)
|
|
|
|
// Verify that serviceable leader (node 1) was chosen for release
|
|
suite.cluster.AssertExpectations(suite.T())
|
|
}
|
|
|
|
func (suite *ExecutorTestSuite) TestReleaseSegmentFallbackToNonServiceableLeader() {
|
|
// Setup collection and replica
|
|
collection := utils.CreateTestCollection(1, 1)
|
|
suite.meta.CollectionManager.PutCollection(suite.ctx, collection)
|
|
suite.meta.CollectionManager.PutPartition(suite.ctx, utils.CreateTestPartition(1, 1))
|
|
|
|
replica := utils.CreateTestReplica(1, 1, []int64{1, 2})
|
|
suite.meta.ReplicaManager.Put(suite.ctx, replica)
|
|
|
|
// Setup nodes
|
|
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 1,
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 2,
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 1)
|
|
suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 2)
|
|
|
|
// Create segment task
|
|
action := NewSegmentAction(1, ActionTypeReduce, "test-channel", 100)
|
|
task, err := NewSegmentTask(suite.ctx, time.Second*10, WrapIDSource(0), 1, replica, action)
|
|
suite.NoError(err)
|
|
|
|
// Setup only non-serviceable leader view
|
|
nonServiceableLeaderView := utils.CreateTestLeaderView(2, 1, "test-channel", map[int64]int64{100: 2}, map[int64]*meta.Segment{})
|
|
nonServiceableLeaderView.UnServiceableError = errors.New("not serviceable") // not serviceable
|
|
suite.dist.LeaderViewManager.Update(2, nonServiceableLeaderView)
|
|
|
|
// Mock cluster response for fallback to non-serviceable leader
|
|
suite.cluster.EXPECT().ReleaseSegments(mock.Anything, int64(2), mock.Anything).Return(&commonpb.Status{
|
|
ErrorCode: commonpb.ErrorCode_Success,
|
|
}, nil).Once()
|
|
|
|
// Execute release segment
|
|
suite.executor.releaseSegment(task, 0)
|
|
|
|
// Verify that non-serviceable leader (node 2) was used as fallback
|
|
suite.cluster.AssertExpectations(suite.T())
|
|
}
|
|
|
|
func (suite *ExecutorTestSuite) TestReleaseSegmentNoLeaderFound() {
|
|
// Setup collection and replica
|
|
collection := utils.CreateTestCollection(1, 1)
|
|
suite.meta.CollectionManager.PutCollection(suite.ctx, collection)
|
|
suite.meta.CollectionManager.PutPartition(suite.ctx, utils.CreateTestPartition(1, 1))
|
|
|
|
replica := utils.CreateTestReplica(1, 1, []int64{1, 2})
|
|
suite.meta.ReplicaManager.Put(suite.ctx, replica)
|
|
|
|
// Setup nodes
|
|
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 1,
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 2,
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 1)
|
|
suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 2)
|
|
|
|
// Create segment task
|
|
action := NewSegmentAction(1, ActionTypeReduce, "test-channel", 100)
|
|
task, err := NewSegmentTask(suite.ctx, time.Second*10, WrapIDSource(0), 1, replica, action)
|
|
suite.NoError(err)
|
|
|
|
// No leader views setup - should trigger early return
|
|
|
|
// Execute release segment - should return early without calling cluster
|
|
suite.executor.releaseSegment(task, 0)
|
|
|
|
// Verify that no cluster calls were made
|
|
suite.cluster.AssertExpectations(suite.T())
|
|
}
|
|
|
|
func (suite *ExecutorTestSuite) TestReleaseSegmentWithNodeNotInReplica() {
|
|
// Setup collection and replica
|
|
collection := utils.CreateTestCollection(1, 1)
|
|
suite.meta.CollectionManager.PutCollection(suite.ctx, collection)
|
|
suite.meta.CollectionManager.PutPartition(suite.ctx, utils.CreateTestPartition(1, 1))
|
|
|
|
replica := utils.CreateTestReplica(1, 1, []int64{1, 2})
|
|
suite.meta.ReplicaManager.Put(suite.ctx, replica)
|
|
|
|
// Setup nodes
|
|
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 3, // Node 3 not in replica
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 3)
|
|
|
|
// Create segment task with node not in replica
|
|
action := NewSegmentAction(3, ActionTypeReduce, "test-channel", 100)
|
|
task, err := NewSegmentTask(suite.ctx, time.Second*10, WrapIDSource(0), 1, replica, action)
|
|
suite.NoError(err)
|
|
|
|
// Mock cluster response for direct node release
|
|
suite.cluster.EXPECT().ReleaseSegments(mock.Anything, int64(3), mock.Anything).Return(&commonpb.Status{
|
|
ErrorCode: commonpb.ErrorCode_Success,
|
|
}, nil).Once()
|
|
|
|
// Execute release segment
|
|
suite.executor.releaseSegment(task, 0)
|
|
|
|
// Verify that direct node release was used
|
|
suite.cluster.AssertExpectations(suite.T())
|
|
}
|
|
|
|
func (suite *ExecutorTestSuite) TestReleaseSegmentChannelSpecificLookup() {
|
|
// Setup collection and replica
|
|
collection := utils.CreateTestCollection(1, 1)
|
|
suite.meta.CollectionManager.PutCollection(suite.ctx, collection)
|
|
suite.meta.CollectionManager.PutPartition(suite.ctx, utils.CreateTestPartition(1, 1))
|
|
|
|
replica := utils.CreateTestReplica(1, 1, []int64{1, 2})
|
|
suite.meta.ReplicaManager.Put(suite.ctx, replica)
|
|
|
|
// Setup nodes
|
|
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 1,
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 2,
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 1)
|
|
suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 2)
|
|
|
|
// Create segment task
|
|
action := NewSegmentAction(1, ActionTypeReduce, "test-channel", 100)
|
|
task, err := NewSegmentTask(suite.ctx, time.Second*10, WrapIDSource(0), 1, replica, action)
|
|
suite.NoError(err)
|
|
|
|
// Setup leader view for a different channel (should not be found with channel filter)
|
|
leaderViewDifferentChannel := utils.CreateTestLeaderView(1, 1, "other-channel", map[int64]int64{100: 1}, map[int64]*meta.Segment{})
|
|
leaderViewDifferentChannel.UnServiceableError = errors.New("not serviceable")
|
|
suite.dist.LeaderViewManager.Update(1, leaderViewDifferentChannel)
|
|
|
|
// Setup leader view for the correct channel but non-serviceable
|
|
leaderViewCorrectChannel := utils.CreateTestLeaderView(2, 1, "test-channel", map[int64]int64{100: 2}, map[int64]*meta.Segment{})
|
|
leaderViewCorrectChannel.UnServiceableError = errors.New("not serviceable")
|
|
suite.dist.LeaderViewManager.Update(2, leaderViewCorrectChannel)
|
|
|
|
// Mock cluster response
|
|
suite.cluster.EXPECT().ReleaseSegments(mock.Anything, int64(2), mock.Anything).Return(&commonpb.Status{
|
|
ErrorCode: commonpb.ErrorCode_Success,
|
|
}, nil).Once()
|
|
|
|
// Execute release segment
|
|
suite.executor.releaseSegment(task, 0)
|
|
|
|
// Verify that channel-specific lookup worked correctly
|
|
suite.cluster.AssertExpectations(suite.T())
|
|
}
|
|
|
|
func (suite *ExecutorTestSuite) TestBalanceTaskWithTwoDelegators() {
|
|
// Setup collection and replica
|
|
collection := utils.CreateTestCollection(1, 1)
|
|
suite.meta.CollectionManager.PutCollection(suite.ctx, collection)
|
|
suite.meta.CollectionManager.PutPartition(suite.ctx, utils.CreateTestPartition(1, 1))
|
|
|
|
replica := utils.CreateTestReplica(1, 1, []int64{1, 2})
|
|
suite.meta.ReplicaManager.Put(suite.ctx, replica)
|
|
|
|
// Setup nodes
|
|
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 1,
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 2,
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 1)
|
|
suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 2)
|
|
|
|
// Create balance task with load and release actions
|
|
loadAction := NewSegmentAction(2, ActionTypeGrow, "test-channel", 100) // Load on node 2
|
|
releaseAction := NewSegmentAction(1, ActionTypeReduce, "test-channel", 100) // Release from node 1
|
|
|
|
task, err := NewSegmentTask(suite.ctx, time.Second*10, WrapIDSource(0), 1, replica, loadAction, releaseAction)
|
|
suite.NoError(err)
|
|
|
|
// Setup old delegator (node 1) - serviceable
|
|
oldDelegatorView := utils.CreateTestLeaderView(1, 1, "test-channel", map[int64]int64{100: 1}, map[int64]*meta.Segment{})
|
|
oldDelegatorView.Version = 1
|
|
oldDelegatorView.UnServiceableError = nil // serviceable
|
|
suite.dist.LeaderViewManager.Update(1, oldDelegatorView)
|
|
|
|
// Mock broker responses
|
|
suite.broker.EXPECT().DescribeCollection(mock.Anything, int64(1)).Return(&milvuspb.DescribeCollectionResponse{
|
|
Schema: &schemapb.CollectionSchema{
|
|
Name: "TestBalanceTask",
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
|
|
},
|
|
},
|
|
}, nil)
|
|
suite.broker.EXPECT().GetSegmentInfo(mock.Anything, int64(100)).Return([]*datapb.SegmentInfo{
|
|
{
|
|
ID: 100,
|
|
CollectionID: 1,
|
|
PartitionID: 1,
|
|
InsertChannel: "test-channel",
|
|
},
|
|
}, nil)
|
|
suite.broker.EXPECT().GetIndexInfo(mock.Anything, int64(1), int64(100)).Return(nil, nil)
|
|
suite.broker.EXPECT().ListIndexes(mock.Anything, int64(1)).Return([]*indexpb.IndexInfo{
|
|
{
|
|
CollectionID: 1,
|
|
},
|
|
}, nil)
|
|
|
|
// Setup target for collection to be loaded
|
|
channel := &datapb.VchannelInfo{
|
|
CollectionID: 1,
|
|
ChannelName: "test-channel",
|
|
}
|
|
segments := []*datapb.SegmentInfo{
|
|
{
|
|
ID: 100,
|
|
CollectionID: 1,
|
|
PartitionID: 1,
|
|
InsertChannel: "test-channel",
|
|
},
|
|
}
|
|
suite.broker.EXPECT().GetRecoveryInfoV2(mock.Anything, int64(1)).Return([]*datapb.VchannelInfo{channel}, segments, nil)
|
|
suite.target.UpdateCollectionNextTarget(suite.ctx, 1)
|
|
|
|
// Expect load to be called on latest serviceable delegator (node 1)
|
|
suite.cluster.EXPECT().LoadSegments(mock.Anything, int64(1), mock.Anything).Return(&commonpb.Status{
|
|
ErrorCode: commonpb.ErrorCode_Success,
|
|
}, nil).Once()
|
|
|
|
// Execute load segment (step 0)
|
|
suite.executor.loadSegment(task, 0)
|
|
|
|
// Setup new delegator (node 2) - serviceable
|
|
newDelegatorView := utils.CreateTestLeaderView(2, 1, "test-channel", map[int64]int64{100: 2}, map[int64]*meta.Segment{})
|
|
newDelegatorView.Version = 2
|
|
newDelegatorView.UnServiceableError = nil // serviceable
|
|
suite.dist.LeaderViewManager.Update(2, newDelegatorView)
|
|
|
|
// Execute release segment (step 1)
|
|
suite.executor.releaseSegment(task, 1)
|
|
|
|
// verify that the task is failed due to shard leader change
|
|
suite.Error(task.Err())
|
|
}
|
|
|
|
func (suite *ExecutorTestSuite) TestBalanceTaskFallbackToLatestWhenNoServiceableDelegator() {
|
|
// Setup collection and replica
|
|
collection := utils.CreateTestCollection(1, 1)
|
|
suite.meta.CollectionManager.PutCollection(suite.ctx, collection)
|
|
suite.meta.CollectionManager.PutPartition(suite.ctx, utils.CreateTestPartition(1, 1))
|
|
|
|
replica := utils.CreateTestReplica(1, 1, []int64{1, 2})
|
|
suite.meta.ReplicaManager.Put(suite.ctx, replica)
|
|
|
|
// Setup nodes
|
|
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 1,
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 2,
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 1)
|
|
suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 2)
|
|
|
|
// Create load action
|
|
loadAction := NewSegmentAction(2, ActionTypeGrow, "test-channel", 100)
|
|
task, err := NewSegmentTask(suite.ctx, time.Second*10, WrapIDSource(0), 1, replica, loadAction)
|
|
suite.NoError(err)
|
|
|
|
// Setup old delegator (node 1) - not serviceable
|
|
oldDelegatorView := utils.CreateTestLeaderView(1, 1, "test-channel", map[int64]int64{100: 1}, map[int64]*meta.Segment{})
|
|
oldDelegatorView.UnServiceableError = errors.New("not serviceable")
|
|
oldDelegatorView.Version = 1
|
|
suite.dist.LeaderViewManager.Update(1, oldDelegatorView)
|
|
|
|
// Setup new delegator (node 2) - not serviceable but latest
|
|
newDelegatorView := utils.CreateTestLeaderView(2, 1, "test-channel", map[int64]int64{100: 2}, map[int64]*meta.Segment{})
|
|
newDelegatorView.UnServiceableError = errors.New("not serviceable")
|
|
newDelegatorView.Version = 2 // latest version
|
|
suite.dist.LeaderViewManager.Update(2, newDelegatorView)
|
|
|
|
// Mock broker responses
|
|
suite.broker.EXPECT().DescribeCollection(mock.Anything, int64(1)).Return(&milvuspb.DescribeCollectionResponse{
|
|
Schema: &schemapb.CollectionSchema{
|
|
Name: "TestBalanceTask",
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
|
|
},
|
|
},
|
|
}, nil)
|
|
suite.broker.EXPECT().GetSegmentInfo(mock.Anything, int64(100)).Return([]*datapb.SegmentInfo{
|
|
{
|
|
ID: 100,
|
|
CollectionID: 1,
|
|
PartitionID: 1,
|
|
InsertChannel: "test-channel",
|
|
},
|
|
}, nil)
|
|
suite.broker.EXPECT().GetIndexInfo(mock.Anything, int64(1), int64(100)).Return(nil, nil)
|
|
suite.broker.EXPECT().ListIndexes(mock.Anything, int64(1)).Return([]*indexpb.IndexInfo{
|
|
{
|
|
CollectionID: 1,
|
|
},
|
|
}, nil)
|
|
|
|
// Setup target for collection to be loaded
|
|
channel := &datapb.VchannelInfo{
|
|
CollectionID: 1,
|
|
ChannelName: "test-channel",
|
|
}
|
|
segments := []*datapb.SegmentInfo{
|
|
{
|
|
ID: 100,
|
|
CollectionID: 1,
|
|
PartitionID: 1,
|
|
InsertChannel: "test-channel",
|
|
},
|
|
}
|
|
suite.broker.EXPECT().GetRecoveryInfoV2(mock.Anything, int64(1)).Return([]*datapb.VchannelInfo{channel}, segments, nil)
|
|
suite.target.UpdateCollectionNextTarget(suite.ctx, 1)
|
|
|
|
// Expect load to be called on the latest delegator (node 2) as fallback
|
|
suite.cluster.EXPECT().LoadSegments(mock.Anything, int64(2), mock.Anything).Return(&commonpb.Status{
|
|
ErrorCode: commonpb.ErrorCode_Success,
|
|
}, nil).Once()
|
|
|
|
// Execute load segment
|
|
suite.executor.loadSegment(task, 0)
|
|
|
|
// Verify that the latest delegator was chosen as fallback
|
|
suite.cluster.AssertExpectations(suite.T())
|
|
}
|
|
|
|
func TestExecutorSuite(t *testing.T) {
|
|
suite.Run(t, new(ExecutorTestSuite))
|
|
}
|
|
|
|
// L0 segment specific tests for delegator selection during load
|
|
func (suite *ExecutorTestSuite) TestLoadL0_SelectDelegatorMissingSegment() {
|
|
t := suite.T()
|
|
mockey.PatchConvey("TestLoadL0_SelectDelegatorMissingSegment", t, func() {
|
|
// Setup collection and replica
|
|
collectionID := int64(1)
|
|
segmentID := int64(100)
|
|
channelName := "test-channel"
|
|
|
|
collection := utils.CreateTestCollection(collectionID, 1)
|
|
suite.meta.CollectionManager.PutCollection(suite.ctx, collection)
|
|
suite.meta.CollectionManager.PutPartition(suite.ctx, utils.CreateTestPartition(collectionID, 1))
|
|
|
|
replica := utils.CreateTestReplica(1, collectionID, []int64{1, 2})
|
|
suite.meta.ReplicaManager.Put(suite.ctx, replica)
|
|
|
|
// Setup nodes
|
|
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 1, Address: "localhost", Hostname: "localhost"}))
|
|
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 2, Address: "localhost", Hostname: "localhost"}))
|
|
suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 1)
|
|
suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 2)
|
|
|
|
// Only one delegator view which is missing the L0 segment
|
|
viewMissing := utils.CreateTestLeaderView(2, collectionID, channelName, map[int64]int64{}, map[int64]*meta.Segment{})
|
|
viewMissing.UnServiceableError = nil
|
|
viewMissing.Version = 1
|
|
suite.dist.LeaderViewManager.Update(2, viewMissing)
|
|
|
|
// Mock broker responses using mockey
|
|
mockey.Mock((*meta.MockBroker).DescribeCollection).Return(&milvuspb.DescribeCollectionResponse{Schema: utils.CreateTestSchema()}, nil).Build()
|
|
mockey.Mock((*meta.MockBroker).GetSegmentInfo).Return([]*datapb.SegmentInfo{{
|
|
ID: segmentID,
|
|
CollectionID: collectionID,
|
|
PartitionID: 1,
|
|
InsertChannel: channelName,
|
|
Level: datapb.SegmentLevel_L0,
|
|
}}, nil).Build()
|
|
mockey.Mock((*meta.MockBroker).GetIndexInfo).Return(nil, nil).Build()
|
|
mockey.Mock((*meta.MockBroker).ListIndexes).Return([]*indexpb.IndexInfo{{CollectionID: collectionID}}, nil).Build()
|
|
|
|
channel := &datapb.VchannelInfo{CollectionID: collectionID, ChannelName: channelName}
|
|
segments := []*datapb.SegmentInfo{{ID: segmentID, CollectionID: collectionID, PartitionID: 1, InsertChannel: channelName}}
|
|
mockey.Mock((*meta.MockBroker).GetRecoveryInfoV2).Return([]*datapb.VchannelInfo{channel}, segments, nil).Build()
|
|
suite.target.UpdateCollectionNextTarget(suite.ctx, collectionID)
|
|
|
|
// Capture LoadSegments target node
|
|
var capturedNode int64 = -1
|
|
mockey.Mock((*session.MockCluster).LoadSegments).To(func(m *session.MockCluster, ctx context.Context, node int64, req *querypb.LoadSegmentsRequest) (*commonpb.Status, error) {
|
|
capturedNode = node
|
|
return &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil
|
|
}).Build()
|
|
|
|
// Execute load segment
|
|
action := NewSegmentAction(1, ActionTypeGrow, channelName, segmentID)
|
|
task, err := NewSegmentTask(suite.ctx, time.Second*10, WrapIDSource(0), collectionID, replica, action)
|
|
suite.NoError(err)
|
|
_ = suite.executor.loadSegment(task, 0)
|
|
|
|
suite.Equal(int64(2), capturedNode)
|
|
})
|
|
}
|
|
|
|
func (suite *ExecutorTestSuite) TestLoadL0_FallbackWhenAllDelegatorsHaveSegment() {
|
|
t := suite.T()
|
|
mockey.PatchConvey("TestLoadL0_FallbackWhenAllDelegatorsHaveSegment", t, func() {
|
|
// Setup collection and replica
|
|
collectionID := int64(1)
|
|
segmentID := int64(101)
|
|
channelName := "test-channel"
|
|
|
|
collection := utils.CreateTestCollection(collectionID, 1)
|
|
suite.meta.CollectionManager.PutCollection(suite.ctx, collection)
|
|
suite.meta.CollectionManager.PutPartition(suite.ctx, utils.CreateTestPartition(collectionID, 1))
|
|
|
|
replica := utils.CreateTestReplica(1, collectionID, []int64{1, 2})
|
|
suite.meta.ReplicaManager.Put(suite.ctx, replica)
|
|
|
|
// Setup nodes
|
|
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 1, Address: "localhost", Hostname: "localhost"}))
|
|
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 2, Address: "localhost", Hostname: "localhost"}))
|
|
suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 1)
|
|
suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 2)
|
|
|
|
// Both delegators already have the L0 segment; should fallback to serviceable latest
|
|
view1 := utils.CreateTestLeaderView(1, collectionID, channelName, map[int64]int64{segmentID: 1}, map[int64]*meta.Segment{})
|
|
view1.UnServiceableError = nil
|
|
view1.Version = 2
|
|
suite.dist.LeaderViewManager.Update(1, view1)
|
|
|
|
view2 := utils.CreateTestLeaderView(2, collectionID, channelName, map[int64]int64{segmentID: 2}, map[int64]*meta.Segment{})
|
|
view2.UnServiceableError = errors.New("not serviceable")
|
|
view2.Version = 3
|
|
suite.dist.LeaderViewManager.Update(2, view2)
|
|
|
|
// Mock broker responses using mockey
|
|
mockey.Mock((*meta.MockBroker).DescribeCollection).Return(&milvuspb.DescribeCollectionResponse{Schema: utils.CreateTestSchema()}, nil).Build()
|
|
mockey.Mock((*meta.MockBroker).GetSegmentInfo).Return([]*datapb.SegmentInfo{{
|
|
ID: segmentID,
|
|
CollectionID: collectionID,
|
|
PartitionID: 1,
|
|
InsertChannel: channelName,
|
|
Level: datapb.SegmentLevel_L0,
|
|
}}, nil).Build()
|
|
mockey.Mock((*meta.MockBroker).GetIndexInfo).Return(nil, nil).Build()
|
|
mockey.Mock((*meta.MockBroker).ListIndexes).Return([]*indexpb.IndexInfo{{CollectionID: collectionID}}, nil).Build()
|
|
|
|
channel := &datapb.VchannelInfo{CollectionID: collectionID, ChannelName: channelName}
|
|
segments := []*datapb.SegmentInfo{{ID: segmentID, CollectionID: collectionID, PartitionID: 1, InsertChannel: channelName}}
|
|
mockey.Mock((*meta.MockBroker).GetRecoveryInfoV2).Return([]*datapb.VchannelInfo{channel}, segments, nil).Build()
|
|
suite.target.UpdateCollectionNextTarget(suite.ctx, collectionID)
|
|
|
|
// Capture LoadSegments target node
|
|
var capturedNode int64 = -1
|
|
mockey.Mock((*session.MockCluster).LoadSegments).To(func(m *session.MockCluster, ctx context.Context, node int64, req *querypb.LoadSegmentsRequest) (*commonpb.Status, error) {
|
|
capturedNode = node
|
|
return &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil
|
|
}).Build()
|
|
|
|
action := NewSegmentAction(1, ActionTypeGrow, channelName, segmentID)
|
|
task, err := NewSegmentTask(suite.ctx, time.Second*10, WrapIDSource(0), collectionID, replica, action)
|
|
suite.NoError(err)
|
|
_ = suite.executor.loadSegment(task, 0)
|
|
|
|
suite.Equal(int64(1), capturedNode)
|
|
})
|
|
}
|
|
|
|
func (suite *ExecutorTestSuite) TestLoadL0_NoDelegatorFound() {
|
|
t := suite.T()
|
|
mockey.PatchConvey("TestLoadL0_NoDelegatorFound", t, func() {
|
|
// Setup collection and replica
|
|
collectionID := int64(1)
|
|
segmentID := int64(102)
|
|
channelName := "test-channel"
|
|
|
|
collection := utils.CreateTestCollection(collectionID, 1)
|
|
suite.meta.CollectionManager.PutCollection(suite.ctx, collection)
|
|
suite.meta.CollectionManager.PutPartition(suite.ctx, utils.CreateTestPartition(collectionID, 1))
|
|
|
|
replica := utils.CreateTestReplica(1, collectionID, []int64{1})
|
|
suite.meta.ReplicaManager.Put(suite.ctx, replica)
|
|
|
|
// Setup node
|
|
suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 1, Address: "localhost", Hostname: "localhost"}))
|
|
suite.meta.ResourceManager.HandleNodeUp(suite.ctx, 1)
|
|
|
|
// No leader views set for the channel -> should return error early
|
|
|
|
// Mock broker responses using mockey
|
|
mockey.Mock((*meta.MockBroker).DescribeCollection).Return(&milvuspb.DescribeCollectionResponse{Schema: utils.CreateTestSchema()}, nil).Build()
|
|
mockey.Mock((*meta.MockBroker).GetSegmentInfo).Return([]*datapb.SegmentInfo{{
|
|
ID: segmentID,
|
|
CollectionID: collectionID,
|
|
PartitionID: 1,
|
|
InsertChannel: channelName,
|
|
Level: datapb.SegmentLevel_L0,
|
|
}}, nil).Build()
|
|
mockey.Mock((*meta.MockBroker).GetIndexInfo).Return(nil, nil).Build()
|
|
mockey.Mock((*meta.MockBroker).ListIndexes).Return([]*indexpb.IndexInfo{{CollectionID: collectionID}}, nil).Build()
|
|
|
|
channel := &datapb.VchannelInfo{CollectionID: collectionID, ChannelName: channelName}
|
|
segments := []*datapb.SegmentInfo{{ID: segmentID, CollectionID: collectionID, PartitionID: 1, InsertChannel: channelName}}
|
|
mockey.Mock((*meta.MockBroker).GetRecoveryInfoV2).Return([]*datapb.VchannelInfo{channel}, segments, nil).Build()
|
|
suite.target.UpdateCollectionNextTarget(suite.ctx, collectionID)
|
|
|
|
action := NewSegmentAction(1, ActionTypeGrow, channelName, segmentID)
|
|
task, err := NewSegmentTask(suite.ctx, time.Second*10, WrapIDSource(0), collectionID, replica, action)
|
|
suite.NoError(err)
|
|
err = suite.executor.loadSegment(task, 0)
|
|
suite.Error(err)
|
|
})
|
|
}
|