milvus/internal/allocator/allocator.go

package allocator

import (
	"context"
	"errors"
	"log"
	"sync"
	"time"

	"github.com/zilliztech/milvus-distributed/internal/proto/masterpb"
	"google.golang.org/grpc"
)

const (
	maxConcurrentRequests = 10000
)

type request interface {
	Wait()
	Notify(error)
	IsValid() bool
}

type baseRequest struct {
	done  chan error
	valid bool
}

func (req *baseRequest) Wait() {
	err := <-req.done
	req.valid = err == nil
}

func (req *baseRequest) IsValid() bool {
	return req.valid
}

func (req *baseRequest) Notify(err error) {
	req.done <- err
}

type idRequest struct {
	baseRequest
	id    UniqueID
	count uint32
}

type tsoRequest struct {
	baseRequest
	timestamp Timestamp
	count     uint32
}

type segRequest struct {
	baseRequest
	count     uint32
	colName   string
	partition string
	segInfo   map[UniqueID]uint32
	channelID int32
}

type syncRequest struct {
	baseRequest
}

type tickerChan interface {
	Chan() <-chan time.Time
	Close()
	Init()
	Reset()
}

type emptyTicker struct {
	tChan <-chan time.Time
}

func (t *emptyTicker) Chan() <-chan time.Time {
	return t.tChan
}

func (t *emptyTicker) Init() {
}

func (t *emptyTicker) Reset() {
}

func (t *emptyTicker) Close() {
}

type ticker struct {
	ticker         *time.Ticker
	updateInterval time.Duration //
}

func (t *ticker) Init() {
	t.ticker = time.NewTicker(t.updateInterval)
}

func (t *ticker) Reset() {
	t.ticker.Reset(t.updateInterval)
}

func (t *ticker) Close() {
	t.ticker.Stop()
}

func (t *ticker) Chan() <-chan time.Time {
	return t.ticker.C
}

type Allocator struct {
	reqs chan request

	wg     sync.WaitGroup
	ctx    context.Context
	cancel context.CancelFunc

	masterAddress string
	masterConn    *grpc.ClientConn
	masterClient  masterpb.MasterClient
	countPerRPC   uint32

	toDoReqs []request

	syncReqs []request

	tChan         tickerChan
	forceSyncChan chan request

	syncFunc    func()
	processFunc func(req request) error
	checkFunc   func(timeout bool) bool
}

func (ta *Allocator) Start() error {
	err := ta.connectMaster()
	if err != nil {
		panic("connect to master failed")
	}
	ta.tChan.Init()
	ta.wg.Add(1)
	go ta.mainLoop()
	return nil
}

func (ta *Allocator) connectMaster() error {
	log.Printf("Connected to master, master_addr=%s", ta.masterAddress)
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	conn, err := grpc.DialContext(ctx, ta.masterAddress, grpc.WithInsecure(), grpc.WithBlock())
	if err != nil {
		log.Printf("Connect to master failed, error= %v", err)
		return err
	}
	log.Printf("Connected to master, master_addr=%s", ta.masterAddress)
	ta.masterConn = conn
	ta.masterClient = masterpb.NewMasterClient(conn)
	return nil
}

func (ta *Allocator) init() {
	ta.forceSyncChan = make(chan request, maxConcurrentRequests)
}

func (ta *Allocator) mainLoop() {
	defer ta.wg.Done()

	loopCtx, loopCancel := context.WithCancel(ta.ctx)
	defer loopCancel()

	for {
		select {

		case first := <-ta.forceSyncChan:
			ta.syncReqs = append(ta.syncReqs, first)
			pending := len(ta.forceSyncChan)
			for i := 0; i < pending; i++ {
				ta.syncReqs = append(ta.syncReqs, <-ta.forceSyncChan)
			}
			ta.sync(true)
			ta.finishSyncRequest()

		case <-ta.tChan.Chan():
			ta.sync(true)

		case first := <-ta.reqs:
			ta.toDoReqs = append(ta.toDoReqs, first)
			pending := len(ta.reqs)
			for i := 0; i < pending; i++ {
				ta.toDoReqs = append(ta.toDoReqs, <-ta.reqs)
			}
			ta.sync(false)

			ta.finishRequest()

		case <-loopCtx.Done():
			return
		}

	}
}

func (ta *Allocator) sync(timeout bool) {
	if ta.syncFunc == nil {
		return
	}
	if ta.checkFunc == nil || !ta.checkFunc(timeout) {
		return
	}

	ta.syncFunc()

	if !timeout {
		ta.tChan.Reset()
	}
}

func (ta *Allocator) finishSyncRequest() {
	for _, req := range ta.syncReqs {
		if req != nil {
			req.Notify(nil)
		}
	}
	ta.syncReqs = ta.syncReqs[0:0]
}

func (ta *Allocator) finishRequest() {
	for _, req := range ta.toDoReqs {
		if req != nil {
			err := ta.processFunc(req)
			req.Notify(err)
		}
	}
	ta.toDoReqs = ta.toDoReqs[0:0]
}

func (ta *Allocator) revokeRequest(err error) {
	n := len(ta.reqs)
	for i := 0; i < n; i++ {
		req := <-ta.reqs
		req.Notify(err)
	}
}

func (ta *Allocator) Close() {
	ta.cancel()
	ta.wg.Wait()
	ta.tChan.Close()
	ta.revokeRequest(errors.New("closing"))
}

func (ta *Allocator) CleanCache() {
	req := &syncRequest{baseRequest: baseRequest{done: make(chan error), valid: false}}
	ta.forceSyncChan <- req
	req.Wait()
}