milvus/internal/proxynode/proxy_node.go

// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software distributed under the License
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
// or implied. See the License for the specific language governing permissions and limitations under the License.

package proxynode

import (
	"context"
	"errors"
	"fmt"
	"math/rand"
	"sync"
	"sync/atomic"
	"time"

	"go.uber.org/zap"

	"github.com/milvus-io/milvus/internal/allocator"
	"github.com/milvus-io/milvus/internal/log"
	"github.com/milvus-io/milvus/internal/msgstream"
	"github.com/milvus-io/milvus/internal/proto/commonpb"
	"github.com/milvus-io/milvus/internal/proto/internalpb"
	"github.com/milvus-io/milvus/internal/proto/milvuspb"
	"github.com/milvus-io/milvus/internal/proto/proxypb"
	"github.com/milvus-io/milvus/internal/proto/querypb"
	"github.com/milvus-io/milvus/internal/types"
	"github.com/milvus-io/milvus/internal/util/funcutil"
	"github.com/milvus-io/milvus/internal/util/sessionutil"
	"github.com/milvus-io/milvus/internal/util/typeutil"
)

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

type ProxyNode struct {
	ctx    context.Context
	cancel func()
	wg     sync.WaitGroup

	initParams *internalpb.InitParams
	ip         string
	port       int

	stateCode atomic.Value

	masterService types.MasterService
	indexService  types.IndexService
	dataService   types.DataService
	proxyService  types.ProxyService
	queryService  types.QueryService

	chMgr channelsMgr

	sched *TaskScheduler
	tick  *timeTick

	chTicker channelsTimeTicker

	idAllocator  *allocator.IDAllocator
	tsoAllocator *TimestampAllocator
	segAssigner  *SegIDAssigner

	session *sessionutil.Session

	queryMsgStream msgstream.MsgStream
	msFactory      msgstream.Factory

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

func NewProxyNode(ctx context.Context, factory msgstream.Factory) (*ProxyNode, error) {
	rand.Seed(time.Now().UnixNano())
	ctx1, cancel := context.WithCancel(ctx)
	node := &ProxyNode{
		ctx:       ctx1,
		cancel:    cancel,
		msFactory: factory,
	}
	node.UpdateStateCode(internalpb.StateCode_Abnormal)
	log.Debug("ProxyNode", zap.Any("State", node.stateCode.Load()))
	return node, nil

}

// Register register proxy node at etcd
func (node *ProxyNode) Register() error {
	node.session = sessionutil.NewSession(node.ctx, Params.MetaRootPath, []string{Params.EtcdAddress})
	node.session.Init(typeutil.ProxyNodeRole, Params.NetworkAddress, false)
	Params.ProxyID = node.session.ServerID
	return nil
}

func (node *ProxyNode) Init() error {
	// todo wait for proxyservice state changed to Healthy
	ctx := context.Background()

	err := funcutil.WaitForComponentHealthy(ctx, node.proxyService, "ProxyService", 1000000, time.Millisecond*200)
	if err != nil {
		return err
	}
	log.Debug("ProxyService is ready ...")

	request := &proxypb.RegisterNodeRequest{
		Address: &commonpb.Address{
			Ip:   Params.IP,
			Port: int64(Params.NetworkPort),
		},
	}

	response, err := node.proxyService.RegisterNode(ctx, request)
	if err != nil {
		log.Debug("ProxyNode RegisterNode failed", zap.Error(err))
		return err
	}
	if response.Status.ErrorCode != commonpb.ErrorCode_Success {
		log.Debug("ProxyNode RegisterNode failed", zap.String("Reason", response.Status.Reason))
		return errors.New(response.Status.Reason)
	}

	err = Params.LoadConfigFromInitParams(response.InitParams)
	if err != nil {
		log.Debug("ProxyNode LoadConfigFromInitParams failed", zap.Error(err))
		return err
	}

	// wait for dataservice state changed to Healthy
	if node.dataService != nil {
		log.Debug("ProxyNode wait for dataService ready")
		err := funcutil.WaitForComponentHealthy(ctx, node.dataService, "DataService", 1000000, time.Millisecond*200)
		if err != nil {
			log.Debug("ProxyNode wait for dataService ready failed", zap.Error(err))
			return err
		}
		log.Debug("ProxyNode dataService is ready")
	}

	// wait for queryService state changed to Healthy
	if node.queryService != nil {
		log.Debug("ProxyNode wait for queryService ready")
		err := funcutil.WaitForComponentHealthy(ctx, node.queryService, "QueryService", 1000000, time.Millisecond*200)
		if err != nil {
			log.Debug("ProxyNode wait for queryService ready failed", zap.Error(err))
			return err
		}
		log.Debug("ProxyNode queryService is ready")
	}

	// wait for indexservice state changed to Healthy
	if node.indexService != nil {
		log.Debug("ProxyNode wait for indexService ready")
		err := funcutil.WaitForComponentHealthy(ctx, node.indexService, "IndexService", 1000000, time.Millisecond*200)
		if err != nil {
			log.Debug("ProxyNode wait for indexService ready failed", zap.Error(err))
			return err
		}
		log.Debug("ProxyNode indexService is ready")
	}

	if node.queryService != nil {
		resp, err := node.queryService.CreateQueryChannel(ctx, &querypb.CreateQueryChannelRequest{})
		if err != nil {
			log.Debug("ProxyNode CreateQueryChannel failed", zap.Error(err))
			return err
		}
		if resp.Status.ErrorCode != commonpb.ErrorCode_Success {
			log.Debug("ProxyNode CreateQueryChannel failed", zap.String("reason", resp.Status.Reason))

			return errors.New(resp.Status.Reason)
		}
		log.Debug("ProxyNode CreateQueryChannel success")

		Params.SearchChannelNames = []string{resp.RequestChannel}
		Params.SearchResultChannelNames = []string{resp.ResultChannel}
		Params.RetrieveChannelNames = []string{resp.RequestChannel}
		Params.RetrieveResultChannelNames = []string{resp.ResultChannel}
		log.Debug("ProxyNode CreateQueryChannel success", zap.Any("SearchChannelNames", Params.SearchChannelNames))
		log.Debug("ProxyNode CreateQueryChannel success", zap.Any("SearchResultChannelNames", Params.SearchResultChannelNames))
		log.Debug("ProxyNode CreateQueryChannel success", zap.Any("RetrieveChannelNames", Params.RetrieveChannelNames))
		log.Debug("ProxyNode CreateQueryChannel success", zap.Any("RetrieveResultChannelNames", Params.RetrieveResultChannelNames))
	}

	// todo
	//Params.InsertChannelNames, err = node.dataService.GetInsertChannels()
	//if err != nil {
	//	return err
	//}

	m := map[string]interface{}{
		"PulsarAddress": Params.PulsarAddress,
		"PulsarBufSize": 1024}
	err = node.msFactory.SetParams(m)
	if err != nil {
		return err
	}

	node.queryMsgStream, _ = node.msFactory.NewQueryMsgStream(node.ctx)
	node.queryMsgStream.AsProducer(Params.SearchChannelNames)
	// FIXME(wxyu): use log.Debug instead
	log.Debug("proxynode", zap.Strings("proxynode AsProducer:", Params.SearchChannelNames))
	log.Debug("create query message stream ...")

	idAllocator, err := allocator.NewIDAllocator(node.ctx, Params.MetaRootPath, []string{Params.EtcdAddress})

	if err != nil {
		return err
	}
	node.idAllocator = idAllocator
	node.idAllocator.PeerID = Params.ProxyID

	tsoAllocator, err := NewTimestampAllocator(node.masterService, Params.ProxyID)
	if err != nil {
		return err
	}
	node.tsoAllocator = tsoAllocator

	segAssigner, err := NewSegIDAssigner(node.ctx, node.dataService, node.lastTick)
	if err != nil {
		panic(err)
	}
	node.segAssigner = segAssigner
	node.segAssigner.PeerID = Params.ProxyID

	getDmlChannelsFunc := func(collectionID UniqueID) (map[vChan]pChan, error) {
		req := &milvuspb.DescribeCollectionRequest{
			Base: &commonpb.MsgBase{
				MsgType:   commonpb.MsgType_DescribeCollection,
				MsgID:     0, // todo
				Timestamp: 0, // todo
				SourceID:  0, // todo
			},
			DbName:         "", // todo
			CollectionName: "", // todo
			CollectionID:   collectionID,
			TimeStamp:      0, // todo
		}
		resp, err := node.masterService.DescribeCollection(node.ctx, req)
		if err != nil {
			log.Warn("DescribeCollection", zap.Error(err))
			return nil, err
		}
		if resp.Status.ErrorCode != 0 {
			log.Warn("DescribeCollection",
				zap.Any("ErrorCode", resp.Status.ErrorCode),
				zap.Any("Reason", resp.Status.Reason))
			return nil, err
		}
		if len(resp.VirtualChannelNames) != len(resp.PhysicalChannelNames) {
			err := fmt.Errorf(
				"len(VirtualChannelNames): %v, len(PhysicalChannelNames): %v",
				len(resp.VirtualChannelNames),
				len(resp.PhysicalChannelNames))
			log.Warn("GetDmlChannels", zap.Error(err))
			return nil, err
		}

		ret := make(map[vChan]pChan)
		for idx, name := range resp.VirtualChannelNames {
			if _, ok := ret[name]; ok {
				err := fmt.Errorf(
					"duplicated virtual channel found, vchan: %v, pchan: %v",
					name,
					resp.PhysicalChannelNames[idx])
				return nil, err
			}
			ret[name] = resp.PhysicalChannelNames[idx]
		}

		return ret, nil
	}
	mockQueryService := newMockGetChannelsService()
	chMgr := newChannelsMgr(getDmlChannelsFunc, mockQueryService.GetChannels, node.msFactory)
	node.chMgr = chMgr

	node.sched, err = NewTaskScheduler(node.ctx, node.idAllocator, node.tsoAllocator, node.msFactory)
	if err != nil {
		return err
	}

	node.tick = newTimeTick(node.ctx, node.tsoAllocator, time.Millisecond*200, node.sched.TaskDoneTest, node.msFactory)

	// TODO(dragondriver): read this from config
	interval := time.Millisecond * 200
	node.chTicker = newChannelsTimeTicker(node.ctx, interval, []string{}, node.sched.getPChanStatistics, tsoAllocator)

	return nil
}

func (node *ProxyNode) sendChannelsTimeTickLoop() {
	node.wg.Add(1)
	go func() {
		defer node.wg.Done()

		// TODO(dragondriver): read this from config
		interval := time.Millisecond * 200
		timer := time.NewTicker(interval)

		for {
			select {
			case <-node.ctx.Done():
				return
			case <-timer.C:
				stats, err := node.chTicker.getMinTsStatistics()
				if err != nil {
					log.Warn("sendChannelsTimeTickLoop.getMinTsStatistics", zap.Error(err))
					continue
				}

				log.Debug("send timestamp statistics of pchan", zap.Any("statistics", stats))

				channels := make([]pChan, 0, len(stats))
				tss := make([]Timestamp, 0, len(stats))

				for channel, ts := range stats {
					channels = append(channels, channel)
					tss = append(tss, ts)
				}

				req := &internalpb.ChannelTimeTickMsg{
					Base: &commonpb.MsgBase{
						MsgType:   commonpb.MsgType_TimeTick, // todo
						MsgID:     0,                         // todo
						Timestamp: 0,                         // todo
						SourceID:  node.session.ServerID,
					},
					ChannelNames: channels,
					Timestamps:   tss,
				}

				status, err := node.masterService.UpdateChannelTimeTick(node.ctx, req)
				if err != nil {
					log.Warn("sendChannelsTimeTickLoop.UpdateChannelTimeTick", zap.Error(err))
					continue
				}
				if status.ErrorCode != 0 {
					log.Warn("sendChannelsTimeTickLoop.UpdateChannelTimeTick",
						zap.Any("ErrorCode", status.ErrorCode),
						zap.Any("Reason", status.Reason))
					continue
				}
			}
		}
	}()
}

func (node *ProxyNode) Start() error {
	err := InitMetaCache(node.masterService)
	if err != nil {
		return err
	}
	log.Debug("init global meta cache ...")

	initGlobalInsertChannelsMap(node)
	log.Debug("init global insert channels map ...")

	node.queryMsgStream.Start()
	log.Debug("start query message stream ...")

	node.sched.Start()
	log.Debug("start scheduler ...")

	node.idAllocator.Start()
	log.Debug("start id allocator ...")

	node.segAssigner.Start()
	log.Debug("start seg assigner ...")

	node.tick.Start()
	log.Debug("start time tick ...")

	err = node.chTicker.start()
	if err != nil {
		return err
	}
	log.Debug("start channelsTimeTicker")

	node.sendChannelsTimeTickLoop()

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

	node.UpdateStateCode(internalpb.StateCode_Healthy)
	log.Debug("ProxyNode", zap.Any("State", node.stateCode.Load()))

	return nil
}

func (node *ProxyNode) Stop() error {
	node.cancel()

	globalInsertChannelsMap.CloseAllMsgStream()
	node.idAllocator.Close()
	node.segAssigner.Close()
	node.sched.Close()
	node.queryMsgStream.Close()
	node.tick.Close()
	err := node.chTicker.close()
	if err != nil {
		return err
	}

	node.wg.Wait()

	for _, cb := range node.closeCallbacks {
		cb()
	}

	return nil
}

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

func (node *ProxyNode) lastTick() Timestamp {
	return node.tick.LastTick()
}

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

func (node *ProxyNode) SetMasterClient(cli types.MasterService) {
	node.masterService = cli
}

func (node *ProxyNode) SetIndexServiceClient(cli types.IndexService) {
	node.indexService = cli
}

func (node *ProxyNode) SetDataServiceClient(cli types.DataService) {
	node.dataService = cli
}

func (node *ProxyNode) SetProxyServiceClient(cli types.ProxyService) {
	node.proxyService = cli
}

func (node *ProxyNode) SetQueryServiceClient(cli types.QueryService) {
	node.queryService = cli
}