milvus/internal/proxy/task_scheduler.go

package proxy

import (
	"container/list"
	"log"
	"sync"

	"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
)

type baseTaskQueue struct {
	unissuedTasks *list.List
	activeTasks   map[typeutil.Timestamp]*task
	utLock        sync.Mutex
	atLock        sync.Mutex
}

type ddTaskQueue struct {
	baseTaskQueue
	lock sync.Mutex
}

type dmTaskQueue struct {
	baseTaskQueue
}

type dqTaskQueue struct {
	baseTaskQueue
}

func (queue *baseTaskQueue) Empty() bool {
	queue.utLock.Lock()
	defer queue.utLock.Unlock()
	queue.atLock.Lock()
	defer queue.atLock.Unlock()
	return queue.unissuedTasks.Len() <= 0 && len(queue.activeTasks) <= 0
}

func (queue *baseTaskQueue) AddUnissuedTask(t *task) {
	queue.utLock.Lock()
	defer queue.utLock.Unlock()
	queue.unissuedTasks.PushBack(t)
}

func (queue *baseTaskQueue) FrontUnissuedTask() *task {
	queue.utLock.Lock()
	defer queue.utLock.Unlock()
	if queue.unissuedTasks.Len() <= 0 {
		log.Fatal("sorry, but the unissued task list is empty!")
		return nil
	}
	return queue.unissuedTasks.Front().Value.(*task)
}

func (queue *baseTaskQueue) PopUnissuedTask() *task {
	queue.utLock.Lock()
	defer queue.utLock.Unlock()
	if queue.unissuedTasks.Len() <= 0 {
		log.Fatal("sorry, but the unissued task list is empty!")
		return nil
	}
	ft := queue.unissuedTasks.Front()
	return queue.unissuedTasks.Remove(ft).(*task)
}

func (queue *baseTaskQueue) AddActiveTask(t *task) {
	queue.atLock.Lock()
	defer queue.atLock.Lock()
	ts := (*t).GetTs()
	_, ok := queue.activeTasks[ts]
	if ok {
		log.Fatalf("task with timestamp %d already in active task list!", ts)
	}
	queue.activeTasks[ts] = t
}

func (queue *baseTaskQueue) PopActiveTask(ts typeutil.Timestamp) *task {
	queue.atLock.Lock()
	defer queue.atLock.Lock()
	t, ok := queue.activeTasks[ts]
	if ok {
		delete(queue.activeTasks, ts)
		return t
	}
	log.Fatalf("sorry, but the timestamp %d was not found in the active task list!", ts)
	return nil
}

func (queue *baseTaskQueue) TaskDoneTest(ts typeutil.Timestamp) bool {
	queue.utLock.Lock()
	defer queue.utLock.Unlock()
	for e := queue.unissuedTasks.Front(); e != nil; e = e.Next() {
		if (*(e.Value.(*task))).GetTs() >= ts {
			return false
		}
	}

	queue.atLock.Lock()
	defer queue.atLock.Unlock()
	for ats := range queue.activeTasks {
		if ats >= ts {
			return false
		}
	}

	return true
}

func (queue *ddTaskQueue) Enqueue(t *task) error {
	queue.lock.Lock()
	defer queue.lock.Unlock()
	// TODO: set Ts, ReqId, ProxyId
	queue.AddUnissuedTask(t)
	return nil
}

func (queue *dmTaskQueue) Enqueue(t *task) error {
	// TODO: set Ts, ReqId, ProxyId
	queue.AddUnissuedTask(t)
	return nil
}

func (queue *dqTaskQueue) Enqueue(t *task) error {
	// TODO: set Ts, ReqId, ProxyId
	queue.AddUnissuedTask(t)
	return nil
}

type taskScheduler struct {
	DdQueue *ddTaskQueue
	DmQueue *dmTaskQueue
	DqQueue *dqTaskQueue

	// tsAllocator, ReqIdAllocator
}

func (sched *taskScheduler) scheduleDdTask() *task {
	return sched.DdQueue.PopUnissuedTask()
}

func (sched *taskScheduler) scheduleDmTask() *task {
	return sched.DmQueue.PopUnissuedTask()
}

func (sched *taskScheduler) scheduleDqTask() *task {
	return sched.DqQueue.PopUnissuedTask()
}

func (sched *taskScheduler) Start() error {
	go func() {
		for {
			if sched.DdQueue.Empty() {
				continue
			}
			t := sched.scheduleDdTask()
			if err := (*t).PreExecute(); err != nil {
				return
			}
			if err := (*t).Execute(); err != nil {
				return
			}
			if err := (*t).PostExecute(); err != nil {
				return
			}
			if err := (*t).WaitToFinish(); err != nil {
				return
			}
			if err := (*t).Notify(); err != nil {
				return
			}
		}
	}()
	go func() {
		for {
			if sched.DdQueue.Empty() {
				continue
			}
			t := sched.scheduleDmTask()
			if err := (*t).PreExecute(); err != nil {
				return
			}
			if err := (*t).Execute(); err != nil {
				return
			}
			if err := (*t).PostExecute(); err != nil {
				return
			}
			if err := (*t).WaitToFinish(); err != nil {
				return
			}
			if err := (*t).Notify(); err != nil {
				return
			}
		}
	}()
	go func() {
		for {
			if sched.DdQueue.Empty() {
				continue
			}
			t := sched.scheduleDqTask()
			if err := (*t).PreExecute(); err != nil {
				return
			}
			if err := (*t).Execute(); err != nil {
				return
			}
			if err := (*t).PostExecute(); err != nil {
				return
			}
			if err := (*t).WaitToFinish(); err != nil {
				return
			}
			if err := (*t).Notify(); err != nil {
				return
			}
		}
	}()
	return nil
}

func (sched *taskScheduler) TaskDoneTest(ts typeutil.Timestamp) bool {
	ddTaskDone := sched.DdQueue.TaskDoneTest(ts)
	dmTaskDone := sched.DmQueue.TaskDoneTest(ts)
	dqTaskDone := sched.DqQueue.TaskDoneTest(ts)
	return ddTaskDone && dmTaskDone && dqTaskDone
}