When the task is canceled, reduce the load of node (#17735)

Signed-off-by: Cai.Zhang <cai.zhang@zilliz.com>

internal/indexcoord/index_coord.go

@@ -546,7 +546,12 @@ func (i *IndexCoord) DropIndex(ctx context.Context, req *indexpb.DropIndexReques
 	ret := &commonpb.Status{
 		ErrorCode: commonpb.ErrorCode_Success,
 	}
-	err := i.metaTable.MarkIndexAsDeleted(req.IndexID)
+	nodeTasks, err := i.metaTable.MarkIndexAsDeleted(req.IndexID)
+	defer func() {
+		for nodeID, taskNum := range nodeTasks {
+			i.nodeManager.pq.IncPriority(nodeID, taskNum*-1)
+		}
+	}()
 	if err != nil {
 		ret.ErrorCode = commonpb.ErrorCode_UnexpectedError
 		ret.Reason = err.Error()
@@ -584,7 +589,13 @@ func (i *IndexCoord) RemoveIndex(ctx context.Context, req *indexpb.RemoveIndexRe
 	ret := &commonpb.Status{
 		ErrorCode: commonpb.ErrorCode_Success,
 	}
-	err := i.metaTable.MarkIndexAsDeletedByBuildIDs(req.GetBuildIDs())
+
+	nodeTasks, err := i.metaTable.MarkIndexAsDeletedByBuildIDs(req.GetBuildIDs())
+	defer func() {
+		for nodeID, taskNum := range nodeTasks {
+			i.nodeManager.pq.IncPriority(nodeID, -1*taskNum)
+		}
+	}()
 	if err != nil {
 		log.Error("IndexCoord MarkIndexAsDeletedByBuildIDs failed", zap.Int64s("buildIDs", req.GetBuildIDs()),
 			zap.Error(err))

internal/indexcoord/meta_table.go

@@ -249,12 +249,13 @@ func (mt *metaTable) UpdateVersion(indexBuildID UniqueID) error {
 }
 
 // MarkIndexAsDeleted will mark the corresponding index as deleted, and recycleUnusedIndexFiles will recycle these tasks.
-func (mt *metaTable) MarkIndexAsDeleted(indexID UniqueID) error {
+func (mt *metaTable) MarkIndexAsDeleted(indexID UniqueID) (map[int64]int, error) {
 	mt.lock.Lock()
 	defer mt.lock.Unlock()
 
 	log.Debug("IndexCoord metaTable MarkIndexAsDeleted ", zap.Int64("indexID", indexID))
 
+	node2TaskNum := make(map[int64]int)
 	for _, meta := range mt.indexBuildID2Meta {
 		if meta.indexMeta.Req.IndexID == indexID && !meta.indexMeta.MarkDeleted {
 			meta.indexMeta.MarkDeleted = true
@@ -273,21 +274,24 @@ func (mt *metaTable) MarkIndexAsDeleted(indexID UniqueID) error {
 				}
 				err2 := retry.Do(context.TODO(), fn, retry.Attempts(5))
 				if err2 != nil {
-					return err2
+					return node2TaskNum, err2
 				}
+			} else {
+				node2TaskNum[meta.indexMeta.NodeID]++
 			}
 		}
 	}
 
-	return nil
+	return node2TaskNum, nil
 }
 
-func (mt *metaTable) MarkIndexAsDeletedByBuildIDs(buildIDs []UniqueID) error {
+func (mt *metaTable) MarkIndexAsDeletedByBuildIDs(buildIDs []UniqueID) (map[int64]int, error) {
 	mt.lock.Lock()
 	defer mt.lock.Unlock()
 
 	log.Debug("IndexCoord metaTable MarkIndexAsDeletedByBuildIDs ", zap.Int64s("buildIDs", buildIDs))
 
+	node2TaskNum := make(map[int64]int)
 	for _, buildID := range buildIDs {
 		if meta, ok := mt.indexBuildID2Meta[buildID]; ok {
 			clonedMeta := &Meta{
@@ -310,12 +314,13 @@ func (mt *metaTable) MarkIndexAsDeletedByBuildIDs(buildIDs []UniqueID) error {
 				}
 				err2 := retry.Do(context.TODO(), fn, retry.Attempts(5))
 				if err2 != nil {
-					return err2
+					return node2TaskNum, err2
 				}
 			}
+			node2TaskNum[meta.indexMeta.NodeID]++
 		}
 	}
-	return nil
+	return node2TaskNum, nil
 }
 
 // GetIndexStates gets the index states from meta table.

internal/indexcoord/meta_table_test.go

@@ -135,7 +135,7 @@ func TestMetaTable(t *testing.T) {
 		key = "indexes/" + strconv.FormatInt(indexMeta1.IndexBuildID, 10)
 		err = etcdKV.Save(key, string(value))
 		assert.Nil(t, err)
-		err = metaTable.MarkIndexAsDeleted(indexMeta1.Req.IndexID)
+		_, err = metaTable.MarkIndexAsDeleted(indexMeta1.Req.IndexID)
 		assert.Nil(t, err)
 	})
 
@@ -153,7 +153,7 @@ func TestMetaTable(t *testing.T) {
 		key = path.Join(indexFilePrefix, strconv.FormatInt(indexMeta1.IndexBuildID, 10))
 		err = etcdKV.Save(key, string(value))
 		assert.Nil(t, err)
-		err = metaTable.MarkIndexAsDeletedByBuildIDs([]UniqueID{indexMeta1.IndexBuildID})
+		_, err = metaTable.MarkIndexAsDeletedByBuildIDs([]UniqueID{indexMeta1.IndexBuildID})
 		assert.Nil(t, err)
 	})