Set triggerCondition when queryCoord reload loadbalanceTask meta (#15380

) Signed-off-by: xige-16 <[email protected]>
milvus-io · Jan 25, 2022 · 6336e23 · 6336e23
1 parent 8ab9d76
commit 6336e23
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 0 deletions.
diff --git a/internal/querycoord/impl.go b/internal/querycoord/impl.go
@@ -847,6 +847,7 @@ func (qc *QueryCoord) LoadBalance(ctx context.Context, req *querypb.LoadBalanceR
 	}
 
 	baseTask := newBaseTask(qc.loopCtx, querypb.TriggerCondition_LoadBalance)
+	req.BalanceReason = querypb.TriggerCondition_LoadBalance
 	loadBalanceTask := &loadBalanceTask{
 		baseTask:           baseTask,
 		LoadBalanceRequest: req,

diff --git a/internal/querycoord/query_coord.go b/internal/querycoord/query_coord.go
@@ -343,6 +343,7 @@ func (qc *QueryCoord) watchNodeLoop() {
 				MsgType:  commonpb.MsgType_LoadBalanceSegments,
 				SourceID: qc.session.ServerID,
 			},
+			BalanceReason: querypb.TriggerCondition_NodeDown,
 			SourceNodeIDs: offlineNodeIDs,
 		}
 

diff --git a/internal/querycoord/task.go b/internal/querycoord/task.go
@@ -72,6 +72,7 @@ type task interface {
 	msgType() commonpb.MsgType
 	timestamp() Timestamp
 	getTriggerCondition() querypb.TriggerCondition
+	setTriggerCondition(trigger querypb.TriggerCondition)
 	preExecute(ctx context.Context) error
 	execute(ctx context.Context) error
 	postExecute(ctx context.Context) error
@@ -109,6 +110,7 @@ type baseTask struct {
 
 	taskID           UniqueID
 	triggerCondition querypb.TriggerCondition
+	triggerMu        sync.RWMutex
 	parentTask       task
 	childTasks       []task
 	childTasksMu     sync.RWMutex
@@ -146,9 +148,19 @@ func (bt *baseTask) traceCtx() context.Context {
 }
 
 func (bt *baseTask) getTriggerCondition() querypb.TriggerCondition {
+	bt.triggerMu.RLock()
+	defer bt.triggerMu.RUnlock()
+
 	return bt.triggerCondition
 }
 
+func (bt *baseTask) setTriggerCondition(trigger querypb.TriggerCondition) {
+	bt.triggerMu.Lock()
+	defer bt.triggerMu.Unlock()
+
+	bt.triggerCondition = trigger
+}
+
 func (bt *baseTask) taskPriority() querypb.TriggerCondition {
 	return bt.triggerCondition
 }

diff --git a/internal/querycoord/task_scheduler.go b/internal/querycoord/task_scheduler.go
@@ -403,6 +403,9 @@ func (scheduler *TaskScheduler) unmarshalTask(taskID UniqueID, t string) (task,
 		if err != nil {
 			return nil, err
 		}
+		// if triggerCondition == nodeDown, and the queryNode resources are insufficient,
+		// queryCoord will waits until queryNode can load the data, ensuring that the data is not lost
+		baseTask.setTriggerCondition(loadReq.BalanceReason)
 		loadBalanceTask := &loadBalanceTask{
 			baseTask:           baseTask,
 			LoadBalanceRequest: &loadReq,