milvus/internal/allocator/allocator.go

package allocator

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

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

const (
	maxMergeRequests = 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
}

func (req *idRequest) Wait() {
	req.baseRequest.Wait()
}

type tsoRequest struct {
	baseRequest
	timestamp Timestamp
	count     uint32
}

func (req *tsoRequest) Wait() {
	req.baseRequest.Wait()
}

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

	tChan       tickerChan
	syncFunc    func()
	processFunc func(req request)
}

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, _ := context.WithTimeout(context.Background(), 2*time.Second)
	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) mainLoop() {
	defer ta.wg.Done()

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

	defaultSize := maxMergeRequests + 1
	reqs := make([]request, defaultSize)
	for {
		select {
		case <-ta.tChan.Chan():
			ta.sync()
		case first := <-ta.reqs:
			pendingPlus1 := len(ta.reqs) + 1
			reqs[0] = first
			for i := 1; i < pendingPlus1; i++ {
				reqs[i] = <-ta.reqs
			}
			ta.finishRequest(reqs[:pendingPlus1])

		case <-loopCtx.Done():
			return
		}

	}
}

func (ta *Allocator) sync() {
	if ta.syncFunc != nil {
		ta.syncFunc()
		ta.tChan.Reset()
		fmt.Println("synced")
	}
}

func (ta *Allocator) finishRequest(reqs []request) {
	for i := 0; i < len(reqs); i++ {
		ta.processFunc(reqs[i])
		if reqs[i] != nil {
			reqs[i].Notify(nil)
		}
	}
}

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"))
}