milvus/internal/master/flush_scheduler.go

package master

import (
	"context"
	"log"
	"time"

	"github.com/zilliztech/milvus-distributed/internal/proto/etcdpb"

	"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
						}
					}
					// Save data to meta table
					segMeta, err := scheduler.metaTable.GetSegmentByID(singleSegmentID)
					if err != nil {
						return err
					}
					segMeta.BinlogFilePaths = make([]*etcdpb.FieldBinlogFiles, 0)
					for k, v := range mapData {
						segMeta.BinlogFilePaths = append(segMeta.BinlogFilePaths, &etcdpb.FieldBinlogFiles{
							FieldID:     k,
							BinlogFiles: v,
						})
					}
					if err = scheduler.metaTable.UpdateSegment(segMeta); err != nil {
						return err
					}
					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()
}