Add flush, build index, load index scheduler

Signed-off-by: godchen <qingxiang.chen@zilliz.com>

internal/master/client.go

@@ -0,0 +1,101 @@
+package master
+
+import (
+	"time"
+
+	buildindexclient "github.com/zilliztech/milvus-distributed/internal/indexbuilder/client"
+	"github.com/zilliztech/milvus-distributed/internal/proto/indexbuilderpb"
+	writerclient "github.com/zilliztech/milvus-distributed/internal/writenode/client"
+)
+
+type WriteNodeClient interface {
+	FlushSegment(segmentID UniqueID) error
+	DescribeSegment(segmentID UniqueID) (*writerclient.SegmentDescription, error)
+	GetInsertBinlogPaths(segmentID UniqueID) (map[UniqueID][]string, error)
+}
+
+type MockWriteNodeClient struct {
+	segmentID UniqueID
+	flushTime time.Time
+}
+
+func (m *MockWriteNodeClient) FlushSegment(segmentID UniqueID) error {
+	m.flushTime = time.Now()
+	m.segmentID = segmentID
+	return nil
+}
+
+func (m *MockWriteNodeClient) DescribeSegment(segmentID UniqueID) (*writerclient.SegmentDescription, error) {
+	now := time.Now()
+	if now.Sub(m.flushTime).Seconds() > 3 {
+		return &writerclient.SegmentDescription{
+			SegmentID: segmentID,
+			IsClosed:  true,
+			OpenTime:  0,
+			CloseTime: 1,
+		}, nil
+	}
+	return &writerclient.SegmentDescription{
+		SegmentID: segmentID,
+		IsClosed:  false,
+		OpenTime:  0,
+		CloseTime: 1,
+	}, nil
+}
+
+func (m *MockWriteNodeClient) GetInsertBinlogPaths(segmentID UniqueID) (map[UniqueID][]string, error) {
+	return map[UniqueID][]string{
+		1:   {"/binlog/insert/file_1"},
+		100: {"/binlog/insert/file_100"},
+	}, nil
+}
+
+type BuildIndexClient interface {
+	BuildIndexWithoutID(columnDataPaths []string, typeParams map[string]string, indexParams map[string]string) (UniqueID, error)
+	DescribeIndex(indexID UniqueID) (*buildindexclient.IndexDescription, error)
+	GetIndexFilePaths(indexID UniqueID) ([]string, error)
+}
+
+type MockBuildIndexClient struct {
+	buildTime time.Time
+}
+
+func (m *MockBuildIndexClient) BuildIndexWithoutID(columnDataPaths []string, typeParams map[string]string, indexParams map[string]string) (UniqueID, error) {
+	m.buildTime = time.Now()
+	return 1, nil
+}
+
+func (m *MockBuildIndexClient) DescribeIndex(indexID UniqueID) (*buildindexclient.IndexDescription, error) {
+	now := time.Now()
+	if now.Sub(m.buildTime).Seconds() > 3 {
+		return &buildindexclient.IndexDescription{
+			ID:                1,
+			Status:            indexbuilderpb.IndexStatus_FINISHED,
+			EnqueueTime:       time.Now(),
+			ScheduleTime:      time.Now(),
+			BuildCompleteTime: time.Now(),
+		}, nil
+	}
+	return &buildindexclient.IndexDescription{
+		ID:                1,
+		Status:            indexbuilderpb.IndexStatus_INPROGRESS,
+		EnqueueTime:       time.Now(),
+		ScheduleTime:      time.Now(),
+		BuildCompleteTime: time.Now(),
+	}, nil
+}
+
+func (m *MockBuildIndexClient) GetIndexFilePaths(indexID UniqueID) ([]string, error) {
+	return []string{"/binlog/index/file_1", "/binlog/index/file_2", "/binlog/index/file_3"}, nil
+}
+
+type LoadIndexClient interface {
+	LoadIndex(indexPaths []string, segmentID int64, fieldID int64, fieldName string) error
+}
+
+type MockLoadIndexClient struct {
+}
+
+func (m *MockLoadIndexClient) LoadIndex(indexPaths []string, segmentID int64, fieldID int64, fieldName string) error {
+	return nil
+}

internal/master/flush_scheduler.go

@@ -0,0 +1,119 @@
+package master
+
+import (
+	"context"
+	"log"
+	"time"
+
+	"github.com/zilliztech/milvus-distributed/internal/errors"
+)
+
+type FlushScheduler struct {
+	client              WriteNodeClient
+	metaTable           *metaTable
+	segmentFlushChan    chan UniqueID
+	segmentDescribeChan chan UniqueID
+	indexBuilderSch     persistenceScheduler
+
+	ctx    context.Context
+	cancel context.CancelFunc
+}
+
+func NewFlushScheduler(ctx context.Context, client WriteNodeClient, metaTable *metaTable, buildScheduler *IndexBuildScheduler) *FlushScheduler {
+	ctx2, cancel := context.WithCancel(ctx)
+
+	return &FlushScheduler{
+		client:              client,
+		metaTable:           metaTable,
+		indexBuilderSch:     buildScheduler,
+		segmentFlushChan:    make(chan UniqueID, 100),
+		segmentDescribeChan: make(chan UniqueID, 100),
+		ctx:                 ctx2,
+		cancel:              cancel,
+	}
+}
+
+func (scheduler *FlushScheduler) schedule(id interface{}) error {
+	segmentID := id.(UniqueID)
+	err := scheduler.client.FlushSegment(segmentID)
+	log.Printf("flush segment %d", segmentID)
+	if err != nil {
+		return err
+	}
+
+	scheduler.segmentDescribeChan <- segmentID
+
+	return nil
+}
+func (scheduler *FlushScheduler) describe() error {
+	for {
+		select {
+		case <-scheduler.ctx.Done():
+			{
+				log.Printf("broadcast context done, exit")
+				return errors.New("broadcast done exit")
+			}
+		case singleSegmentID := <-scheduler.segmentDescribeChan:
+			for {
+				description, err := scheduler.client.DescribeSegment(singleSegmentID)
+				if err != nil {
+					return err
+				}
+				if description.IsClosed {
+					log.Printf("flush segment %d is closed", singleSegmentID)
+					mapData, err := scheduler.client.GetInsertBinlogPaths(singleSegmentID)
+					if err != nil {
+						return err
+					}
+					for fieldID, data := range mapData {
+						info := &IndexBuildInfo{
+							segmentID:      singleSegmentID,
+							fieldID:        fieldID,
+							binlogFilePath: data,
+						}
+						err = scheduler.indexBuilderSch.Enqueue(info)
+						log.Printf("segment %d field %d enqueue build index scheduler", singleSegmentID, fieldID)
+						if err != nil {
+							return err
+						}
+					}
+					//TODO: Save data to meta table
+					log.Printf("flush segment %d finished", singleSegmentID)
+					break
+				}
+				time.Sleep(1 * time.Second)
+			}
+		}
+	}
+
+}
+
+func (scheduler *FlushScheduler) scheduleLoop() {
+	for {
+		select {
+		case id := <-scheduler.segmentFlushChan:
+			err := scheduler.schedule(id)
+			if err != nil {
+				log.Println(err)
+			}
+		case <-scheduler.ctx.Done():
+			log.Print("server is closed, exit flush scheduler loop")
+			return
+		}
+	}
+}
+
+func (scheduler *FlushScheduler) Enqueue(id interface{}) error {
+	scheduler.segmentFlushChan <- id.(UniqueID)
+	return nil
+}
+
+func (scheduler *FlushScheduler) Start() error {
+	go scheduler.scheduleLoop()
+	go scheduler.describe()
+	return nil
+}
+
+func (scheduler *FlushScheduler) Close() {
+	scheduler.cancel()
+}

internal/master/index_builder_scheduler.go

@@ -0,0 +1,147 @@
+package master
+
+import (
+	"context"
+	"log"
+	"time"
+
+	"github.com/zilliztech/milvus-distributed/internal/errors"
+	"github.com/zilliztech/milvus-distributed/internal/proto/indexbuilderpb"
+)
+
+type IndexBuildInfo struct {
+	segmentID      UniqueID
+	fieldID        UniqueID
+	binlogFilePath []string
+}
+type IndexBuildChannelInfo struct {
+	id   UniqueID
+	info *IndexBuildInfo
+}
+
+type IndexBuildScheduler struct {
+	client          BuildIndexClient
+	metaTable       *metaTable
+	indexBuildChan  chan *IndexBuildInfo
+	indexLoadSch    persistenceScheduler
+	indexDescribeID chan UniqueID
+	indexDescribe   chan *IndexBuildChannelInfo
+
+	ctx    context.Context
+	cancel context.CancelFunc
+}
+
+func NewIndexBuildScheduler(ctx context.Context, client BuildIndexClient, metaTable *metaTable, indexLoadScheduler *IndexLoadScheduler) *IndexBuildScheduler {
+	ctx2, cancel := context.WithCancel(ctx)
+
+	return &IndexBuildScheduler{
+		client:         client,
+		metaTable:      metaTable,
+		indexLoadSch:   indexLoadScheduler,
+		indexBuildChan: make(chan *IndexBuildInfo, 100),
+		indexDescribe:  make(chan *IndexBuildChannelInfo, 100),
+		ctx:            ctx2,
+		cancel:         cancel,
+	}
+}
+
+func (scheduler *IndexBuildScheduler) schedule(info interface{}) error {
+	indexBuildInfo := info.(*IndexBuildInfo)
+	indexID, err := scheduler.client.BuildIndexWithoutID(indexBuildInfo.binlogFilePath, nil, nil)
+	log.Printf("build index for segment %d field %d", indexBuildInfo.segmentID, indexBuildInfo.fieldID)
+	if err != nil {
+		return err
+	}
+	scheduler.indexDescribe <- &IndexBuildChannelInfo{
+		id:   indexID,
+		info: indexBuildInfo,
+	}
+	return nil
+}
+
+func (scheduler *IndexBuildScheduler) describe() error {
+	for {
+		select {
+		case <-scheduler.ctx.Done():
+			{
+				log.Printf("broadcast context done, exit")
+				return errors.New("broadcast done exit")
+			}
+		case channelInfo := <-scheduler.indexDescribe:
+			indexID := channelInfo.id
+			indexBuildInfo := channelInfo.info
+			for {
+				description, err := scheduler.client.DescribeIndex(channelInfo.id)
+				if err != nil {
+					return err
+				}
+				if description.Status == indexbuilderpb.IndexStatus_FINISHED {
+					log.Printf("build index for segment %d field %d is finished", indexBuildInfo.segmentID, indexBuildInfo.fieldID)
+					filePaths, err := scheduler.client.GetIndexFilePaths(indexID)
+					if err != nil {
+						return err
+					}
+
+					//TODO: remove fileName
+					var fieldName string
+					segMeta := scheduler.metaTable.segID2Meta[indexBuildInfo.segmentID]
+					collMeta := scheduler.metaTable.collID2Meta[segMeta.CollectionID]
+					if collMeta.Schema != nil {
+						for _, field := range collMeta.Schema.Fields {
+							if field.FieldID == indexBuildInfo.fieldID {
+								fieldName = field.Name
+							}
+						}
+					}
+
+					info := &IndexLoadInfo{
+						segmentID:      indexBuildInfo.segmentID,
+						fieldID:        indexBuildInfo.fieldID,
+						fieldName:      fieldName,
+						indexFilePaths: filePaths,
+					}
+					//TODO: Save data to meta table
+
+					err = scheduler.indexLoadSch.Enqueue(info)
+					log.Printf("build index for segment %d field %d enqueue load index", indexBuildInfo.segmentID, indexBuildInfo.fieldID)
+					if err != nil {
+						return err
+					}
+					log.Printf("build index for segment %d field %d finished", indexBuildInfo.segmentID, indexBuildInfo.fieldID)
+				}
+				time.Sleep(1 * time.Second)
+			}
+		}
+	}
+
+}
+
+func (scheduler *IndexBuildScheduler) scheduleLoop() {
+	for {
+		select {
+		case info := <-scheduler.indexBuildChan:
+			err := scheduler.schedule(info)
+			if err != nil {
+				log.Println(err)
+			}
+		case <-scheduler.ctx.Done():
+			log.Print("server is closed, exit index build loop")
+			return
+		}
+	}
+}
+
+func (scheduler *IndexBuildScheduler) Enqueue(info interface{}) error {
+	scheduler.indexBuildChan <- info.(*IndexBuildInfo)
+	return nil
+}
+
+func (scheduler *IndexBuildScheduler) Start() error {
+	go scheduler.scheduleLoop()
+	go scheduler.describe()
+	return nil
+}
+
+func (scheduler *IndexBuildScheduler) Close() {
+	scheduler.cancel()
+}

internal/master/index_load_scheduler.go

@@ -0,0 +1,74 @@
+package master
+
+import (
+	"context"
+	"log"
+)
+
+type IndexLoadInfo struct {
+	segmentID      UniqueID
+	fieldID        UniqueID
+	fieldName      string
+	indexFilePaths []string
+}
+
+type IndexLoadScheduler struct {
+	indexLoadChan chan *IndexLoadInfo
+	client        LoadIndexClient
+	metaTable     *metaTable
+
+	ctx    context.Context
+	cancel context.CancelFunc
+}
+
+func NewIndexLoadScheduler(ctx context.Context, client LoadIndexClient, metaTable *metaTable) *IndexLoadScheduler {
+	ctx2, cancel := context.WithCancel(ctx)
+	indexLoadChan := make(chan *IndexLoadInfo, 100)
+
+	return &IndexLoadScheduler{
+		client:        client,
+		metaTable:     metaTable,
+		indexLoadChan: indexLoadChan,
+		ctx:           ctx2,
+		cancel:        cancel,
+	}
+}
+
+func (scheduler *IndexLoadScheduler) schedule(info interface{}) error {
+	indexLoadInfo := info.(*IndexLoadInfo)
+	err := scheduler.client.LoadIndex(indexLoadInfo.indexFilePaths, indexLoadInfo.segmentID, indexLoadInfo.fieldID, indexLoadInfo.fieldName)
+	//TODO: Save data to meta table
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+func (scheduler *IndexLoadScheduler) scheduleLoop() {
+	for {
+		select {
+		case info := <-scheduler.indexLoadChan:
+			err := scheduler.schedule(info)
+			if err != nil {
+				log.Println(err)
+			}
+		case <-scheduler.ctx.Done():
+			log.Print("server is closed, exit flush scheduler loop")
+			return
+		}
+	}
+}
+
+func (scheduler *IndexLoadScheduler) Enqueue(info interface{}) error {
+	scheduler.indexLoadChan <- info.(*IndexLoadInfo)
+	return nil
+}
+
+func (scheduler *IndexLoadScheduler) Start() error {
+	go scheduler.scheduleLoop()
+	return nil
+}
+
+func (scheduler *IndexLoadScheduler) Close() {
+	scheduler.cancel()
+}

internal/master/master.go

@@ -44,6 +44,9 @@ type Master struct {
 
 	kvBase               *etcdkv.EtcdKV
 	scheduler            *ddRequestScheduler
+	flushSch             *FlushScheduler
+	indexBuildSch        *IndexBuildScheduler
+	indexLoadSch         *IndexLoadScheduler
 	metaTable            *metaTable
 	timesSyncMsgProducer *timeSyncMsgProducer
 
@@ -176,15 +179,24 @@ func CreateServer(ctx context.Context) (*Master, error) {
 	m.scheduler.SetDDMsgStream(pulsarDDStream)
 	m.scheduler.SetIDAllocator(func() (UniqueID, error) { return m.idAllocator.AllocOne() })
 
+	flushClient := &MockWriteNodeClient{}
+	buildIndexClient := &MockBuildIndexClient{}
+	loadIndexClient := &MockLoadIndexClient{}
+
+	m.indexLoadSch = NewIndexLoadScheduler(ctx, loadIndexClient, m.metaTable)
+	m.indexBuildSch = NewIndexBuildScheduler(ctx, buildIndexClient, m.metaTable, m.indexLoadSch)
+	m.flushSch = NewFlushScheduler(ctx, flushClient, m.metaTable, m.indexBuildSch)
+
 	m.segmentAssigner = NewSegmentAssigner(ctx, metakv,
 		func() (Timestamp, error) { return m.tsoAllocator.AllocOne() },
 		proxyTtBarrierWatcher,
 	)
+
 	m.segmentManager, err = NewSegmentManager(ctx, metakv,
 		func() (UniqueID, error) { return m.idAllocator.AllocOne() },
 		func() (Timestamp, error) { return m.tsoAllocator.AllocOne() },
 		writeNodeTtBarrierWatcher,
-		&MockFlushScheduler{}, // todo replace mock with real flush scheduler
+		m.flushSch,
 		m.segmentAssigner)
 
 	if err != nil {
@@ -283,6 +295,18 @@ func (s *Master) startServerLoop(ctx context.Context, grpcPort int64) error {
 	if err := s.scheduler.Start(); err != nil {
 		return err
 	}
+	s.serverLoopWg.Add(1)
+	if err := s.indexLoadSch.Start(); err != nil {
+		return err
+	}
+	s.serverLoopWg.Add(1)
+	if err := s.indexBuildSch.Start(); err != nil {
+		return err
+	}
+	s.serverLoopWg.Add(1)
+	if err := s.flushSch.Start(); err != nil {
+		return err
+	}
 
 	s.serverLoopWg.Add(1)
 	go s.grpcLoop(grpcPort)
@@ -305,6 +329,12 @@ func (s *Master) stopServerLoop() {
 	s.serverLoopWg.Done()
 	s.scheduler.Close()
 	s.serverLoopWg.Done()
+	s.flushSch.Close()
+	s.serverLoopWg.Done()
+	s.indexBuildSch.Close()
+	s.serverLoopWg.Done()
+	s.indexLoadSch.Close()
+	s.serverLoopWg.Done()
 
 	if s.grpcServer != nil {
 		s.grpcServer.GracefulStop()

internal/master/persistence_scheduler_test.go

@@ -0,0 +1,85 @@
+package master
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	etcdkv "github.com/zilliztech/milvus-distributed/internal/kv/etcd"
+	"github.com/zilliztech/milvus-distributed/internal/proto/indexbuilderpb"
+	"go.etcd.io/etcd/clientv3"
+)
+
+func TestPersistenceScheduler(t *testing.T) {
+	//Init environment Params
+	Init()
+
+	ctx := context.Background()
+
+	//Init client, use Mock instead
+	flushClient := &MockWriteNodeClient{}
+	buildIndexClient := &MockBuildIndexClient{}
+	loadIndexClient := &MockLoadIndexClient{}
+
+	etcdAddr := Params.EtcdAddress
+	cli, err := clientv3.New(clientv3.Config{Endpoints: []string{etcdAddr}})
+	assert.Nil(t, err)
+	etcdKV := etcdkv.NewEtcdKV(cli, "/etcd/test/root")
+
+	_, err = cli.Delete(context.TODO(), "/etcd/test/root", clientv3.WithPrefix())
+	assert.Nil(t, err)
+
+	meta, err := NewMetaTable(etcdKV)
+	assert.Nil(t, err)
+	defer meta.client.Close()
+
+	//Init scheduler
+	indexLoadSch := NewIndexLoadScheduler(ctx, loadIndexClient, meta)
+	indexBuildSch := NewIndexBuildScheduler(ctx, buildIndexClient, meta, indexLoadSch)
+	flushSch := NewFlushScheduler(ctx, flushClient, meta, indexBuildSch)
+
+	//scheduler start
+	err = indexLoadSch.Start()
+	assert.Nil(t, err)
+	defer indexLoadSch.Close()
+
+	err = indexBuildSch.Start()
+	assert.Nil(t, err)
+	defer indexBuildSch.Close()
+
+	err = flushSch.Start()
+	assert.Nil(t, err)
+	defer flushSch.Close()
+
+	//start from flush scheduler
+	err = flushSch.Enqueue(UniqueID(1))
+	assert.Nil(t, err)
+	//wait flush segment request sent to write node
+	time.Sleep(100 * time.Millisecond)
+	segDes, err := flushClient.DescribeSegment(UniqueID(1))
+	assert.Nil(t, err)
+	assert.Equal(t, false, segDes.IsClosed)
+
+	//wait flush to finish
+	time.Sleep(3 * time.Second)
+
+	segDes, err = flushClient.DescribeSegment(UniqueID(1))
+	assert.Nil(t, err)
+	assert.Equal(t, UniqueID(1), segDes.SegmentID)
+	assert.Equal(t, true, segDes.IsClosed)
+
+	//wait flush segment request sent to build index node
+	time.Sleep(100 * time.Microsecond)
+	idxDes, err := buildIndexClient.DescribeIndex(UniqueID(1))
+	assert.Nil(t, err)
+	assert.Equal(t, indexbuilderpb.IndexStatus_INPROGRESS, idxDes.Status)
+
+	//wait build index to finish
+	time.Sleep(3 * time.Second)
+
+	idxDes, err = buildIndexClient.DescribeIndex(UniqueID(1))
+	assert.Nil(t, err)
+	assert.Equal(t, indexbuilderpb.IndexStatus_FINISHED, idxDes.Status)
+
+}