milvus/internal/indexnode/indexnode.go

package indexnode

import (
	"context"
	"log"
	"math/rand"
	"time"

	"github.com/zilliztech/milvus-distributed/internal/errors"
	"github.com/zilliztech/milvus-distributed/internal/kv"
	miniokv "github.com/zilliztech/milvus-distributed/internal/kv/minio"
	"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
	"github.com/zilliztech/milvus-distributed/internal/proto/indexpb"
	"github.com/zilliztech/milvus-distributed/internal/util/retry"
	"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
)

const (
	reqTimeoutInterval = time.Second * 10
)

type UniqueID = typeutil.UniqueID
type Timestamp = typeutil.Timestamp

type IndexNode struct {
	loopCtx    context.Context
	loopCancel func()

	sched *TaskScheduler

	kv kv.Base

	serviceClient typeutil.IndexServiceInterface // method factory

	// Add callback functions at different stages
	startCallbacks []func()
	closeCallbacks []func()
}

func Init() {
	rand.Seed(time.Now().UnixNano())
	Params.Init()
}

func CreateIndexNode(ctx context.Context) (*IndexNode, error) {
	ctx1, cancel := context.WithCancel(ctx)
	b := &IndexNode{
		loopCtx:    ctx1,
		loopCancel: cancel,
	}

	connectMinIOFn := func() error {
		option := &miniokv.Option{
			Address:           Params.MinIOAddress,
			AccessKeyID:       Params.MinIOAccessKeyID,
			SecretAccessKeyID: Params.MinIOSecretAccessKey,
			UseSSL:            Params.MinIOUseSSL,
			BucketName:        Params.MinioBucketName,
			CreateBucket:      true,
		}
		var err error
		b.kv, err = miniokv.NewMinIOKV(b.loopCtx, option)
		if err != nil {
			return err
		}
		return nil
	}

	err := retry.Retry(10, time.Millisecond*200, connectMinIOFn)
	if err != nil {
		return nil, err
	}

	b.sched, err = NewTaskScheduler(b.loopCtx, b.kv)
	if err != nil {
		return nil, err
	}

	return b, nil
}

// AddStartCallback adds a callback in the startServer phase.
func (i *IndexNode) AddStartCallback(callbacks ...func()) {
	i.startCallbacks = append(i.startCallbacks, callbacks...)
}

// AddCloseCallback adds a callback in the Close phase.
func (i *IndexNode) AddCloseCallback(callbacks ...func()) {
	i.closeCallbacks = append(i.closeCallbacks, callbacks...)
}

func (i *IndexNode) Init() error {
	return nil
}

func (i *IndexNode) Start() error {
	i.sched.Start()

	// Start callbacks
	for _, cb := range i.startCallbacks {
		cb()
	}

	return nil
}

func (i *IndexNode) stopIndexNodeLoop() {
	i.loopCancel()

	i.sched.Close()
}

// Close closes the server.
func (i *IndexNode) Stop() error {
	i.stopIndexNodeLoop()

	for _, cb := range i.closeCallbacks {
		cb()
	}
	log.Print("IndexNode  closed.")
	return nil
}

func (i *IndexNode) SetServiceClient(serviceClient typeutil.IndexServiceInterface) {
	i.serviceClient = serviceClient
}

func (i *IndexNode) BuildIndex(request *indexpb.BuildIndexCmd) (*commonpb.Status, error) {
	t := newIndexBuildTask()
	t.cmd = request
	t.kv = i.kv
	t.serviceClient = i.serviceClient
	t.nodeID = Params.NodeID
	ctx, cancel := context.WithTimeout(context.Background(), reqTimeoutInterval)
	defer cancel()

	fn := func() error {
		select {
		case <-ctx.Done():
			return errors.New("Enqueue BuildQueue timeout")
		default:
			return i.sched.IndexBuildQueue.Enqueue(t)
		}
	}
	ret := &commonpb.Status{
		ErrorCode: commonpb.ErrorCode_SUCCESS,
	}

	err := fn()
	if err != nil {
		ret.ErrorCode = commonpb.ErrorCode_UNEXPECTED_ERROR
		ret.Reason = err.Error()
		return ret, nil
	}
	log.Println("index scheduler successfully with indexID = ", request.IndexID)
	err = t.WaitToFinish()
	log.Println("build index finish ...err = ", err)
	if err != nil {
		ret.ErrorCode = commonpb.ErrorCode_UNEXPECTED_ERROR
		ret.Reason = err.Error()
		return ret, nil
	}
	return ret, nil
}