Repository: milvus-io/milvus
File: internal/querycoordv2/task/scheduler.go (Go)

Code before:
func (scheduler *taskScheduler) RemoveByNode(node int64) {
	scheduler.rwmutex.Lock()
	defer scheduler.rwmutex.Unlock()

	for _, task := range scheduler.segmentTasks {
		if scheduler.isRelated(task, node) {
			scheduler.remove(task)
		}
	}
	for _, task := range scheduler.channelTasks {
		if scheduler.isRelated(task, node) {
			scheduler.remove(task)
		}
	}
}

func (scheduler *taskScheduler) remove(task Task) {
	log := log.With(
		zap.Int64("taskID", task.ID()),
		zap.Int64("collectionID", task.CollectionID()),
		zap.Int64("replicaID", task.ReplicaID()),
		zap.Int32("taskStatus", task.Status()),
	)
	task.Cancel()
	scheduler.tasks.Remove(task.ID())
	scheduler.waitQueue.Remove(task)
	scheduler.processQueue.Remove(task)

	switch task := task.(type) {
	case *SegmentTask:
		index := NewReplicaSegmentIndex(task)
		delete(scheduler.segmentTasks, index)
		log = log.With(zap.Int64("segmentID", task.SegmentID()))

	case *ChannelTask:
		index := replicaChannelIndex{task.ReplicaID(), task.Channel()}
		delete(scheduler.channelTasks, index)
		log = log.With(zap.String("channel", task.Channel()))
	}

	metrics.QueryCoordTaskNum.WithLabelValues().Set(float64(scheduler.tasks.Len()))
	log.Debug("task removed")
}

func (scheduler *taskScheduler) checkCanceled(task Task) bool {
	log := log.With(
		zap.Int64("taskID", task.ID()),
		zap.Int64("collectionID", task.CollectionID()),
		zap.Int64("replicaID", task.ReplicaID()),
		zap.Int64("source", task.SourceID()),
	)

Review comment (suggestion):
We would better to make this work for all types of errors, such as `NoSuchKey`, so user can retry (which may fix the error) quickly

Code after:

func (scheduler *taskScheduler) RemoveByNode(node int64) {
	scheduler.rwmutex.Lock()
	defer scheduler.rwmutex.Unlock()

	for _, task := range scheduler.segmentTasks {
		if scheduler.isRelated(task, node) {
			scheduler.remove(task)
		}
	}
	for _, task := range scheduler.channelTasks {
		if scheduler.isRelated(task, node) {
			scheduler.remove(task)
		}
	}
}

func (scheduler *taskScheduler) recordSegmentTaskError(task *SegmentTask) {
	var errCode commonpb.ErrorCode
	if errors.Is(task.Err(), ErrInsufficientMemory) {
		errCode = commonpb.ErrorCode_InsufficientMemoryToLoad
	} else {
		errCode = commonpb.ErrorCode_UnexpectedError
	}
	meta.GlobalFailedLoadCache.Put(task.collectionID, errCode, task.Err())
}

func (scheduler *taskScheduler) remove(task Task) {
	log := log.With(
		zap.Int64("taskID", task.ID()),
		zap.Int64("collectionID", task.CollectionID()),
		zap.Int64("replicaID", task.ReplicaID()),
		zap.Int32("taskStatus", task.Status()),
	)
	task.Cancel()
	scheduler.tasks.Remove(task.ID())
	scheduler.waitQueue.Remove(task)
	scheduler.processQueue.Remove(task)

	switch task := task.(type) {
	case *SegmentTask:
		index := NewReplicaSegmentIndex(task)
		delete(scheduler.segmentTasks, index)
		log = log.With(zap.Int64("segmentID", task.SegmentID()))
		if task.Err() != nil {
			log.Warn("task scheduler recordSegmentTaskError", zap.Error(task.err))
			scheduler.recordSegmentTaskError(task)
		}

	case *ChannelTask:
		index := replicaChannelIndex{task.ReplicaID(), task.Channel()}