milvus/internal/indexbuilder/task.go

package indexbuilder

import (
	"context"
	"log"
	"time"

	"github.com/zilliztech/milvus-distributed/internal/allocator"
	"github.com/zilliztech/milvus-distributed/internal/errors"
	"github.com/zilliztech/milvus-distributed/internal/proto/indexbuilderpb"
)

type task interface {
	ID() UniqueID       // return ReqID
	SetID(uid UniqueID) // set ReqID
	PreExecute() error
	Execute() error
	PostExecute() error
	WaitToFinish() error
	Notify(err error)
	OnEnqueue() error
}

type BaseTask struct {
	done  chan error
	ctx   context.Context
	id    UniqueID
	table *metaTable
}

func (bt *BaseTask) ID() UniqueID {
	return bt.id
}

func (bt *BaseTask) setID(id UniqueID) {
	bt.id = id
}

func (bt *BaseTask) WaitToFinish() error {
	select {
	case <-bt.ctx.Done():
		return errors.New("timeout")
	case err := <-bt.done:
		return err
	}
}

func (bt *BaseTask) Notify(err error) {
	bt.done <- err
}

type IndexAddTask struct {
	BaseTask
	req         *indexbuilderpb.BuildIndexRequest
	indexID     UniqueID
	idAllocator *allocator.IDAllocator
	buildQueue  TaskQueue
}

func (it *IndexAddTask) SetID(ID UniqueID) {
	it.BaseTask.setID(ID)
}

func (it *IndexAddTask) OnEnqueue() error {
	var err error
	it.indexID, err = it.idAllocator.AllocOne()
	if err != nil {
		return err
	}
	return nil
}

func (it *IndexAddTask) PreExecute() error {
	log.Println("pretend to check Index Req")
	err := it.table.AddIndex(it.indexID, it.req)
	if err != nil {
		return err
	}
	return nil
}

func (it *IndexAddTask) Execute() error {
	t := newIndexBuildTask()
	t.table = it.table
	t.indexID = it.indexID
	var cancel func()
	t.ctx, cancel = context.WithTimeout(it.ctx, reqTimeoutInterval)
	defer cancel()

	fn := func() error {
		select {
		case <-t.ctx.Done():
			return errors.New("index add timeout")
		default:
			return it.buildQueue.Enqueue(t)
		}
	}
	return fn()
}

func (it *IndexAddTask) PostExecute() error {
	return nil
}

func NewIndexAddTask() *IndexAddTask {
	return &IndexAddTask{
		BaseTask: BaseTask{
			done: make(chan error),
		},
	}
}

type IndexBuildTask struct {
	BaseTask
	indexID   UniqueID
	indexMeta *indexbuilderpb.IndexMeta
}

func newIndexBuildTask() *IndexBuildTask {
	return &IndexBuildTask{
		BaseTask: BaseTask{
			done: make(chan error, 1), // intend to do this
		},
	}
}

func (it *IndexBuildTask) SetID(ID UniqueID) {
	it.BaseTask.setID(ID)
}

func (it *IndexBuildTask) OnEnqueue() error {
	return it.table.UpdateIndexEnqueTime(it.indexID, time.Now())
}

func (it *IndexBuildTask) PreExecute() error {
	return it.table.UpdateIndexScheduleTime(it.indexID, time.Now())
}

func (it *IndexBuildTask) Execute() error {
	err := it.table.UpdateIndexStatus(it.indexID, indexbuilderpb.IndexStatus_INPROGRESS)
	if err != nil {
		return err
	}
	time.Sleep(time.Second)
	log.Println("Pretend to Execute for 1 second")
	return nil
}

func (it *IndexBuildTask) PostExecute() error {
	dataPaths := []string{"file1", "file2"}
	return it.table.CompleteIndex(it.indexID, dataPaths)
}