milvus/internal/queryservice/queryservice.go

package queryservice

import (
	"context"
	"log"
	"strconv"
	"sync/atomic"
	"time"

	"github.com/zilliztech/milvus-distributed/internal/distributed/masterservice"
	grpcquerynodeclient "github.com/zilliztech/milvus-distributed/internal/distributed/querynode/client"
	"github.com/zilliztech/milvus-distributed/internal/errors"
	"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
	"github.com/zilliztech/milvus-distributed/internal/proto/internalpb2"
	"github.com/zilliztech/milvus-distributed/internal/proto/querypb"
	"github.com/zilliztech/milvus-distributed/internal/querynode"
	"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
)

type Interface = typeutil.QueryServiceInterface

type QueryService struct {
	loopCtx    context.Context
	loopCancel context.CancelFunc

	QueryServiceID uint64
	replica        metaReplica

	// type(masterServiceClient) should be interface
	// masterServiceClient masterService.Service
	masterServiceClient *masterservice.GrpcClient
	queryNodeClient     map[int]querynode.Node
	numQueryNode        int
	numQueryChannel     int

	stateCode atomic.Value
	isInit    atomic.Value
}

type InitParams struct {
	Distributed bool
}

//serverBase interface
func (qs *QueryService) Init() error {
	if Params.Distributed {
		var err error
		//TODO:: alter 2*second
		qs.masterServiceClient, err = masterservice.NewGrpcClient(Params.MasterServiceAddress, 2*time.Second)
		if err != nil {
			return err

		}
	} else {
		//TODO:: create masterService.Core{}
		log.Fatal(errors.New("should not use grpc client"))
	}

	qs.isInit.Store(true)
	return nil
}

func (qs *QueryService) InitParams(params *InitParams) {
	Params.Distributed = params.Distributed
}

func (qs *QueryService) Start() error {
	isInit := qs.isInit.Load().(bool)
	if !isInit {
		return errors.New("call start before init")
	}
	qs.stateCode.Store(internalpb2.StateCode_HEALTHY)
	return nil
}

func (qs *QueryService) Stop() error {
	qs.loopCancel()
	qs.stateCode.Store(internalpb2.StateCode_ABNORMAL)
	return nil
}

//func (qs *QueryService) SetDataService(p querynode.DataServiceInterface) error {
//	for k, v := range qs.queryNodeClient {
//		v.Set
//	}
//	return c.SetDataService(p)
//}
//
//func (qs *QueryService) SetIndexService(p querynode.IndexServiceInterface) error {
//	c, ok := s.core.(*cms.Core)
//	if !ok {
//		return errors.Errorf("set index service failed")
//	}
//	return c.SetIndexService(p)
//}

func (qs *QueryService) GetComponentStates() (*internalpb2.ComponentStates, error) {
	serviceComponentInfo := &internalpb2.ComponentInfo{
		NodeID:    Params.QueryServiceID,
		StateCode: qs.stateCode.Load().(internalpb2.StateCode),
	}
	subComponentInfos := make([]*internalpb2.ComponentInfo, 0)
	for nodeID, nodeClient := range qs.queryNodeClient {
		componentStates, err := nodeClient.GetComponentStates()
		if err != nil {
			subComponentInfos = append(subComponentInfos, &internalpb2.ComponentInfo{
				NodeID:    int64(nodeID),
				StateCode: internalpb2.StateCode_ABNORMAL,
			})
			continue
		}
		subComponentInfos = append(subComponentInfos, componentStates.State)
	}
	return &internalpb2.ComponentStates{
		State:              serviceComponentInfo,
		SubcomponentStates: subComponentInfos,
	}, nil
}

func (qs *QueryService) GetTimeTickChannel() (string, error) {
	return Params.TimeTickChannelName, nil
}

func (qs *QueryService) GetStatisticsChannel() (string, error) {
	return Params.StatsChannelName, nil
}

func (qs *QueryService) RegisterNode(req *querypb.RegisterNodeRequest) (*querypb.RegisterNodeResponse, error) {
	allocatedID := qs.numQueryNode
	qs.numQueryNode++

	registerNodeAddress := req.Address.Ip + ":" + strconv.FormatInt(req.Address.Port, 10)
	var client querynode.Node
	if Params.Distributed {
		client = grpcquerynodeclient.NewClient(registerNodeAddress)
	} else {
		log.Fatal(errors.New("should be queryNodeImpl.QueryNode"))
	}
	qs.queryNodeClient[allocatedID] = client

	return &querypb.RegisterNodeResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_SUCCESS,
		},
		InitParams: &internalpb2.InitParams{
			NodeID: int64(allocatedID),
		},
	}, nil
}

func (qs *QueryService) ShowCollections(req *querypb.ShowCollectionRequest) (*querypb.ShowCollectionResponse, error) {
	dbID := req.DbID
	collectionIDs, err := qs.replica.getCollectionIDs(dbID)
	if err != nil {
		return nil, err
	}
	return &querypb.ShowCollectionResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_SUCCESS,
		},
		CollectionIDs: collectionIDs,
	}, nil
}

func (qs *QueryService) LoadCollection(req *querypb.LoadCollectionRequest) (*commonpb.Status, error) {
	panic("implement me")
}

func (qs *QueryService) ReleaseCollection(req *querypb.ReleaseCollectionRequest) (*commonpb.Status, error) {
	panic("implement me")
}

func (qs *QueryService) ShowPartitions(req *querypb.ShowPartitionRequest) (*querypb.ShowPartitionResponse, error) {
	dbID := req.DbID
	collectionID := req.CollectionID
	partitionIDs, err := qs.replica.getPartitionIDs(dbID, collectionID)
	if err != nil {
		return nil, err
	}
	return &querypb.ShowPartitionResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_SUCCESS,
		},
		PartitionIDs: partitionIDs,
	}, nil

}

func (qs *QueryService) LoadPartitions(req *querypb.LoadPartitionRequest) (*commonpb.Status, error) {
	panic("implement me")
}

func (qs *QueryService) ReleasePartitions(req *querypb.ReleasePartitionRequest) (*commonpb.Status, error) {
	panic("implement me")
}

func (qs *QueryService) CreateQueryChannel() (*querypb.CreateQueryChannelResponse, error) {
	channelID := qs.numQueryChannel
	qs.numQueryChannel++
	allocatedQueryChannel := "query-" + strconv.FormatInt(int64(channelID), 10)
	allocatedQueryResultChannel := "queryResult-" + strconv.FormatInt(int64(channelID), 10)

	return &querypb.CreateQueryChannelResponse{
		Status: &commonpb.Status{
			ErrorCode: commonpb.ErrorCode_SUCCESS,
		},
		RequestChannel: allocatedQueryChannel,
		ResultChannel:  allocatedQueryResultChannel,
	}, nil
}

func (qs *QueryService) GetPartitionStates(req *querypb.PartitionStatesRequest) (*querypb.PartitionStatesResponse, error) {
	panic("implement me")
}

func NewQueryService(ctx context.Context) (Interface, error) {
	Params.Init()
	nodeClients := make(map[int]querynode.Node)
	ctx1, cancel := context.WithCancel(ctx)
	service := &QueryService{
		loopCtx:         ctx1,
		loopCancel:      cancel,
		numQueryNode:    0,
		queryNodeClient: nodeClients,
		numQueryChannel: 0,
	}
	service.stateCode.Store(internalpb2.StateCode_INITIALIZING)
	service.isInit.Store(false)
	return service, nil
}