fix: check collection health(queryable) fail for releasing collection (…

…milvus-io#34947) issue: milvus-io#34946 Signed-off-by: jaime <[email protected]>
bigsheeper · Aug 2, 2024 · fcec4c2 · fcec4c2
1 parent 3641ae6
commit fcec4c2
Show file tree

Hide file tree

Showing 20 changed files with 173 additions and 75 deletions.
diff --git a/configs/milvus.yaml b/configs/milvus.yaml
@@ -367,6 +367,7 @@ queryCoord:
   channelExclusiveNodeFactor: 4 # the least node number for enable channel's exclusive mode
   collectionObserverInterval: 200 # the interval of collection observer
   checkExecutedFlagInterval: 100 # the interval of check executed flag to force to pull dist
+  updateCollectionLoadStatusInterval: 5 # 5m, max interval of updating collection loaded status for check health
   cleanExcludeSegmentInterval: 60 # the time duration of clean pipeline exclude segment which used for filter invalid data, in seconds
   ip:  # TCP/IP address of queryCoord. If not specified, use the first unicastable address
   port: 19531 # TCP port of queryCoord

diff --git a/internal/querycoordv2/balance/channel_level_score_balancer.go b/internal/querycoordv2/balance/channel_level_score_balancer.go
@@ -41,7 +41,7 @@ func NewChannelLevelScoreBalancer(scheduler task.Scheduler,
 	nodeManager *session.NodeManager,
 	dist *meta.DistributionManager,
 	meta *meta.Meta,
-	targetMgr *meta.TargetManager,
+	targetMgr meta.TargetManagerInterface,
 ) *ChannelLevelScoreBalancer {
 	return &ChannelLevelScoreBalancer{
 		ScoreBasedBalancer: NewScoreBasedBalancer(scheduler, nodeManager, dist, meta, targetMgr),

diff --git a/internal/querycoordv2/balance/multi_target_balance.go b/internal/querycoordv2/balance/multi_target_balance.go
@@ -452,7 +452,7 @@ func (g *randomPlanGenerator) generatePlans() []SegmentAssignPlan {
 type MultiTargetBalancer struct {
 	*ScoreBasedBalancer
 	dist      *meta.DistributionManager
-	targetMgr *meta.TargetManager
+	targetMgr meta.TargetManagerInterface
 }
 
 func (b *MultiTargetBalancer) BalanceReplica(replica *meta.Replica) ([]SegmentAssignPlan, []ChannelAssignPlan) {
@@ -548,7 +548,7 @@ func (b *MultiTargetBalancer) genPlanByDistributions(nodeSegments, globalNodeSeg
 	return plans
 }
 
-func NewMultiTargetBalancer(scheduler task.Scheduler, nodeManager *session.NodeManager, dist *meta.DistributionManager, meta *meta.Meta, targetMgr *meta.TargetManager) *MultiTargetBalancer {
+func NewMultiTargetBalancer(scheduler task.Scheduler, nodeManager *session.NodeManager, dist *meta.DistributionManager, meta *meta.Meta, targetMgr meta.TargetManagerInterface) *MultiTargetBalancer {
 	return &MultiTargetBalancer{
 		ScoreBasedBalancer: NewScoreBasedBalancer(scheduler, nodeManager, dist, meta, targetMgr),
 		dist:               dist,

diff --git a/internal/querycoordv2/balance/rowcount_based_balancer.go b/internal/querycoordv2/balance/rowcount_based_balancer.go
@@ -36,7 +36,7 @@ type RowCountBasedBalancer struct {
 	*RoundRobinBalancer
 	dist      *meta.DistributionManager
 	meta      *meta.Meta
-	targetMgr *meta.TargetManager
+	targetMgr meta.TargetManagerInterface
 }
 
 // AssignSegment, when row count based balancer assign segments, it will assign segment to node with least global row count.
@@ -354,7 +354,7 @@ func NewRowCountBasedBalancer(
 	nodeManager *session.NodeManager,
 	dist *meta.DistributionManager,
 	meta *meta.Meta,
-	targetMgr *meta.TargetManager,
+	targetMgr meta.TargetManagerInterface,
 ) *RowCountBasedBalancer {
 	return &RowCountBasedBalancer{
 		RoundRobinBalancer: NewRoundRobinBalancer(scheduler, nodeManager),

diff --git a/internal/querycoordv2/balance/score_based_balancer.go b/internal/querycoordv2/balance/score_based_balancer.go
@@ -41,7 +41,7 @@ func NewScoreBasedBalancer(scheduler task.Scheduler,
 	nodeManager *session.NodeManager,
 	dist *meta.DistributionManager,
 	meta *meta.Meta,
-	targetMgr *meta.TargetManager,
+	targetMgr meta.TargetManagerInterface,
 ) *ScoreBasedBalancer {
 	return &ScoreBasedBalancer{
 		RowCountBasedBalancer: NewRowCountBasedBalancer(scheduler, nodeManager, dist, meta, targetMgr),

diff --git a/internal/querycoordv2/dist/dist_controller.go b/internal/querycoordv2/dist/dist_controller.go
@@ -41,7 +41,7 @@ type ControllerImpl struct {
 	client      session.Cluster
 	nodeManager *session.NodeManager
 	dist        *meta.DistributionManager
-	targetMgr   *meta.TargetManager
+	targetMgr   meta.TargetManagerInterface
 	scheduler   task.Scheduler
 }
 
@@ -98,7 +98,7 @@ func NewDistController(
 	client session.Cluster,
 	nodeManager *session.NodeManager,
 	dist *meta.DistributionManager,
-	targetMgr *meta.TargetManager,
+	targetMgr meta.TargetManagerInterface,
 	scheduler task.Scheduler,
 ) *ControllerImpl {
 	return &ControllerImpl{

diff --git a/internal/querycoordv2/job/job_load.go b/internal/querycoordv2/job/job_load.go
@@ -48,7 +48,7 @@ type LoadCollectionJob struct {
 	meta               *meta.Meta
 	broker             meta.Broker
 	cluster            session.Cluster
-	targetMgr          *meta.TargetManager
+	targetMgr          meta.TargetManagerInterface
 	targetObserver     *observers.TargetObserver
 	collectionObserver *observers.CollectionObserver
 	nodeMgr            *session.NodeManager
@@ -61,7 +61,7 @@ func NewLoadCollectionJob(
 	meta *meta.Meta,
 	broker meta.Broker,
 	cluster session.Cluster,
-	targetMgr *meta.TargetManager,
+	targetMgr meta.TargetManagerInterface,
 	targetObserver *observers.TargetObserver,
 	collectionObserver *observers.CollectionObserver,
 	nodeMgr *session.NodeManager,
@@ -239,7 +239,7 @@ type LoadPartitionJob struct {
 	meta               *meta.Meta
 	broker             meta.Broker
 	cluster            session.Cluster
-	targetMgr          *meta.TargetManager
+	targetMgr          meta.TargetManagerInterface
 	targetObserver     *observers.TargetObserver
 	collectionObserver *observers.CollectionObserver
 	nodeMgr            *session.NodeManager
@@ -252,7 +252,7 @@ func NewLoadPartitionJob(
 	meta *meta.Meta,
 	broker meta.Broker,
 	cluster session.Cluster,
-	targetMgr *meta.TargetManager,
+	targetMgr meta.TargetManagerInterface,
 	targetObserver *observers.TargetObserver,
 	collectionObserver *observers.CollectionObserver,
 	nodeMgr *session.NodeManager,

diff --git a/internal/querycoordv2/job/job_release.go b/internal/querycoordv2/job/job_release.go
@@ -39,7 +39,7 @@ type ReleaseCollectionJob struct {
 	meta              *meta.Meta
 	broker            meta.Broker
 	cluster           session.Cluster
-	targetMgr         *meta.TargetManager
+	targetMgr         meta.TargetManagerInterface
 	targetObserver    *observers.TargetObserver
 	checkerController *checkers.CheckerController
 }
@@ -50,7 +50,7 @@ func NewReleaseCollectionJob(ctx context.Context,
 	meta *meta.Meta,
 	broker meta.Broker,
 	cluster session.Cluster,
-	targetMgr *meta.TargetManager,
+	targetMgr meta.TargetManagerInterface,
 	targetObserver *observers.TargetObserver,
 	checkerController *checkers.CheckerController,
 ) *ReleaseCollectionJob {
@@ -114,7 +114,7 @@ type ReleasePartitionJob struct {
 	meta              *meta.Meta
 	broker            meta.Broker
 	cluster           session.Cluster
-	targetMgr         *meta.TargetManager
+	targetMgr         meta.TargetManagerInterface
 	targetObserver    *observers.TargetObserver
 	checkerController *checkers.CheckerController
 }
@@ -125,7 +125,7 @@ func NewReleasePartitionJob(ctx context.Context,
 	meta *meta.Meta,
 	broker meta.Broker,
 	cluster session.Cluster,
-	targetMgr *meta.TargetManager,
+	targetMgr meta.TargetManagerInterface,
 	targetObserver *observers.TargetObserver,
 	checkerController *checkers.CheckerController,
 ) *ReleasePartitionJob {

diff --git a/internal/querycoordv2/job/undo.go b/internal/querycoordv2/job/undo.go
@@ -38,12 +38,12 @@ type UndoList struct {
 	ctx            context.Context
 	meta           *meta.Meta
 	cluster        session.Cluster
-	targetMgr      *meta.TargetManager
+	targetMgr      meta.TargetManagerInterface
 	targetObserver *observers.TargetObserver
 }
 
 func NewUndoList(ctx context.Context, meta *meta.Meta,
-	cluster session.Cluster, targetMgr *meta.TargetManager, targetObserver *observers.TargetObserver,
+	cluster session.Cluster, targetMgr meta.TargetManagerInterface, targetObserver *observers.TargetObserver,
 ) *UndoList {
 	return &UndoList{
 		ctx:            ctx,

diff --git a/internal/querycoordv2/meta/mock_target_manager.go b/internal/querycoordv2/meta/mock_target_manager.go
diff --git a/internal/querycoordv2/meta/target_manager.go b/internal/querycoordv2/meta/target_manager.go
@@ -71,6 +71,7 @@ type TargetManagerInterface interface {
 	IsNextTargetExist(collectionID int64) bool
 	SaveCurrentTarget(catalog metastore.QueryCoordCatalog)
 	Recover(catalog metastore.QueryCoordCatalog) error
+	CanSegmentBeMoved(collectionID, segmentID int64) bool
 }
 
 type TargetManager struct {

diff --git a/internal/querycoordv2/observers/collection_observer.go b/internal/querycoordv2/observers/collection_observer.go
@@ -43,7 +43,7 @@ type CollectionObserver struct {
 
 	dist                 *meta.DistributionManager
 	meta                 *meta.Meta
-	targetMgr            *meta.TargetManager
+	targetMgr            meta.TargetManagerInterface
 	targetObserver       *TargetObserver
 	checkerController    *checkers.CheckerController
 	partitionLoadedCount map[int64]int
@@ -62,7 +62,7 @@ type LoadTask struct {
 func NewCollectionObserver(
 	dist *meta.DistributionManager,
 	meta *meta.Meta,
-	targetMgr *meta.TargetManager,
+	targetMgr meta.TargetManagerInterface,
 	targetObserver *TargetObserver,
 	checherController *checkers.CheckerController,
 ) *CollectionObserver {

diff --git a/internal/querycoordv2/observers/target_observer.go b/internal/querycoordv2/observers/target_observer.go
@@ -55,7 +55,7 @@ type TargetObserver struct {
 	cancel    context.CancelFunc
 	wg        sync.WaitGroup
 	meta      *meta.Meta
-	targetMgr *meta.TargetManager
+	targetMgr meta.TargetManagerInterface
 	distMgr   *meta.DistributionManager
 	broker    meta.Broker
 	cluster   session.Cluster
@@ -76,7 +76,7 @@ type TargetObserver struct {
 
 func NewTargetObserver(
 	meta *meta.Meta,
-	targetMgr *meta.TargetManager,
+	targetMgr meta.TargetManagerInterface,
 	distMgr *meta.DistributionManager,
 	broker meta.Broker,
 	cluster session.Cluster,

diff --git a/internal/querycoordv2/server.go b/internal/querycoordv2/server.go
@@ -90,7 +90,7 @@ type Server struct {
 	store     metastore.QueryCoordCatalog
 	meta      *meta.Meta
 	dist      *meta.DistributionManager
-	targetMgr *meta.TargetManager
+	targetMgr meta.TargetManagerInterface
 	broker    meta.Broker
 
 	// Session

diff --git a/internal/querycoordv2/services_test.go b/internal/querycoordv2/services_test.go
@@ -1610,44 +1610,65 @@ func (suite *ServiceSuite) TestGetReplicasWhenNoAvailableNodes() {
 }
 
 func (suite *ServiceSuite) TestCheckHealth() {
+	suite.loadAll()
 	ctx := context.Background()
 	server := suite.server
 
-	// Test for server is not healthy
-	server.UpdateStateCode(commonpb.StateCode_Initializing)
-	resp, err := server.CheckHealth(ctx, &milvuspb.CheckHealthRequest{})
-	suite.NoError(err)
-	suite.Equal(resp.IsHealthy, false)
-	suite.NotEmpty(resp.Reasons)
+	assertCheckHealthResult := func(isHealthy bool) {
+		resp, err := server.CheckHealth(ctx, &milvuspb.CheckHealthRequest{})
+		suite.NoError(err)
+		suite.Equal(resp.IsHealthy, isHealthy)
+		if !isHealthy {
+			suite.NotEmpty(resp.Reasons)
+		} else {
+			suite.Empty(resp.Reasons)
+		}
+	}
 
-	// Test for components state fail
-	for _, node := range suite.nodes {
-		suite.cluster.EXPECT().GetComponentStates(mock.Anything, node).Return(
+	setNodeSate := func(state commonpb.StateCode) {
+		// Test for components state fail
+		suite.cluster.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Unset()
+		suite.cluster.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(
 			&milvuspb.ComponentStates{
-				State:  &milvuspb.ComponentInfo{StateCode: commonpb.StateCode_Abnormal},
+				State:  &milvuspb.ComponentInfo{StateCode: state},
 				Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
 			},
-			nil).Once()
+			nil).Maybe()
 	}
+
+	// Test for server is not healthy
+	server.UpdateStateCode(commonpb.StateCode_Initializing)
+	assertCheckHealthResult(false)
+
+	// Test for components state fail
+	setNodeSate(commonpb.StateCode_Abnormal)
 	server.UpdateStateCode(commonpb.StateCode_Healthy)
-	resp, err = server.CheckHealth(ctx, &milvuspb.CheckHealthRequest{})
-	suite.NoError(err)
-	suite.Equal(resp.IsHealthy, false)
-	suite.NotEmpty(resp.Reasons)
+	assertCheckHealthResult(false)
 
-	// Test for server is healthy
+	// Test for check load percentage fail
+	setNodeSate(commonpb.StateCode_Healthy)
+	assertCheckHealthResult(true)
+
+	// Test for check channel ok
+	for _, collection := range suite.collections {
+		suite.updateCollectionStatus(collection, querypb.LoadStatus_Loaded)
+		suite.updateChannelDist(collection)
+	}
+	assertCheckHealthResult(true)
+
+	// Test for check channel fail
+	tm := meta.NewMockTargetManager(suite.T())
+	tm.EXPECT().GetDmChannelsByCollection(mock.Anything, mock.Anything).Return(nil).Maybe()
+	otm := server.targetMgr
+	server.targetMgr = tm
+	assertCheckHealthResult(true)
+
+	// Test for get shard leader fail
+	server.targetMgr = otm
 	for _, node := range suite.nodes {
-		suite.cluster.EXPECT().GetComponentStates(mock.Anything, node).Return(
-			&milvuspb.ComponentStates{
-				State:  &milvuspb.ComponentInfo{StateCode: commonpb.StateCode_Healthy},
-				Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
-			},
-			nil).Once()
+		suite.nodeMgr.Suspend(node)
 	}
-	resp, err = server.CheckHealth(ctx, &milvuspb.CheckHealthRequest{})
-	suite.NoError(err)
-	suite.Equal(resp.IsHealthy, true)
-	suite.Empty(resp.Reasons)
+	assertCheckHealthResult(true)
 }
 
 func (suite *ServiceSuite) TestGetShardLeaders() {

diff --git a/internal/querycoordv2/task/executor.go b/internal/querycoordv2/task/executor.go
@@ -57,7 +57,7 @@ type Executor struct {
 	meta      *meta.Meta
 	dist      *meta.DistributionManager
 	broker    meta.Broker
-	targetMgr *meta.TargetManager
+	targetMgr meta.TargetManagerInterface
 	cluster   session.Cluster
 	nodeMgr   *session.NodeManager
 
@@ -69,7 +69,7 @@ type Executor struct {
 func NewExecutor(meta *meta.Meta,
 	dist *meta.DistributionManager,
 	broker meta.Broker,
-	targetMgr *meta.TargetManager,
+	targetMgr meta.TargetManagerInterface,
 	cluster session.Cluster,
 	nodeMgr *session.NodeManager,
 ) *Executor {

diff --git a/internal/querycoordv2/task/scheduler.go b/internal/querycoordv2/task/scheduler.go
@@ -157,7 +157,7 @@ type taskScheduler struct {
 
 	distMgr   *meta.DistributionManager
 	meta      *meta.Meta
-	targetMgr *meta.TargetManager
+	targetMgr meta.TargetManagerInterface
 	broker    meta.Broker
 	cluster   session.Cluster
 	nodeMgr   *session.NodeManager
@@ -177,7 +177,7 @@ type taskScheduler struct {
 func NewScheduler(ctx context.Context,
 	meta *meta.Meta,
 	distMgr *meta.DistributionManager,
-	targetMgr *meta.TargetManager,
+	targetMgr meta.TargetManagerInterface,
 	broker meta.Broker,
 	cluster session.Cluster,
 	nodeMgr *session.NodeManager,