fix restart querycoord blocking issue (#18801)

Signed-off-by: Enwei Jiao <enwei.jiao@zilliz.com>

Signed-off-by: Enwei Jiao <enwei.jiao@zilliz.com>
This commit is contained in:
Enwei Jiao 2022-08-25 18:42:54 +08:00 committed by GitHub
parent 305601ad25
commit 232f48784e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 35 additions and 14 deletions

View File

@ -60,6 +60,13 @@ func (queue *taskQueue) taskEmpty() bool {
return queue.tasks.Len() == 0
}
func (queue *taskQueue) Len() int {
queue.Lock()
defer queue.Unlock()
return queue.tasks.Len()
}
func (queue *taskQueue) taskFull() bool {
return int64(queue.tasks.Len()) >= queue.maxTask
}
@ -123,6 +130,9 @@ func (queue *taskQueue) addTask(t task) {
func (queue *taskQueue) addTaskToFront(t task) {
queue.taskChan <- 1
queue.Lock()
defer queue.Unlock()
if queue.tasks.Len() == 0 {
queue.tasks.PushBack(t)
} else {
@ -272,6 +282,9 @@ func (scheduler *TaskScheduler) reloadFromKV() error {
triggerTasks[taskID].setState(state)
}
// triggerTaskQueue's size is 1024, if the size of triggerTasks if large than 1024 the loop will be blocked,
// so run it in a standalone goroutine
go func() {
var doneTriggerTask task
for _, t := range triggerTasks {
if t.getState() == taskDone {
@ -289,6 +302,7 @@ func (scheduler *TaskScheduler) reloadFromKV() error {
if doneTriggerTask != nil {
scheduler.triggerTaskQueue.addTaskToFront(doneTriggerTask)
}
}()
return nil
}

View File

@ -462,6 +462,13 @@ func TestReloadTaskFromKV(t *testing.T) {
taskScheduler.reloadFromKV()
// wait for the addtask goroutine finished
assert.Eventually(t,
func() bool {
return taskScheduler.triggerTaskQueue.Len() == len(kvs)-2
},
10*time.Second, 100*time.Millisecond)
task := taskScheduler.triggerTaskQueue.popTask()
assert.Equal(t, taskDone, task.getState())
assert.Equal(t, 1, len(task.getChildTask()))